1use crate::agent::compact::estimate_messages_tokens;
11use crate::agent::retrieval::RetrievalState;
12use crate::config::AgentConfig;
13use crate::error::{AacodeError, Result};
14use crate::llm::types::{ChatMessage, ToolCall};
15use crate::llm::LlmClient;
16use crate::session::{SessionManager, SessionMessage};
17use crate::stream::{CollectingSink, EventSink};
18use crate::tools::ToolRegistry;
19use serde_json::Value;
20use std::collections::HashMap;
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::sync::Arc;
23use tokio::sync::Notify;
24
25#[derive(Debug, Clone, PartialEq)]
27pub enum RunStatus {
28 Completed,
29 MaxIterations,
30 Cancelled,
31 Error(String),
32}
33
34impl RunStatus {
35 pub fn as_str(&self) -> &'static str {
38 match self {
39 RunStatus::Completed => "completed",
40 RunStatus::MaxIterations => "max_iterations",
41 RunStatus::Cancelled => "cancelled",
42 RunStatus::Error(_) => "error",
43 }
44 }
45}
46
47pub struct RunResult {
48 pub status: RunStatus,
49 pub iterations: u32,
50 pub final_text: String,
51}
52
53impl RunResult {
54 pub fn to_result_json(&self) -> serde_json::Value {
57 serde_json::json!({
58 "status": self.status.as_str(),
59 "iterations": self.iterations,
60 "final_text": self.final_text,
61 })
62 }
63}
64
65pub struct ReactLoop<'a> {
67 pub llm: &'a dyn LlmClient,
68 pub tools: &'a ToolRegistry,
69 pub config: &'a AgentConfig,
70 pub native_tools: Vec<Value>,
72}
73
74impl<'a> ReactLoop<'a> {
75 pub fn new(
76 llm: &'a dyn LlmClient,
77 tools: &'a ToolRegistry,
78 config: &'a AgentConfig,
79 native_tools: Vec<Value>,
80 ) -> Self {
81 ReactLoop {
82 llm,
83 tools,
84 config,
85 native_tools,
86 }
87 }
88
89 pub async fn run(
92 &self,
93 mut messages: Vec<ChatMessage>,
94 session: &mut SessionManager,
95 emitter: &dyn EventSink,
96 cancel: &AtomicBool,
97 ) -> Result<RunResult> {
98 let cancel_notify = Arc::new(Notify::new());
99 let mut stale = StaleTracker::default();
100 let project_path = session.project_path();
104 let mut retrieval = RetrievalState::from_value(session.context_checkpoint());
105 let session_id = session.current_session_id.clone().unwrap_or_default();
106 let usage_tracker = crate::usage::UsageTracker::new(&session.project_path());
108 let run_uuid = uuid::Uuid::new_v4().simple().to_string();
109 let run_id = &run_uuid[..12.min(run_uuid.len())];
110 let mut last_text = String::new();
112
113 for iteration in 0..self.config.max_iterations {
114 if cancel.load(Ordering::SeqCst) {
115 emitter.error("cancelled");
116 return Ok(RunResult {
117 status: RunStatus::Cancelled,
118 iterations: iteration,
119 final_text: String::new(),
120 });
121 }
122
123 let (mut view, compacted, tokens, checkpoint_changed) =
127 retrieval.model_view(&messages, &self.config.context, &project_path, &session_id);
128 if checkpoint_changed {
129 let _ = session.set_context_checkpoint(retrieval.to_value());
130 }
131 if tokens > self.config.context.max_context_tokens {
135 return Err(AacodeError::Api(format!(
136 "context too large: {} tokens exceeds limit of {} ({} messages). Try reducing history or restarting the session.",
137 tokens,
138 self.config.context.max_context_tokens,
139 messages.len(),
140 )));
141 }
142 if compacted {
146 crate::agent::sanitize::sanitize_history(&mut view);
147 }
148
149 let resp = match self
150 .chat_with_retry(&view, emitter, cancel, &cancel_notify)
151 .await
152 {
153 Ok(r) => r,
154 Err(AacodeError::Cancelled) => {
155 emitter.error("cancelled");
156 return Ok(RunResult {
157 status: RunStatus::Cancelled,
158 iterations: iteration,
159 final_text: String::new(),
160 });
161 }
162 Err(e) => {
163 let classified = classify_api_error(&e);
164 emitter.error(&classified);
165 return Ok(RunResult {
166 status: RunStatus::Error(classified),
167 iterations: iteration,
168 final_text: String::new(),
169 });
170 }
171 };
172
173 {
177 let gateway = match self.config.model.gateway {
178 crate::config::Gateway::Anthropic => "anthropic",
179 crate::config::Gateway::Openai => "openai",
180 };
181 let mut event = crate::usage::usage_event(
182 &session_id,
183 run_id,
184 iteration + 1,
185 gateway,
186 &self.config.model.name,
187 tokens,
188 view.len(),
189 compacted,
190 resp.usage.as_ref(),
191 );
192 let tools: Vec<String> = self
196 .native_tools
197 .iter()
198 .filter_map(|t| {
199 t.get("function")
200 .and_then(|f| f.get("name"))
201 .and_then(|n| n.as_str())
202 .map(str::to_string)
203 })
204 .collect();
205 event["tools"] = serde_json::json!(tools);
206 if let Some(cp) = retrieval.checkpoint.as_ref() {
207 event["retrieval_cutoff"] = serde_json::json!(cp.cutoff_index);
208 event["retrieval_stub_chars"] = serde_json::json!(cp.stub.chars().count());
209 }
210 if usage_tracker.record(&event) && usage_tracker.verbose {
211 if let Some(u) = &resp.usage {
212 eprintln!(
213 "📊 usage iter {}: input={} (cache {:.1}%) output={} reasoning={}",
214 iteration + 1,
215 u.input_tokens,
216 u.cache_hit_rate() * 100.0,
217 u.output_tokens,
218 u.reasoning_tokens
219 );
220 }
221 }
222 }
223
224 last_text = resp.text.clone();
225
226 if resp.tool_calls.is_empty() {
228 if resp.is_truncated() {
232 let assistant = ChatMessage {
233 role: "assistant".into(),
234 content: last_text.clone(),
235 tool_calls: None,
236 tool_call_id: None,
237 reasoning_content: resp.reasoning_content.clone(),
238 };
239 messages.push(assistant.clone());
240 let _ = session.add_message(SessionMessage::from_chat(&assistant));
241 let cont = ChatMessage::user("continue");
242 messages.push(cont.clone());
243 let _ = session.add_message(SessionMessage::from_chat(&cont));
244 continue;
245 }
246
247 let assistant = ChatMessage {
248 role: "assistant".into(),
249 content: last_text.clone(),
250 tool_calls: None,
251 tool_call_id: None,
252 reasoning_content: resp.reasoning_content.clone(),
253 };
254 messages.push(assistant.clone());
255 let _ = session.add_message(SessionMessage::from_chat(&assistant));
256 let _ = session.flush();
259 emitter.done_result(&session_id, "completed", iteration + 1, &last_text);
260 return Ok(RunResult {
261 status: RunStatus::Completed,
262 iterations: iteration + 1,
263 final_text: last_text,
264 });
265 }
266
267 if let Some(w) = stale.detect(&resp.tool_calls) {
270 let warn_msg = ChatMessage::system(format!("[SYSTEM WARNING]: {w}"));
271 let _ = session.add_message(SessionMessage::from_chat(&warn_msg));
272 messages.push(warn_msg);
273 }
274
275 let assistant =
277 ChatMessage::assistant_with_tools(resp.text.clone(), resp.tool_calls.clone());
278 let mut assistant_msg = assistant.clone();
279 assistant_msg.reasoning_content = resp.reasoning_content.clone();
280 messages.push(assistant_msg.clone());
281 let _ = session.add_message(SessionMessage::from_chat(&assistant_msg));
282
283 let tool_count = resp.tool_calls.len();
289 let has_lane = resp
290 .tool_calls
291 .iter()
292 .any(|tc| self.tools.lane_of(&tc.name).is_some());
293 if tool_count > 1 && !has_lane {
294 use futures::future::join_all;
295 let cancel_arc = Arc::new(AtomicBool::new(false));
296 let cancel_ref = cancel as &AtomicBool; let futures: Vec<_> = resp
299 .tool_calls
300 .iter()
301 .map(|tc| {
302 let name = tc.name.clone();
303 let args = tc.parsed_args();
304 let cancel = cancel_arc.clone();
305 async move {
306 let started = std::time::Instant::now();
307 if cancel_ref.load(Ordering::SeqCst) || cancel.load(Ordering::SeqCst) {
308 return (
309 tc.id.clone(),
310 tc.name.clone(),
311 "cancelled".to_string(),
312 0u64,
313 );
314 }
315 let obs = self
316 .execute_with_retry(&name, args, emitter, cancel.as_ref())
317 .await;
318 (
319 tc.id.clone(),
320 tc.name.clone(),
321 obs,
322 started.elapsed().as_millis() as u64,
323 )
324 }
325 })
326 .collect();
327
328 let results = join_all(futures).await;
329
330 for (tc_id, tc_name, observation, duration_ms) in results {
331 if cancel.load(Ordering::SeqCst) {
332 emitter.error("cancelled");
333 return Ok(RunResult {
334 status: RunStatus::Cancelled,
335 iterations: iteration + 1,
336 final_text: String::new(),
337 });
338 }
339 if tc_name == "fetch_url" {
340 let url = resp
341 .tool_calls
342 .iter()
343 .find(|tc| tc.id == tc_id)
344 .and_then(|tc| {
345 tc.parsed_args()
346 .get("url")
347 .and_then(|v| v.as_str())
348 .map(|s| s.to_string())
349 })
350 .unwrap_or_default();
351 stale.record_fetch(&url, &observation);
352 }
353 let (status, code) = classify_tool_status(&observation);
354 emitter.tool_status(
355 &tc_id,
356 &tc_name,
357 status,
358 duration_ms,
359 code,
360 if code.is_some() {
361 Some(observation.as_str())
362 } else {
363 None
364 },
365 );
366 emitter.seg_observation(&observation, self.config.limits.display_preview_chars);
367 let tool_msg = ChatMessage::tool_result(tc_id, observation);
368 messages.push(tool_msg.clone());
369 let _ = session.add_message(SessionMessage::from_chat(&tool_msg));
370 }
371 } else {
372 for tc in &resp.tool_calls {
374 if cancel.load(Ordering::SeqCst) {
375 emitter.error("cancelled");
376 return Ok(RunResult {
377 status: RunStatus::Cancelled,
378 iterations: iteration + 1,
379 final_text: String::new(),
380 });
381 }
382 let args = tc.parsed_args();
383 let started = std::time::Instant::now();
384 let observation = self
385 .execute_with_retry(&tc.name, args, emitter, cancel)
386 .await;
387 let duration_ms = started.elapsed().as_millis() as u64;
388 if tc.name == "fetch_url" {
389 let url = tc
390 .parsed_args()
391 .get("url")
392 .and_then(|v| v.as_str())
393 .unwrap_or("")
394 .to_string();
395 stale.record_fetch(&url, &observation);
396 }
397 let (status, code) = classify_tool_status(&observation);
398 emitter.tool_status(
399 &tc.id,
400 &tc.name,
401 status,
402 duration_ms,
403 code,
404 if code.is_some() {
405 Some(observation.as_str())
406 } else {
407 None
408 },
409 );
410 emitter.seg_observation(&observation, self.config.limits.display_preview_chars);
411 let tool_msg = ChatMessage::tool_result(tc.id.clone(), observation);
412 messages.push(tool_msg.clone());
413 let _ = session.add_message(SessionMessage::from_chat(&tool_msg));
414 }
415 }
416
417 let _ = estimate_messages_tokens(&messages);
419 }
420
421 let _ = session.flush();
422 emitter.done_result(
423 &session_id,
424 "max_iterations",
425 self.config.max_iterations,
426 &last_text,
427 );
428 Ok(RunResult {
429 status: RunStatus::MaxIterations,
430 iterations: self.config.max_iterations,
431 final_text: last_text,
432 })
433 }
434
435 async fn chat_with_retry(
437 &self,
438 view: &[ChatMessage],
439 emitter: &dyn EventSink,
440 cancel: &AtomicBool,
441 cancel_notify: &Notify,
442 ) -> Result<crate::llm::types::LlmResponse> {
443 let max = self.config.limits.max_retries.max(1);
444 let mut last: Option<AacodeError> = None;
445 for attempt in 0..max {
446 if cancel.load(Ordering::SeqCst) {
447 return Err(AacodeError::Cancelled);
448 }
449 let use_buffered = attempt > 0 && !emitter.is_tty();
457 let (result, buffered): (_, Option<CollectingSink>) = if use_buffered {
460 let b = CollectingSink::new(false);
461 let fut = self.llm.chat_stream(view, &self.native_tools, &b, cancel);
462 let r = tokio::select! {
463 r = fut => r,
464 _ = cancel_checker(cancel, cancel_notify) => Err(AacodeError::Cancelled),
465 };
466 (r, Some(b))
467 } else {
468 let fut = self
469 .llm
470 .chat_stream(view, &self.native_tools, emitter, cancel);
471 let r = tokio::select! {
472 r = fut => r,
473 _ = cancel_checker(cancel, cancel_notify) => Err(AacodeError::Cancelled),
474 };
475 (r, None)
476 };
477
478 match result {
479 Ok(r) => {
480 if let Some(collected) = buffered {
481 for line in collected.lines() {
482 emitter.emit_line(&line);
483 }
484 }
485 return Ok(r);
486 }
487 Err(AacodeError::Cancelled) => return Err(AacodeError::Cancelled),
488 Err(e) => {
489 let retryable = e.is_retryable();
490 if retryable && attempt + 1 < max {
494 emitter.tool_progress(
495 "building",
496 &format!("llm retry {}/{} ({})", attempt + 2, max, brief(&e)),
497 0,
498 );
499 }
500 last = Some(e);
501 if !retryable || attempt + 1 >= max {
502 break;
503 }
504 let delay = 1u64 << attempt.min(4);
507 let jitter = {
508 let nanos = std::time::SystemTime::now()
509 .duration_since(std::time::UNIX_EPOCH)
510 .unwrap_or_default()
511 .subsec_nanos() as u64;
512 let half_range = (delay * 250).max(1); nanos % (half_range * 2).max(1)
514 };
515 let actual = if jitter > 0 {
516 (delay * 1000 + jitter - delay * 250).max(100)
517 } else {
518 delay * 1000
519 };
520 tokio::time::sleep(std::time::Duration::from_millis(actual)).await;
521 }
522 }
523 }
524 Err(last.unwrap_or_else(|| AacodeError::Api("model call failed".into())))
525 }
526
527 async fn execute_with_retry(
529 &self,
530 name: &str,
531 args: Value,
532 _emitter: &dyn EventSink,
533 cancel: &AtomicBool,
534 ) -> String {
535 if self.config.safety.permission_mode == crate::config::PermissionMode::ReadOnly
539 && !crate::tools::registry::is_read_only(name)
540 {
541 return serde_json::json!({
542 "success": false,
543 "error": format!(
544 "read-only mode: tool '{name}' is not allowed (switch to full access to run it)"
545 ),
546 "tool": name,
547 })
548 .to_string();
549 }
550 let max_retries = self.config.limits.max_retries.max(1);
551 let mut last = String::new();
552 for attempt in 0..max_retries {
553 if cancel.load(Ordering::SeqCst) {
554 return "cancelled".to_string();
555 }
556 let notify = Notify::new();
562 let obs = tokio::select! {
563 obs = self.tools.execute(name, args.clone(), cancel) => obs,
564 _ = cancel_checker(cancel, ¬ify) => return "cancelled".to_string(),
565 };
566 let lower = obs.to_lowercase();
567 let retryable = lower.contains("timeout")
568 || lower.contains("timed out")
569 || lower.contains("connection")
570 || lower.contains("temporary")
571 || lower.contains("network");
572 last = obs;
573 if !retryable || attempt + 1 >= max_retries {
574 break;
575 }
576 }
577 last
578 }
579}
580
581fn classify_tool_status(observation: &str) -> (&'static str, Option<&'static str>) {
584 let l = observation.to_lowercase();
585 let failed = l.contains("parameter validation failed")
586 || l.contains("execution error")
587 || l.contains("execution cancelled")
588 || l.contains("unknown tool")
589 || observation.starts_with('❌')
590 || l.contains("\"success\":false")
591 || l.contains("\"success\": false")
592 || l.contains("command rejected by safety guard")
593 || l.contains("traceback");
594 if !failed {
595 return ("success", None);
596 }
597 if l.contains("safety guard") || l.contains("dangerous operation") {
598 ("error", Some("safety_guard"))
599 } else if l.contains("timed out") || l.contains("timeout") {
600 ("error", Some("timeout"))
601 } else if l.contains("parameter validation")
602 || l.contains("parameter parsing")
603 || l.contains("missing required")
604 {
605 ("error", Some("validation"))
606 } else if l.contains("unknown tool") || l.contains("not found") {
607 ("error", Some("not_found"))
608 } else if l.contains("cancel") {
609 ("error", Some("cancelled"))
610 } else {
611 ("error", Some("tool_error"))
612 }
613}
614
615fn classify_api_error(e: &AacodeError) -> String {
617 let s = e.to_string();
618 let l = s.to_lowercase();
619 if l.contains("401") || l.contains("authentication") || l.contains("invalid api key") {
620 format!("API authentication failed (check LLM_API_KEY / endpoint): {s}")
621 } else if l.contains("quota") || l.contains("rate limit") || l.contains("429") {
622 format!("API quota/rate limit reached (check usage or wait): {s}")
623 } else if l.contains("connection") || l.contains("timeout") || l.contains("network") {
624 format!("Network error reaching the model (check connectivity/endpoint): {s}")
625 } else if l.contains("config") || l.contains("api key not configured") {
626 format!("Configuration error: {s}")
627 } else {
628 s
629 }
630}
631
632#[derive(Default)]
634struct StaleTracker {
635 by_domain: HashMap<String, Vec<bool>>, warned: std::collections::HashSet<String>,
637}
638
639impl StaleTracker {
640 fn detect(&mut self, tool_calls: &[ToolCall]) -> Option<String> {
641 if !tool_calls.iter().any(|tc| tc.name == "fetch_url") {
643 return None;
644 }
645 let mut candidate: Option<String> = None;
647 for (domain, entries) in &self.by_domain {
648 if entries.len() < 3 {
649 continue;
650 }
651 let last3 = &entries[entries.len() - 3..];
652 if last3.iter().all(|has| !has) {
653 let key = format!("stale_{domain}");
654 if !self.warned.contains(&key) {
655 candidate = Some(domain.clone());
656 break;
657 }
658 }
659 }
660 let domain = candidate?;
661 self.warned.insert(format!("stale_{domain}"));
662 Some(format!(
663 "Last 3 fetch_url calls to '{domain}' returned unreadable content. Try search_web or proceed with existing knowledge."
664 ))
665 }
666
667 fn record_fetch(&mut self, url: &str, observation: &str) {
668 let domain = domain_of(url);
669 let text: String = observation
671 .chars()
672 .filter(|c| *c != '<' && *c != '>')
673 .collect();
674 let has_content = text.trim().len() >= 200;
675 let e = self.by_domain.entry(domain).or_default();
676 e.push(has_content);
677 if e.len() > 5 {
678 let drain = e.len() - 5;
679 e.drain(0..drain);
680 }
681 }
682}
683
684fn domain_of(url: &str) -> String {
685 let after_scheme = url.split("://").nth(1).unwrap_or(url);
687 after_scheme
688 .split('/')
689 .next()
690 .unwrap_or(after_scheme)
691 .to_string()
692}
693
694#[cfg(test)]
695mod tests {
696 use super::*;
697 use crate::config::AgentConfig;
698 use crate::llm::types::LlmResponse;
699 use crate::stream::CollectingSink;
700 use crate::tools::registry::Tool;
701 use crate::tools::schema::{ParamType, ToolParameter, ToolSchema};
702 use std::sync::Mutex;
703
704 struct ScriptedLlm {
706 responses: Mutex<Vec<LlmResponse>>,
707 }
708 #[async_trait::async_trait]
709 impl LlmClient for ScriptedLlm {
710 async fn chat_stream(
711 &self,
712 _m: &[ChatMessage],
713 _t: &[Value],
714 emitter: &dyn EventSink,
715 _c: &AtomicBool,
716 ) -> Result<LlmResponse> {
717 let mut q = self.responses.lock().unwrap();
718 if q.is_empty() {
719 return Err(AacodeError::Api("no more scripted responses".into()));
720 }
721 let r = q.remove(0);
722 emitter.seg("thought", &r.text);
723 Ok(r)
724 }
725 async fn validate(&self) -> Result<()> {
726 Ok(())
727 }
728 }
729
730 struct RecordingLlm {
732 responses: Mutex<Vec<LlmResponse>>,
733 seen: Mutex<Vec<Vec<ChatMessage>>>,
734 }
735 #[async_trait::async_trait]
736 impl LlmClient for RecordingLlm {
737 async fn chat_stream(
738 &self,
739 m: &[ChatMessage],
740 _t: &[Value],
741 emitter: &dyn EventSink,
742 _c: &AtomicBool,
743 ) -> Result<LlmResponse> {
744 self.seen.lock().unwrap().push(m.to_vec());
745 let mut q = self.responses.lock().unwrap();
746 if q.is_empty() {
747 return Err(AacodeError::Api("no more scripted responses".into()));
748 }
749 let r = q.remove(0);
750 emitter.seg("thought", &r.text);
751 Ok(r)
752 }
753 async fn validate(&self) -> Result<()> {
754 Ok(())
755 }
756 }
757
758 struct EchoTool;
759 #[async_trait::async_trait]
760 impl Tool for EchoTool {
761 fn schema(&self) -> ToolSchema {
762 ToolSchema::new(
763 "echo",
764 "echo",
765 vec![ToolParameter::new(
766 "text",
767 ParamType::String,
768 true,
769 "t",
770 &[],
771 )],
772 )
773 }
774 async fn call(&self, args: &Value, _c: &AtomicBool) -> Result<String> {
775 Ok(format!("echoed: {}", args["text"].as_str().unwrap_or("")))
776 }
777 }
778
779 struct SleepyTool {
780 name: &'static str,
781 delay_ms: u64,
782 lane: Option<&'static str>,
783 log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
784 }
785 #[async_trait::async_trait]
786 impl Tool for SleepyTool {
787 fn schema(&self) -> ToolSchema {
788 ToolSchema::new(self.name, "sleep", vec![])
789 }
790 fn lane(&self) -> Option<&'static str> {
791 self.lane
792 }
793 async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
794 self.log
795 .lock()
796 .unwrap()
797 .push(format!("start:{}", self.name));
798 tokio::time::sleep(std::time::Duration::from_millis(self.delay_ms)).await;
799 self.log.lock().unwrap().push(format!("end:{}", self.name));
800 Ok(format!("{} done", self.name))
801 }
802 }
803
804 fn tmp_proj() -> std::path::PathBuf {
805 let d = std::env::temp_dir().join(format!(
806 "aacode_react_{}_{}",
807 std::process::id(),
808 uuid::Uuid::new_v4().simple()
809 ));
810 std::fs::create_dir_all(&d).unwrap();
811 d
812 }
813
814 fn tc(name: &str, args: &str) -> ToolCall {
815 ToolCall {
816 id: "c1".into(),
817 name: name.into(),
818 arguments: args.into(),
819 }
820 }
821
822 fn named_call(id: &str, name: &str) -> ToolCall {
823 ToolCall {
824 id: id.into(),
825 name: name.into(),
826 arguments: "{}".into(),
827 }
828 }
829
830 #[tokio::test]
832 async fn lane_batch_runs_sequentially() {
833 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
834 let mut reg = ToolRegistry::new();
835 reg.register(Box::new(SleepyTool {
836 name: "a",
837 delay_ms: 30,
838 lane: Some("x"),
839 log: log.clone(),
840 }));
841 reg.register(Box::new(SleepyTool {
842 name: "b",
843 delay_ms: 30,
844 lane: Some("x"),
845 log: log.clone(),
846 }));
847 let llm = ScriptedLlm {
848 responses: Mutex::new(vec![
849 LlmResponse {
850 tool_calls: vec![named_call("c1", "a"), named_call("c2", "b")],
851 ..Default::default()
852 },
853 LlmResponse {
854 text: "done".into(),
855 ..Default::default()
856 },
857 ]),
858 };
859 let mut cfg = AgentConfig::default();
860 cfg.max_iterations = 3;
861 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
862 let proj = tmp_proj();
863 let mut sm = SessionManager::new(&proj);
864 sm.create_session("t", None).unwrap();
865 let sink = CollectingSink::new(false);
866 let cancel = AtomicBool::new(false);
867 loop_
868 .run(
869 vec![ChatMessage::system("s"), ChatMessage::user("u")],
870 &mut sm,
871 &sink,
872 &cancel,
873 )
874 .await
875 .unwrap();
876
877 let events = log.lock().unwrap().clone();
878 assert_eq!(events, vec!["start:a", "end:a", "start:b", "end:b"]);
879 let ids: Vec<String> = sm
880 .messages
881 .iter()
882 .filter(|m| m.role == "tool")
883 .map(|m| m.tool_call_id.clone().unwrap_or_default())
884 .collect();
885 assert_eq!(ids, vec!["c1".to_string(), "c2".to_string()]);
886 }
887
888 #[tokio::test]
891 async fn concurrent_batch_submits_results_in_call_order() {
892 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
893 let mut reg = ToolRegistry::new();
894 reg.register(Box::new(SleepyTool {
895 name: "slow",
896 delay_ms: 80,
897 lane: None,
898 log: log.clone(),
899 }));
900 reg.register(Box::new(SleepyTool {
901 name: "fast",
902 delay_ms: 5,
903 lane: None,
904 log: log.clone(),
905 }));
906 let llm = ScriptedLlm {
907 responses: Mutex::new(vec![
908 LlmResponse {
909 tool_calls: vec![named_call("c1", "slow"), named_call("c2", "fast")],
910 ..Default::default()
911 },
912 LlmResponse {
913 text: "done".into(),
914 ..Default::default()
915 },
916 ]),
917 };
918 let mut cfg = AgentConfig::default();
919 cfg.max_iterations = 3;
920 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
921 let proj = tmp_proj();
922 let mut sm = SessionManager::new(&proj);
923 sm.create_session("t", None).unwrap();
924 let sink = CollectingSink::new(false);
925 let cancel = AtomicBool::new(false);
926 loop_
927 .run(
928 vec![ChatMessage::system("s"), ChatMessage::user("u")],
929 &mut sm,
930 &sink,
931 &cancel,
932 )
933 .await
934 .unwrap();
935
936 let ids: Vec<String> = sm
937 .messages
938 .iter()
939 .filter(|m| m.role == "tool")
940 .map(|m| m.tool_call_id.clone().unwrap_or_default())
941 .collect();
942 assert_eq!(ids, vec!["c1".to_string(), "c2".to_string()]);
943 let events = log.lock().unwrap().clone();
944 let end_fast = events.iter().position(|e| e == "end:fast");
945 let end_slow = events.iter().position(|e| e == "end:slow");
946 assert!(
947 end_fast < end_slow,
948 "independent tools should overlap: {events:?}"
949 );
950 }
951
952 #[tokio::test]
954 async fn read_only_mode_blocks_mutating_tools() {
955 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
956 let mut reg = ToolRegistry::new();
957 reg.register(Box::new(SleepyTool {
958 name: "run_shell",
959 delay_ms: 1,
960 lane: None,
961 log: log.clone(),
962 }));
963 reg.register(Box::new(SleepyTool {
964 name: "get_page_title",
965 delay_ms: 1,
966 lane: None,
967 log: log.clone(),
968 }));
969 let llm = ScriptedLlm {
970 responses: Mutex::new(vec![
971 LlmResponse {
972 tool_calls: vec![
973 named_call("c1", "run_shell"),
974 named_call("c2", "get_page_title"),
975 ],
976 ..Default::default()
977 },
978 LlmResponse {
979 text: "done".into(),
980 ..Default::default()
981 },
982 ]),
983 };
984 let mut cfg = AgentConfig::default();
985 cfg.max_iterations = 3;
986 cfg.safety.permission_mode = crate::config::PermissionMode::ReadOnly;
987 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
988 let proj = tmp_proj();
989 let mut sm = SessionManager::new(&proj);
990 sm.create_session("t", None).unwrap();
991 let sink = CollectingSink::new(false);
992 let cancel = AtomicBool::new(false);
993 loop_
994 .run(
995 vec![ChatMessage::system("s"), ChatMessage::user("u")],
996 &mut sm,
997 &sink,
998 &cancel,
999 )
1000 .await
1001 .unwrap();
1002
1003 let events = log.lock().unwrap().clone();
1004 assert!(
1005 !events.iter().any(|e| e == "start:run_shell"),
1006 "mutating tool ran in read-only mode: {events:?}"
1007 );
1008 assert!(events.iter().any(|e| e == "start:get_page_title"));
1009 let blocked = sm
1010 .messages
1011 .iter()
1012 .find(|m| m.role == "tool" && m.tool_call_id.as_deref() == Some("c1"))
1013 .expect("blocked tool result");
1014 assert!(
1015 blocked.content.contains("read-only mode"),
1016 "{}",
1017 blocked.content
1018 );
1019 }
1020
1021 #[test]
1022 fn classifies_tool_status() {
1023 assert_eq!(super::classify_tool_status("all good"), ("success", None));
1024 assert_eq!(
1025 super::classify_tool_status(
1026 "{\"success\":false,\"hint\":\"command timed out — split it\"}"
1027 ),
1028 ("error", Some("timeout"))
1029 );
1030 assert_eq!(
1031 super::classify_tool_status("Parameter validation failed\n\nmissing required"),
1032 ("error", Some("validation"))
1033 );
1034 assert_eq!(
1035 super::classify_tool_status("Execution error: boom"),
1036 ("error", Some("tool_error"))
1037 );
1038 assert_eq!(
1039 super::classify_tool_status("Error: Unknown tool 'x'"),
1040 ("error", Some("not_found"))
1041 );
1042 assert_eq!(
1043 super::classify_tool_status("Execution cancelled"),
1044 ("error", Some("cancelled"))
1045 );
1046 assert_eq!(
1047 super::classify_tool_status(
1048 "{\"success\":false,\"error\":\"Command rejected by safety guard\"}"
1049 ),
1050 ("error", Some("safety_guard"))
1051 );
1052 }
1053
1054 #[tokio::test]
1057 async fn retrieval_compacts_view_and_persists_checkpoint() {
1058 let mut cfg = AgentConfig::default();
1059 cfg.max_iterations = 2;
1060 cfg.context.max_context_tokens = 2_000_000;
1061 cfg.context.retrieval_trigger_tokens = 50;
1062 cfg.context.retrieval_keep_tokens = 40;
1063 cfg.context.retrieval_keep_min_rounds = 1;
1064 cfg.context.retrieval_keep_recent_user_messages = 1;
1065
1066 let llm = RecordingLlm {
1067 responses: Mutex::new(vec![LlmResponse {
1068 text: "done".into(),
1069 ..Default::default()
1070 }]),
1071 seen: Mutex::new(vec![]),
1072 };
1073 let reg = ToolRegistry::new();
1074 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1075
1076 let proj = tmp_proj();
1077 let mut sm = SessionManager::new(&proj);
1078 sm.create_session("seed", None).unwrap();
1079
1080 let mut msgs = vec![ChatMessage::system("SYS")];
1081 for i in 0..20 {
1082 msgs.push(ChatMessage::user(format!("task {i} {}", "x".repeat(120))));
1083 msgs.push(ChatMessage::assistant("ok"));
1084 }
1085
1086 let sink = CollectingSink::new(false);
1087 let cancel = AtomicBool::new(false);
1088 let res = loop_
1089 .run(msgs.clone(), &mut sm, &sink, &cancel)
1090 .await
1091 .unwrap();
1092 assert_eq!(res.status, RunStatus::Completed);
1093
1094 let seen = llm.seen.lock().unwrap();
1095 let first = &seen[0];
1096 assert!(first.len() < msgs.len(), "view should be compacted");
1097 assert_eq!(first[0].role, "system");
1098 assert!(
1099 first[1].content.starts_with("<conversation-checkpoint>"),
1100 "checkpoint stub missing: {}",
1101 first[1].content
1102 );
1103 let cp = sm.context_checkpoint().expect("checkpoint persisted");
1104 assert!(cp["cutoff_index"].as_u64().unwrap() > 0);
1105 }
1106
1107 #[tokio::test]
1108 async fn completes_when_no_tool_calls() {
1109 let llm = ScriptedLlm {
1110 responses: Mutex::new(vec![LlmResponse {
1111 text: "all done".into(),
1112 ..Default::default()
1113 }]),
1114 };
1115 let mut reg = ToolRegistry::new();
1116 reg.register(Box::new(EchoTool));
1117 let cfg = AgentConfig::default();
1118 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1119
1120 let proj = tmp_proj();
1121 let mut sm = SessionManager::new(&proj);
1122 sm.create_session("t", None).unwrap();
1123 let sink = CollectingSink::new(false);
1124 let cancel = AtomicBool::new(false);
1125
1126 let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("hi")];
1127 let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
1128 assert_eq!(res.status, RunStatus::Completed);
1129 assert_eq!(res.final_text, "all done");
1130 assert!(sink.lines().iter().any(|l| l.contains(r#""type":"done""#)));
1131 }
1132
1133 #[tokio::test]
1134 async fn executes_tool_then_completes() {
1135 let llm = ScriptedLlm {
1136 responses: Mutex::new(vec![
1137 LlmResponse {
1138 text: "let me echo".into(),
1139 tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
1140 ..Default::default()
1141 },
1142 LlmResponse {
1143 text: "done".into(),
1144 ..Default::default()
1145 },
1146 ]),
1147 };
1148 let mut reg = ToolRegistry::new();
1149 reg.register(Box::new(EchoTool));
1150 let cfg = AgentConfig::default();
1151 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1152
1153 let proj = tmp_proj();
1154 let mut sm = SessionManager::new(&proj);
1155 sm.create_session("t", None).unwrap();
1156 let sink = CollectingSink::new(false);
1157 let cancel = AtomicBool::new(false);
1158
1159 let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
1160 let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
1161 assert_eq!(res.status, RunStatus::Completed);
1162 assert!(sink.lines().iter().any(|l| l.contains("echoed: hey")));
1164 assert!(sm.messages.iter().any(|m| m.tool_calls.is_some()));
1166 assert!(sm.messages.iter().any(|m| m.role == "tool"));
1167 }
1168
1169 #[tokio::test]
1170 async fn persists_turn_messages_without_segment_duplication() {
1171 let llm = ScriptedLlm {
1172 responses: Mutex::new(vec![
1173 LlmResponse {
1174 text: "let me echo".into(),
1175 tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
1176 reasoning_content: Some("I need to echo".into()),
1177 ..Default::default()
1178 },
1179 LlmResponse {
1180 text: "all done".into(),
1181 reasoning_content: Some("summarize".into()),
1182 ..Default::default()
1183 },
1184 ]),
1185 };
1186 let mut reg = ToolRegistry::new();
1187 reg.register(Box::new(EchoTool));
1188 let cfg = AgentConfig::default();
1189 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1190
1191 let proj = tmp_proj();
1192 let mut sm = SessionManager::new(&proj);
1193 sm.create_session("t", None).unwrap();
1194 let sink = CollectingSink::new(false);
1195 let cancel = AtomicBool::new(false);
1196
1197 let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
1198 let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
1199 assert_eq!(res.status, RunStatus::Completed);
1200
1201 let tool_asst = sm
1205 .messages
1206 .iter()
1207 .find(|m| m.tool_calls.is_some())
1208 .expect("assistant with tool_calls present");
1209 assert_eq!(
1210 tool_asst.reasoning_content.as_deref(),
1211 Some("I need to echo")
1212 );
1213 assert_eq!(tool_asst.content, "let me echo");
1214 assert_eq!(tool_asst.tool_calls.as_ref().unwrap()[0].name, "echo");
1215
1216 let tool_msg = sm
1217 .messages
1218 .iter()
1219 .find(|m| m.role == "tool")
1220 .expect("tool result present");
1221 assert!(tool_msg.content.contains("echoed: hey"));
1222
1223 let final_asst = sm
1224 .messages
1225 .iter()
1226 .rev()
1227 .find(|m| m.role == "assistant")
1228 .expect("final assistant present");
1229 assert_eq!(final_asst.reasoning_content.as_deref(), Some("summarize"));
1230 assert_eq!(final_asst.content, "all done");
1231 }
1232
1233 #[tokio::test]
1237 async fn llm_request_has_exactly_one_copy_of_each_tool_message() {
1238 let llm = RecordingLlm {
1239 responses: Mutex::new(vec![
1240 LlmResponse {
1241 text: "let me echo".into(),
1242 tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
1243 reasoning_content: Some("I need to echo".into()),
1244 ..Default::default()
1245 },
1246 LlmResponse {
1247 text: "all done".into(),
1248 reasoning_content: Some("summarize".into()),
1249 ..Default::default()
1250 },
1251 ]),
1252 seen: Mutex::new(vec![]),
1253 };
1254 let mut reg = ToolRegistry::new();
1255 reg.register(Box::new(EchoTool));
1256 let cfg = AgentConfig::default();
1257 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1258
1259 let proj = tmp_proj();
1260 let mut sm = SessionManager::new(&proj);
1261 sm.create_session("t", None).unwrap();
1262 let sink = CollectingSink::new(false);
1263 let cancel = AtomicBool::new(false);
1264 let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
1265 let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
1266 assert_eq!(res.status, RunStatus::Completed);
1267
1268 let seen = llm.seen.lock().unwrap();
1269 assert_eq!(seen.len(), 2, "one request per model turn");
1270 let second = &seen[1];
1271
1272 let with_calls: Vec<&ChatMessage> =
1273 second.iter().filter(|m| m.tool_calls.is_some()).collect();
1274 assert_eq!(with_calls.len(), 1, "exactly one assistant tool_calls msg");
1275 let tool_msgs: Vec<&ChatMessage> = second.iter().filter(|m| m.role == "tool").collect();
1276 assert_eq!(
1277 tool_msgs.len(),
1278 1,
1279 "exactly one tool result (no duplication)"
1280 );
1281 assert_eq!(tool_msgs[0].tool_call_id.as_deref(), Some("c1"));
1282 assert_eq!(tool_msgs[0].content, "echoed: hey");
1283
1284 assert_eq!(with_calls[0].content, "let me echo");
1285 assert_eq!(
1286 with_calls[0].reasoning_content.as_deref(),
1287 Some("I need to echo")
1288 );
1289 let ai = second.iter().position(|m| m.tool_calls.is_some()).unwrap();
1290 assert_eq!(second[ai + 1].role, "tool", "result adjacent to its call");
1291
1292 assert_eq!(second.iter().filter(|m| m.role == "user").count(), 1);
1294 assert_eq!(second.iter().filter(|m| m.role == "system").count(), 1);
1295
1296 assert_eq!(sm.history_chat().len(), sm.messages.len());
1298 }
1299
1300 #[tokio::test]
1301 async fn respects_cancel_before_start() {
1302 let llm = ScriptedLlm {
1303 responses: Mutex::new(vec![]),
1304 };
1305 let reg = ToolRegistry::new();
1306 let cfg = AgentConfig::default();
1307 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1308 let proj = tmp_proj();
1309 let mut sm = SessionManager::new(&proj);
1310 sm.create_session("t", None).unwrap();
1311 let sink = CollectingSink::new(false);
1312 let cancel = AtomicBool::new(true);
1313 let res = loop_
1314 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1315 .await
1316 .unwrap();
1317 assert_eq!(res.status, RunStatus::Cancelled);
1318 }
1319
1320 #[tokio::test]
1321 async fn max_iterations_reached() {
1322 let responses: Vec<LlmResponse> = (0..5)
1324 .map(|_| LlmResponse {
1325 text: "again".into(),
1326 tool_calls: vec![tc("echo", "{\"text\":\"x\"}")],
1327 ..Default::default()
1328 })
1329 .collect();
1330 let llm = ScriptedLlm {
1331 responses: Mutex::new(responses),
1332 };
1333 let mut reg = ToolRegistry::new();
1334 reg.register(Box::new(EchoTool));
1335 let mut cfg = AgentConfig::default();
1336 cfg.max_iterations = 3;
1337 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1338 let proj = tmp_proj();
1339 let mut sm = SessionManager::new(&proj);
1340 sm.create_session("t", None).unwrap();
1341 let sink = CollectingSink::new(false);
1342 let cancel = AtomicBool::new(false);
1343 let res = loop_
1344 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1345 .await
1346 .unwrap();
1347 assert_eq!(res.status, RunStatus::MaxIterations);
1348 assert_eq!(res.iterations, 3);
1349 }
1350
1351 #[tokio::test]
1352 async fn domain_parse() {
1353 assert_eq!(domain_of("https://example.com/path"), "example.com");
1354 assert_eq!(domain_of("http://a.b.c/x/y"), "a.b.c");
1355 }
1356
1357 #[tokio::test]
1358 async fn classify_errors() {
1359 assert!(
1360 classify_api_error(&AacodeError::Api("HTTP 401 no".into())).contains("authentication")
1361 );
1362 assert!(classify_api_error(&AacodeError::Api("rate limit".into())).contains("quota"));
1363 assert!(classify_api_error(&AacodeError::Network("reset".into())).contains("Network"));
1364 assert!(classify_api_error(&AacodeError::Config("no key".into())).contains("Configuration"));
1365 }
1366
1367 struct FlakyLlm {
1369 remaining_failures: Mutex<u32>,
1370 }
1371 #[async_trait::async_trait]
1372 impl LlmClient for FlakyLlm {
1373 async fn chat_stream(
1374 &self,
1375 _m: &[ChatMessage],
1376 _t: &[Value],
1377 _e: &dyn EventSink,
1378 _c: &AtomicBool,
1379 ) -> Result<LlmResponse> {
1380 let mut n = self.remaining_failures.lock().unwrap();
1381 if *n > 0 {
1382 *n -= 1;
1383 return Err(AacodeError::Network("temporary reset".into()));
1384 }
1385 Ok(LlmResponse {
1386 text: "recovered".into(),
1387 ..Default::default()
1388 })
1389 }
1390 async fn validate(&self) -> Result<()> {
1391 Ok(())
1392 }
1393 }
1394
1395 #[tokio::test]
1396 async fn retries_transient_llm_errors() {
1397 let llm = FlakyLlm {
1398 remaining_failures: Mutex::new(1),
1399 };
1400 let reg = ToolRegistry::new();
1401 let mut cfg = AgentConfig::default();
1402 cfg.limits.max_retries = 3;
1403 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1404 let sink = CollectingSink::new(false);
1405 let cancel = AtomicBool::new(false);
1406 let proj = tmp_proj();
1408 let mut sm = SessionManager::new(&proj);
1409 sm.create_session("t", None).unwrap();
1410 let res = loop_
1411 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1412 .await
1413 .unwrap();
1414 assert_eq!(res.status, RunStatus::Completed);
1415 assert_eq!(res.final_text, "recovered");
1416 }
1417
1418 #[tokio::test]
1419 async fn truncated_response_not_completed() {
1420 let responses: Vec<LlmResponse> = (0..5)
1424 .map(|_| LlmResponse {
1425 text: "partial".into(),
1426 tool_calls: vec![],
1427 finish_reason: Some("length".to_string()),
1428 ..Default::default()
1429 })
1430 .collect();
1431 let llm = ScriptedLlm {
1432 responses: Mutex::new(responses),
1433 };
1434 let mut reg = ToolRegistry::new();
1435 reg.register(Box::new(EchoTool));
1436 let mut cfg = AgentConfig::default();
1437 cfg.max_iterations = 5;
1438 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1439 let proj = tmp_proj();
1440 let mut sm = SessionManager::new(&proj);
1441 sm.create_session("t", None).unwrap();
1442 let sink = CollectingSink::new(false);
1443 let cancel = AtomicBool::new(false);
1444 let res = loop_
1445 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1446 .await
1447 .unwrap();
1448 assert_eq!(
1451 res.status,
1452 RunStatus::MaxIterations,
1453 "truncated response must not be treated as completed"
1454 );
1455 }
1456
1457 #[tokio::test]
1458 async fn max_iterations_returns_last_text() {
1459 let responses: Vec<LlmResponse> = (0..5)
1460 .map(|i| LlmResponse {
1461 text: format!("iter {}", i),
1462 tool_calls: vec![tc("echo", "{\"text\":\"x\"}")],
1463 ..Default::default()
1464 })
1465 .collect();
1466 let llm = ScriptedLlm {
1467 responses: Mutex::new(responses),
1468 };
1469 let mut reg = ToolRegistry::new();
1470 reg.register(Box::new(EchoTool));
1471 let mut cfg = AgentConfig::default();
1472 cfg.max_iterations = 2;
1473 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1474 let proj = tmp_proj();
1475 let mut sm = SessionManager::new(&proj);
1476 sm.create_session("t", None).unwrap();
1477 let sink = CollectingSink::new(false);
1478 let cancel = AtomicBool::new(false);
1479 let res = loop_
1480 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1481 .await
1482 .unwrap();
1483 assert_eq!(res.status, RunStatus::MaxIterations);
1484 assert!(
1485 !res.final_text.is_empty(),
1486 "max_iterations must return last response text"
1487 );
1488 }
1489
1490 struct FailThenStreamLlm {
1494 remaining_failures: Mutex<u32>,
1495 }
1496 #[async_trait::async_trait]
1497 impl LlmClient for FailThenStreamLlm {
1498 async fn chat_stream(
1499 &self,
1500 _m: &[ChatMessage],
1501 _t: &[Value],
1502 emitter: &dyn EventSink,
1503 _c: &AtomicBool,
1504 ) -> Result<LlmResponse> {
1505 let mut n = self.remaining_failures.lock().unwrap();
1506 if *n > 0 {
1507 *n -= 1;
1508 return Err(AacodeError::Network("temporary reset".into()));
1509 }
1510 emitter.delta("thought", "streaming");
1511 emitter.seg("thinking", "reasoning done");
1512 emitter.seg("thought", "complete response");
1513 Ok(LlmResponse {
1514 text: "complete response".into(),
1515 ..Default::default()
1516 })
1517 }
1518 async fn validate(&self) -> Result<()> {
1519 Ok(())
1520 }
1521 }
1522
1523 struct EmitThenFailLlm {
1527 call: Mutex<u32>,
1528 }
1529 #[async_trait::async_trait]
1530 impl LlmClient for EmitThenFailLlm {
1531 async fn chat_stream(
1532 &self,
1533 _m: &[ChatMessage],
1534 _t: &[Value],
1535 emitter: &dyn EventSink,
1536 _c: &AtomicBool,
1537 ) -> Result<LlmResponse> {
1538 let mut call = self.call.lock().unwrap();
1539 *call += 1;
1540 if *call == 1 {
1541 emitter.delta("thought", "partial tok");
1543 emitter.seg("thinking", "partial reasoning");
1544 emitter.seg("thought", "partial text");
1545 return Err(AacodeError::Network("mid-stream disconnect".into()));
1546 }
1547 emitter.delta("thought", "complete");
1549 emitter.seg("thinking", "full reasoning");
1550 emitter.seg("thought", "final complete text");
1551 Ok(LlmResponse {
1552 text: "final complete text".into(),
1553 ..Default::default()
1554 })
1555 }
1556 async fn validate(&self) -> Result<()> {
1557 Ok(())
1558 }
1559 }
1560
1561 struct AlwaysFailLlm;
1563 #[async_trait::async_trait]
1564 impl LlmClient for AlwaysFailLlm {
1565 async fn chat_stream(
1566 &self,
1567 _m: &[ChatMessage],
1568 _t: &[Value],
1569 _emitter: &dyn EventSink,
1570 _c: &AtomicBool,
1571 ) -> Result<LlmResponse> {
1572 Err(AacodeError::Network("always fails".into()))
1573 }
1574 async fn validate(&self) -> Result<()> {
1575 Ok(())
1576 }
1577 }
1578
1579 struct TimeoutTool;
1581 #[async_trait::async_trait]
1582 impl Tool for TimeoutTool {
1583 fn schema(&self) -> ToolSchema {
1584 ToolSchema::new(
1585 "timeout_cmd",
1586 "timeout",
1587 vec![ToolParameter::new("x", ParamType::String, true, "x", &[])],
1588 )
1589 }
1590 async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
1591 Ok("command timed out after 10s".to_string())
1592 }
1593 }
1594
1595 struct CountingTool {
1597 count: std::sync::Arc<Mutex<u32>>,
1598 tool_name: &'static str,
1599 response: String,
1600 }
1601 impl CountingTool {
1602 fn new(name: &'static str, response: String) -> (Self, std::sync::Arc<Mutex<u32>>) {
1603 let count = std::sync::Arc::new(Mutex::new(0u32));
1604 (
1605 CountingTool {
1606 count: count.clone(),
1607 tool_name: name,
1608 response,
1609 },
1610 count,
1611 )
1612 }
1613 }
1614 #[async_trait::async_trait]
1615 impl Tool for CountingTool {
1616 fn schema(&self) -> ToolSchema {
1617 ToolSchema::new(self.tool_name, self.tool_name, vec![])
1618 }
1619 async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
1620 *self.count.lock().unwrap() += 1;
1621 Ok(self.response.clone())
1622 }
1623 }
1624
1625 #[tokio::test]
1626 async fn execute_with_retry_retries_on_network() {
1627 let mut reg = ToolRegistry::new();
1628 let (tool, counter) = CountingTool::new(
1629 "understand_image",
1630 r#"{"success":false,"error":"network error: vision: error sending request for url (https://api.moonshot.cn/v1/chat/completions)"}"#.to_string(),
1631 );
1632 reg.register(Box::new(tool));
1633 let mut cfg = AgentConfig::default();
1634 cfg.limits.max_retries = 3;
1635 cfg.max_iterations = 1;
1636 let responses = vec![LlmResponse {
1637 text: "try".into(),
1638 tool_calls: vec![tc(
1639 "understand_image",
1640 r#"{"image_path":"a.jpg","prompt":"desc"}"#,
1641 )],
1642 ..Default::default()
1643 }];
1644 let llm = ScriptedLlm {
1645 responses: Mutex::new(responses),
1646 };
1647 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1648 let proj = tmp_proj();
1649 let mut sm = SessionManager::new(&proj);
1650 sm.create_session("t", None).unwrap();
1651 let sink = CollectingSink::new(false);
1652 let cancel = AtomicBool::new(false);
1653 let _ = loop_
1654 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1655 .await;
1656 let invocations = *counter.lock().unwrap();
1657 assert!(
1658 invocations >= 2,
1659 "network error must trigger retry, got {invocations} invocations"
1660 );
1661 }
1662
1663 #[tokio::test]
1664 async fn execute_with_retry_does_not_retry_on_config_error() {
1665 let mut reg = ToolRegistry::new();
1666 let (tool, counter) = CountingTool::new(
1667 "understand_image",
1668 r#"{"success":false,"error":"config error: no multimodal model configured"}"#
1669 .to_string(),
1670 );
1671 reg.register(Box::new(tool));
1672 let mut cfg = AgentConfig::default();
1673 cfg.limits.max_retries = 3;
1674 cfg.max_iterations = 1;
1675 let responses = vec![LlmResponse {
1676 text: "try".into(),
1677 tool_calls: vec![tc(
1678 "understand_image",
1679 r#"{"image_path":"a.jpg","prompt":"desc"}"#,
1680 )],
1681 ..Default::default()
1682 }];
1683 let llm = ScriptedLlm {
1684 responses: Mutex::new(responses),
1685 };
1686 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1687 let proj = tmp_proj();
1688 let mut sm = SessionManager::new(&proj);
1689 sm.create_session("t", None).unwrap();
1690 let sink = CollectingSink::new(false);
1691 let cancel = AtomicBool::new(false);
1692 let _ = loop_
1693 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1694 .await;
1695 assert_eq!(
1696 *counter.lock().unwrap(),
1697 1,
1698 "config error must not trigger retry"
1699 );
1700 }
1701
1702 #[tokio::test]
1703 async fn execute_with_retry_timeout_retried() {
1704 let mut reg = ToolRegistry::new();
1705 let (tool, counter) =
1706 CountingTool::new("timeout_cmd", "command timed out after 10s".to_string());
1707 reg.register(Box::new(tool));
1708 let mut cfg = AgentConfig::default();
1709 cfg.limits.max_retries = 3;
1710 cfg.max_iterations = 1;
1711 let responses = vec![LlmResponse {
1712 text: "try".into(),
1713 tool_calls: vec![tc("timeout_cmd", r#"{"x":"1"}"#)],
1714 ..Default::default()
1715 }];
1716 let llm = ScriptedLlm {
1717 responses: Mutex::new(responses),
1718 };
1719 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1720 let proj = tmp_proj();
1721 let mut sm = SessionManager::new(&proj);
1722 sm.create_session("t", None).unwrap();
1723 let sink = CollectingSink::new(false);
1724 let cancel = AtomicBool::new(false);
1725 let _ = loop_
1726 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1727 .await;
1728 assert!(*counter.lock().unwrap() >= 2, "timeout must trigger retry");
1729 }
1730
1731 #[tokio::test]
1732 async fn execute_with_retry_connection_retried() {
1733 let mut reg = ToolRegistry::new();
1734 let (tool, counter) = CountingTool::new("fetch_url", "connection reset".to_string());
1735 reg.register(Box::new(tool));
1736 let mut cfg = AgentConfig::default();
1737 cfg.limits.max_retries = 3;
1738 cfg.max_iterations = 1;
1739 let responses = vec![LlmResponse {
1740 text: "try".into(),
1741 tool_calls: vec![tc("fetch_url", r#"{"url":"http://x"}"#)],
1742 ..Default::default()
1743 }];
1744 let llm = ScriptedLlm {
1745 responses: Mutex::new(responses),
1746 };
1747 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1748 let proj = tmp_proj();
1749 let mut sm = SessionManager::new(&proj);
1750 sm.create_session("t", None).unwrap();
1751 let sink = CollectingSink::new(false);
1752 let cancel = AtomicBool::new(false);
1753 let _ = loop_
1754 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1755 .await;
1756 assert!(
1757 *counter.lock().unwrap() >= 2,
1758 "connection must trigger retry"
1759 );
1760 }
1761
1762 #[tokio::test]
1763 async fn execute_with_retry_respects_cancel() {
1764 let responses = vec![LlmResponse {
1765 text: "try".into(),
1766 tool_calls: vec![tc("timeout_cmd", "{\"x\":\"1\"}")],
1767 ..Default::default()
1768 }];
1769 let llm = ScriptedLlm {
1770 responses: Mutex::new(responses),
1771 };
1772 let mut reg = ToolRegistry::new();
1773 reg.register(Box::new(TimeoutTool));
1774 let mut cfg = AgentConfig::default();
1775 cfg.limits.max_retries = 10;
1776 cfg.max_iterations = 5;
1777 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1778 let proj = tmp_proj();
1779 let mut sm = SessionManager::new(&proj);
1780 sm.create_session("t", None).unwrap();
1781 let sink = CollectingSink::new(false);
1782 let cancel = AtomicBool::new(true); let res = loop_
1784 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1785 .await
1786 .unwrap();
1787 assert_eq!(res.status, RunStatus::Cancelled);
1790 }
1791
1792 #[tokio::test]
1793 async fn max_context_tokens_enforced_and_errors() {
1794 let llm = ScriptedLlm {
1795 responses: Mutex::new(vec![LlmResponse {
1796 text: "done".into(),
1797 ..Default::default()
1798 }]),
1799 };
1800 let reg = ToolRegistry::new();
1801 let mut cfg = AgentConfig::default();
1802 cfg.max_iterations = 5;
1803 cfg.context.max_context_tokens = 10;
1804 cfg.context.compact_trigger_tokens = 5;
1805 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1806
1807 let proj = tmp_proj();
1808 let mut sm = SessionManager::new(&proj);
1809 sm.create_session("t", None).unwrap();
1810 let sink = CollectingSink::new(false);
1811 let cancel = AtomicBool::new(false);
1812
1813 let msgs = vec![
1814 ChatMessage::system("you are a helpful assistant who provides detailed responses"),
1815 ChatMessage::user("do something"),
1816 ];
1817 let res = loop_.run(msgs, &mut sm, &sink, &cancel).await;
1818 assert!(
1819 res.is_err(),
1820 "context over max must error, not silently proceed"
1821 );
1822 let msg = format!("{}", res.err().unwrap());
1823 assert!(
1824 msg.contains("context too large") || msg.contains("exceeds limit"),
1825 "error must mention context size: {msg}"
1826 );
1827 }
1828
1829 #[tokio::test]
1832 async fn pipe_mode_retry_buffers_and_replays() {
1833 let llm = FailThenStreamLlm {
1837 remaining_failures: Mutex::new(1),
1838 };
1839 let reg = ToolRegistry::new();
1840 let mut cfg = AgentConfig::default();
1841 cfg.limits.max_retries = 3;
1842 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1843
1844 let proj = tmp_proj();
1845 let mut sm = SessionManager::new(&proj);
1846 sm.create_session("t", None).unwrap();
1847 let sink = CollectingSink::new(false);
1849 let cancel = AtomicBool::new(false);
1850
1851 let res = loop_
1852 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1853 .await
1854 .unwrap();
1855 assert_eq!(res.status, RunStatus::Completed);
1856 assert_eq!(res.final_text, "complete response");
1857
1858 let lines = sink.lines();
1859 assert!(
1861 lines
1862 .iter()
1863 .any(|l| l.contains(r#""type":"delta""#) && l.contains("streaming")),
1864 "delta event must be replayed to real sink"
1865 );
1866 assert!(
1867 lines
1868 .iter()
1869 .any(|l| l.contains(r#""type":"seg_content""#) && l.contains("complete response")),
1870 "seg_content event must be replayed to real sink"
1871 );
1872 }
1873
1874 #[tokio::test]
1875 async fn tty_mode_never_buffers_retries() {
1876 let llm = FailThenStreamLlm {
1878 remaining_failures: Mutex::new(1),
1879 };
1880 let reg = ToolRegistry::new();
1881 let mut cfg = AgentConfig::default();
1882 cfg.limits.max_retries = 3;
1883 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1884
1885 let proj = tmp_proj();
1886 let mut sm = SessionManager::new(&proj);
1887 sm.create_session("t", None).unwrap();
1888 let sink = CollectingSink::new(true);
1890 let cancel = AtomicBool::new(false);
1891
1892 let res = loop_
1893 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1894 .await
1895 .unwrap();
1896 assert_eq!(res.status, RunStatus::Completed);
1897
1898 let lines = sink.lines();
1903 assert!(
1905 lines.iter().any(|l| l == "streaming"),
1906 "TTY delta must emit raw text: {lines:?}"
1907 );
1908 }
1909
1910 #[tokio::test]
1911 async fn multiple_retries_only_last_success_replayed() {
1912 let llm = FailThenStreamLlm {
1915 remaining_failures: Mutex::new(2),
1916 };
1917 let reg = ToolRegistry::new();
1918 let mut cfg = AgentConfig::default();
1919 cfg.limits.max_retries = 4;
1920 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1921
1922 let proj = tmp_proj();
1923 let mut sm = SessionManager::new(&proj);
1924 sm.create_session("t", None).unwrap();
1925 let sink = CollectingSink::new(false);
1926 let cancel = AtomicBool::new(false);
1927
1928 let res = loop_
1929 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1930 .await
1931 .unwrap();
1932 assert_eq!(res.status, RunStatus::Completed);
1933
1934 let lines = sink.lines();
1935 let seg_count = lines
1939 .iter()
1940 .filter(|l| l.contains(r#""type":"seg_content""#) && l.contains("complete response"))
1941 .count();
1942 assert_eq!(
1943 seg_count, 1,
1944 "only one seg_content for the final thought; found {seg_count}. lines: {lines:?}"
1945 );
1946 }
1947
1948 #[tokio::test]
1949 async fn first_attempt_streams_directly_leak_is_allowed() {
1950 let llm = EmitThenFailLlm {
1956 call: Mutex::new(0),
1957 };
1958 let reg = ToolRegistry::new();
1959 let mut cfg = AgentConfig::default();
1960 cfg.limits.max_retries = 3;
1961 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
1962
1963 let proj = tmp_proj();
1964 let mut sm = SessionManager::new(&proj);
1965 sm.create_session("t", None).unwrap();
1966 let sink = CollectingSink::new(false);
1967 let cancel = AtomicBool::new(false);
1968
1969 let res = loop_
1970 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1971 .await
1972 .unwrap();
1973 assert_eq!(res.status, RunStatus::Completed);
1974 assert_eq!(res.final_text, "final complete text");
1975
1976 let lines = sink.lines();
1977 assert!(
1979 lines.iter().any(|l| l.contains("partial text")),
1980 "first attempt leaked: {lines:?}"
1981 );
1982 assert!(
1984 lines.iter().any(|l| l.contains("final complete text")),
1985 "retry must replay: {lines:?}"
1986 );
1987 }
1990
1991 #[tokio::test]
1992 async fn buffered_retry_replays_all_event_types() {
1993 let llm = FailThenStreamLlm {
1996 remaining_failures: Mutex::new(1),
1997 };
1998 let reg = ToolRegistry::new();
1999 let mut cfg = AgentConfig::default();
2000 cfg.limits.max_retries = 3;
2001 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
2002
2003 let proj = tmp_proj();
2004 let mut sm = SessionManager::new(&proj);
2005 sm.create_session("t", None).unwrap();
2006 let sink = CollectingSink::new(false);
2007 let cancel = AtomicBool::new(false);
2008
2009 let res = loop_
2010 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
2011 .await
2012 .unwrap();
2013 assert_eq!(res.status, RunStatus::Completed);
2014
2015 let lines = sink.lines();
2016 assert!(
2018 lines
2019 .iter()
2020 .any(|l| l.contains(r#""type":"delta""#) && l.contains("streaming")),
2021 "delta missing: {lines:?}"
2022 );
2023 assert!(
2025 lines.iter().any(|l| l.contains(r#""type":"seg_content""#) && l.contains(r#""seg":"thinking""#)),
2026 "thinking seg missing: {lines:?}"
2027 );
2028 assert!(
2030 lines
2031 .iter()
2032 .any(|l| l.contains(r#""type":"seg_content""#) && l.contains(r#""seg":"thought""#)),
2033 "thought seg missing: {lines:?}"
2034 );
2035 assert!(
2038 lines.iter().any(|l| l.contains(r#""type":"done""#)),
2039 "done event missing: {lines:?}"
2040 );
2041 }
2042
2043 #[tokio::test]
2044 async fn buffered_retry_respects_cancel() {
2045 let llm = FailThenStreamLlm {
2048 remaining_failures: Mutex::new(0),
2049 };
2050 let reg = ToolRegistry::new();
2051 let mut cfg = AgentConfig::default();
2052 cfg.limits.max_retries = 3;
2053 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
2054
2055 let proj = tmp_proj();
2056 let mut sm = SessionManager::new(&proj);
2057 sm.create_session("t", None).unwrap();
2058 let sink = CollectingSink::new(false);
2059 let cancel = AtomicBool::new(true);
2060
2061 let res = loop_
2062 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
2063 .await
2064 .unwrap();
2065 assert_eq!(res.status, RunStatus::Cancelled);
2066 }
2067
2068 #[tokio::test]
2069 async fn cancel_during_retry_backoff_stops_cleanly() {
2070 use std::sync::Arc;
2071 let llm = FailThenStreamLlm {
2072 remaining_failures: Mutex::new(2),
2073 };
2074 let reg = ToolRegistry::new();
2075 let mut cfg = AgentConfig::default();
2076 cfg.limits.max_retries = 5;
2077 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
2078
2079 let proj = tmp_proj();
2080 let mut sm = SessionManager::new(&proj);
2081 sm.create_session("t", None).unwrap();
2082 let sink = CollectingSink::new(false);
2083 let cancel = Arc::new(AtomicBool::new(false));
2084
2085 let cancel2 = Arc::clone(&cancel);
2086 std::thread::spawn(move || {
2087 std::thread::sleep(std::time::Duration::from_millis(500));
2089 cancel2.store(true, Ordering::SeqCst);
2090 });
2091
2092 let res = loop_
2093 .run(
2094 vec![ChatMessage::user("x")],
2095 &mut sm,
2096 &sink,
2097 cancel.as_ref(),
2098 )
2099 .await
2100 .unwrap();
2101 assert_eq!(
2102 res.status,
2103 RunStatus::Cancelled,
2104 "cancel during retry backoff must return Cancelled"
2105 );
2106 }
2107
2108 #[tokio::test]
2109 async fn exhausted_retries_buffers_discarded_no_events() {
2110 let llm = AlwaysFailLlm;
2113 let reg = ToolRegistry::new();
2114 let mut cfg = AgentConfig::default();
2115 cfg.limits.max_retries = 3;
2116 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
2117
2118 let proj = tmp_proj();
2119 let mut sm = SessionManager::new(&proj);
2120 sm.create_session("t", None).unwrap();
2121 let sink = CollectingSink::new(false);
2122 let cancel = AtomicBool::new(false);
2123
2124 let res = loop_
2125 .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
2126 .await
2127 .unwrap();
2128 assert!(matches!(res.status, RunStatus::Error(_)));
2129 let lines = sink.lines();
2131 assert!(
2132 lines.iter().any(|l| l.contains(r#""type":"error""#)),
2133 "error event must be emitted: {lines:?}"
2134 );
2135 let deltas = lines
2137 .iter()
2138 .filter(|l| l.contains(r#""type":"delta""#))
2139 .count();
2140 let segs = lines
2141 .iter()
2142 .filter(|l| l.contains(r#""type":"seg_content""#))
2143 .count();
2144 assert_eq!(deltas, 0, "no delta events from failed retries: {lines:?}");
2145 assert_eq!(
2146 segs, 0,
2147 "no seg_content events from failed retries: {lines:?}"
2148 );
2149 }
2150
2151 struct BlockingTool;
2153 #[async_trait::async_trait]
2154 impl Tool for BlockingTool {
2155 fn schema(&self) -> ToolSchema {
2156 ToolSchema::new("blocking_tool", "never returns", vec![])
2157 }
2158 async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
2159 futures::future::pending::<()>().await;
2160 Ok(String::new())
2161 }
2162 }
2163
2164 #[tokio::test]
2167 async fn cancel_interrupts_a_blocking_tool() {
2168 let llm = ScriptedLlm {
2169 responses: Mutex::new(vec![LlmResponse {
2170 text: "work".into(),
2171 tool_calls: vec![tc("blocking_tool", "{}")],
2172 ..Default::default()
2173 }]),
2174 };
2175 let mut reg = ToolRegistry::new();
2176 reg.register(Box::new(BlockingTool));
2177 let cfg = AgentConfig::default();
2178 let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
2179 let proj = tmp_proj();
2180 let mut sm = SessionManager::new(&proj);
2181 sm.create_session("t", None).unwrap();
2182 let sink = CollectingSink::new(false);
2183 let cancel = Arc::new(AtomicBool::new(false));
2184 let c2 = cancel.clone();
2185 std::thread::spawn(move || {
2186 std::thread::sleep(std::time::Duration::from_millis(200));
2187 c2.store(true, Ordering::SeqCst);
2188 });
2189 let start = std::time::Instant::now();
2190 let res = loop_
2191 .run(
2192 vec![ChatMessage::user("x")],
2193 &mut sm,
2194 &sink,
2195 cancel.as_ref(),
2196 )
2197 .await
2198 .unwrap();
2199 assert!(
2200 matches!(res.status, RunStatus::Cancelled),
2201 "expected Cancelled, got {:?}",
2202 res.status
2203 );
2204 assert!(
2205 start.elapsed() < std::time::Duration::from_secs(3),
2206 "cancel should be prompt, took {:?}",
2207 start.elapsed()
2208 );
2209 }
2210}
2211
2212async fn cancel_checker(cancel: &AtomicBool, notify: &Notify) {
2215 use std::time::Duration;
2216 loop {
2217 tokio::select! {
2218 _ = notify.notified() => {},
2219 _ = tokio::time::sleep(Duration::from_secs(1)) => {},
2220 }
2221 if cancel.load(Ordering::SeqCst) {
2222 return;
2223 }
2224 }
2225}
2226
2227fn brief(e: &crate::error::AacodeError) -> String {
2229 let s = e.to_string();
2230 let first = s.lines().next().unwrap_or(&s);
2231 let mut out: String = first.chars().take(60).collect();
2232 if first.chars().count() > 60 {
2233 out.push('…');
2234 }
2235 out
2236}