Skip to main content

aacode_rs/agent/
react_loop.rs

1// Copyright (c) 2026 xiefujin <490021684@qq.com>
2// Licensed under GPL-3.0, see LICENSE file for full license terms.
3
4//! The ReAct loop — iterate model calls + tool execution until completion.
5//!
6//! Ported from Python `core/react_loop.py` (core control flow). Completion is
7//! detected exactly as in Python: the model returns no tool_calls → the final
8//! text is the summary and the loop ends.
9
10use 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/// Outcome of a run.
26#[derive(Debug, Clone, PartialEq)]
27pub enum RunStatus {
28    Completed,
29    MaxIterations,
30    Cancelled,
31    Error(String),
32}
33
34impl RunStatus {
35    /// Stable wire name for the status (used by the `done` event and the
36    /// terminal JSON returned by the FFI `wait`).
37    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    /// Serialize the terminal outcome as JSON (shared by `done` enrichment and
55    /// the FFI `wait` return value).
56    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
65/// The ReAct loop driver. Borrows the collaborators for the duration of a run.
66pub struct ReactLoop<'a> {
67    pub llm: &'a dyn LlmClient,
68    pub tools: &'a ToolRegistry,
69    pub config: &'a AgentConfig,
70    /// Preformatted native tools for the active gateway.
71    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    /// Run the loop. `messages` starts with the system prompt + history + task.
90    /// New messages are persisted to `session` incrementally.
91    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        // Retrieval cutoff (replaces summary compaction): stored messages stay
101        // complete; only the model view is trimmed. The cutoff is frozen so the
102        // request prefix stays byte-stable for provider prompt caching.
103        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        // Usage telemetry (token counts only; never prompt/response text).
107        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        // Track the last response text so MaxIterations can return it.
111        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            // Build the model view (does not mutate stored messages). Older
124            // history is omitted at a frozen cutoff and stays retrievable via
125            // the history tools.
126            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            // Enforce hard context limit: if even the compacted view exceeds
132            // max_context_tokens the request would be rejected or silently
133            // truncated by the API — fail fast instead.
134            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            // Only re-sanitize after compaction — uncompacted views were
143            // already sanitized in MainAgent (history repair) and the ReAct
144            // loop never creates broken tool_calls/tool pairings on its own.
145            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            // Record provider usage for this call (JSONL; no text). Missing
174            // usage (gateways that don't report it) is recorded too, so the
175            // estimator can be calibrated against real counts later.
176            {
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                // Request manifest (no prompt text): the actual exposed tool
193                // names plus the retrieval cutoff. Makes cache misses and tool
194                // drift diagnosable from the usage log alone.
195                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            // Completion: no tool calls.
227            if resp.tool_calls.is_empty() {
228                // If the response was truncated (max_tokens / length), the
229                // model may have intended to output tool calls that got cut
230                // off. Inject a continuation so the model can keep going.
231                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                // Flush before the terminal event so a host re-reading the
257                // session right after `done` sees the final assistant message.
258                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            // Stale-loop detection warning (fetch_url stale domains).
268            // Persisted so cross-execute() prefix stays stable.
269            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            // Append the assistant message carrying tool_calls.
276            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            // Execute tool calls. A batch containing a non-reentrant lane
284            // (shell/browser/python/delegate) runs sequentially; otherwise
285            // independent calls run concurrently. `join_all` preserves input
286            // order, so tool results are always submitted in the model's
287            // original call order regardless of completion timing.
288            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; // reference from outer scope
297
298                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                // Single tool: sequential path
373                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            // Context growth check (informational; compact view is built each iter).
418            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    /// Call the model with retry on transient (network/5xx/timeout) errors.
436    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            // On retry in pipe mode (mobile FFI), buffer output in a
450            // CollectingSink so partial content from a failed stream never
451            // leaks to the client (avoids "回退" — segment content shrinking
452            // mid-stream). Only replay when the stream succeeds.
453            // First attempt still streams directly for real-time UX.
454            // TTY mode (CLI) never buffers — it has no segment tracking
455            // and stale partial text is harmless.
456            let use_buffered = attempt > 0 && !emitter.is_tty();
457            // Compute result first — LlmResponse is self-contained (all
458            // owned), so the CollectingSink borrow ends before we move it.
459            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                    // Surface the retry to the UI (tool_progress renders as a
491                    // live status line) — a silent backoff looks like a hang,
492                    // especially on flaky mobile networks.
493                    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                    // exponential backoff: 1s, 2s, 4s, ... with ±25% jitter
505                    // so concurrent retries don't form a thundering herd.
506                    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); // 250ms * delay
513                        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    /// Execute a tool, retrying transient failures per config.limits.max_retries.
528    async fn execute_with_retry(
529        &self,
530        name: &str,
531        args: Value,
532        _emitter: &dyn EventSink,
533        cancel: &AtomicBool,
534    ) -> String {
535        // Read-only mode: block mutating/executing tools before any dispatch.
536        // The structured observation tells the model why (and is classified as
537        // a tool error by `tool_status`).
538        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            // Race the tool against cancellation so a long/blocking tool (e.g.
557            // a 60s `run_shell`) does not delay a cancel. The tool future is
558            // dropped on cancel; a `spawn_blocking` command may finish in the
559            // background but the agent stops promptly (and releases the session
560            // guard). The local Notify is never notified → 1s polling fallback.
561            let notify = Notify::new();
562            let obs = tokio::select! {
563                obs = self.tools.execute(name, args.clone(), cancel) => obs,
564                _ = cancel_checker(cancel, &notify) => 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
581/// Classify a tool observation into a structured `(status, error_code)` pair
582/// for the additive `tool_status` event.
583fn 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
615/// Classify an LLM/API error into a friendly, actionable message.
616fn 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/// Tracks fetch_url results per domain to detect stale loops (mirrors Python).
633#[derive(Default)]
634struct StaleTracker {
635    by_domain: HashMap<String, Vec<bool>>, // has_content flags
636    warned: std::collections::HashSet<String>,
637}
638
639impl StaleTracker {
640    fn detect(&mut self, tool_calls: &[ToolCall]) -> Option<String> {
641        // Only meaningful when the model is issuing fetch_url calls.
642        if !tool_calls.iter().any(|tc| tc.name == "fetch_url") {
643            return None;
644        }
645        // Find a domain with 3 recent no-content fetches that we haven't warned about.
646        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        // strip tags, check readable length
670        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    // naive scheme://host/... parse
686    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    // A scripted LLM that returns a queue of responses.
705    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    // Scripted LLM that also records every request's messages.
731    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    /// Same-lane tools must not overlap (deterministic event log, no timing).
831    #[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, &reg, &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    /// Independent tools run concurrently, but results are still submitted in
889    /// the model's original call order.
890    #[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, &reg, &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    /// Read-only mode blocks mutating tools before dispatch but allows reads.
953    #[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, &reg, &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    /// End-to-end: the loop must send a compacted view and persist the
1055    /// checkpoint, leaving the full message list untouched.
1056    #[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, &reg, &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, &reg, &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, &reg, &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        // observation emitted
1163        assert!(sink.lines().iter().any(|l| l.contains("echoed: hey")));
1164        // session persisted assistant(tool_calls) + tool + final assistant
1165        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, &reg, &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        // The turn is persisted as plain messages only: assistant(tool_calls)
1202        // carries reasoning + content (~ thinking/thought/action), the tool
1203        // message carries the observation — no duplicated render segments.
1204        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    /// The exact messages handed to the model must contain each tool call and
1234    /// each tool result exactly once (no duplication), adjacent, with no missing
1235    /// result and no stray duplicated system/user rows.
1236    #[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, &reg, &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        // No duplicated user/system rows slipped in.
1293        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        // Persisted history maps 1:1 (no duplicate rows on disk either).
1297        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, &reg, &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        // Always returns a tool call → never completes.
1323        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, &reg, &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    // An LLM that fails N times with a retryable error, then succeeds.
1368    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, &reg, &cfg, vec![]);
1404        let sink = CollectingSink::new(false);
1405        let cancel = AtomicBool::new(false);
1406        // Directly exercise chat_with_retry via a completing run.
1407        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        // Truncated response with no tool calls → injects "continue" →
1421        // needs another LLM response. We provide a chain so the loop can
1422        // exhaust iterations without crashing on empty ScriptedLlm.
1423        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, &reg, &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        // Truncated response with no tool calls should NOT complete — it
1449        // should loop until max_iterations and then return MaxIterations.
1450        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, &reg, &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    // ── LLMs for retry / buffering tests ──────────────────────────────
1491
1492    /// Fails N times with a retryable error, then succeeds with stream events.
1493    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    /// Emits events on attempt 0 then fails, succeeds on attempt 1 with
1524    /// different content. Simulates a mid-stream disconnect followed by
1525    /// a successful retry.
1526    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                // Simulate partial stream output before a mid-stream error.
1542                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            // Retry: completely different (correct) content.
1548            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    /// Fails every call (never succeeds). Used for exhausted-retry tests.
1562    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    /// A tool that always returns "timeout" observation — triggers retry.
1580    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    /// Counts how many times its tool was invoked.
1596    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, &reg, &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, &reg, &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, &reg, &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, &reg, &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, &reg, &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); // pre-set cancel
1783        let res = loop_
1784            .run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
1785            .await
1786            .unwrap();
1787        // Must not retry forever — cancel check in first iteration returns
1788        // Cancelled before tool execution even starts.
1789        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, &reg, &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    // ── retry buffering tests ────────────────────────────────────────
1830
1831    #[tokio::test]
1832    async fn pipe_mode_retry_buffers_and_replays() {
1833        // FlakyLlm fails once then returns Ok (no events emitted).
1834        // With buffering, the emitter must only see output from the
1835        // retry attempt — nothing leaked from the failed first attempt.
1836        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, &reg, &cfg, vec![]);
1843
1844        let proj = tmp_proj();
1845        let mut sm = SessionManager::new(&proj);
1846        sm.create_session("t", None).unwrap();
1847        // Pipe mode (non-TTY) — buffering should engage on retry.
1848        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        // The retry emitted delta + seg events. Verify they are replayed.
1860        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        // In TTY mode, attempts > 0 still use the real emitter directly.
1877        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, &reg, &cfg, vec![]);
1884
1885        let proj = tmp_proj();
1886        let mut sm = SessionManager::new(&proj);
1887        sm.create_session("t", None).unwrap();
1888        // TTY mode — buffering should NOT engage.
1889        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        // The delta("thought", ...) goes through token_line → emit_line in
1899        // TTY mode. The seg("thought", ...) default impl emits JSONL
1900        // regardless of TTY flag on CollectingSink. So we should find
1901        // both raw text and JSONL in the output.
1902        let lines = sink.lines();
1903        // In TTY mode, delta emits raw text via token_line.
1904        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        // Fail twice, succeed on third. Each failure must be discarded;
1913        // only the successful attempt's events reach the real emitter.
1914        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, &reg, &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        // Each retry's buffer was discarded on failure. Only the final
1936        // success should appear. The exact content comes from attempt 2
1937        // (index 2, the third call).
1938        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        // Attempt 0 emits events DIRECTLY to real emitter then fails.
1951        // Attempt 1 (buffered) retries. The real emitter will contain
1952        // BOTH the leaked attempt-0 events AND the replayed attempt-1
1953        // events. This is by design: real-time streaming on first
1954        // attempt trades perfect isolation for live UI feedback.
1955        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, &reg, &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        // Attempt 0 leaked "partial text" directly to the real emitter.
1978        assert!(
1979            lines.iter().any(|l| l.contains("partial text")),
1980            "first attempt leaked: {lines:?}"
1981        );
1982        // Attempt 1 replayed "final complete text" to the real emitter.
1983        assert!(
1984            lines.iter().any(|l| l.contains("final complete text")),
1985            "retry must replay: {lines:?}"
1986        );
1987        // Both sets appear, but retry content is the authoritative final
1988        // seg_content that the iOS side should use.
1989    }
1990
1991    #[tokio::test]
1992    async fn buffered_retry_replays_all_event_types() {
1993        // Verify that the buffered CollectingSink captures and replays
1994        // all event types: delta, seg_content (thinking + thought).
1995        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, &reg, &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        // Delta event replayed.
2017        assert!(
2018            lines
2019                .iter()
2020                .any(|l| l.contains(r#""type":"delta""#) && l.contains("streaming")),
2021            "delta missing: {lines:?}"
2022        );
2023        // Thinking seg_content replayed.
2024        assert!(
2025            lines.iter().any(|l| l.contains(r#""type":"seg_content""#) && l.contains(r#""seg":"thinking""#)),
2026            "thinking seg missing: {lines:?}"
2027        );
2028        // Thought seg_content replayed.
2029        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        // done event from the react_loop (emitted after chat_with_retry
2036        // returns, not from the LLM).
2037        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        // Pre-set cancel: the cancel check at the top of each iteration
2046        // (in ReactLoop::run) catches it before chat_with_retry is called.
2047        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, &reg, &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, &reg, &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            // Wait for the first failure + backoff sleep to start.
2088            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        // When all retries fail, no buffered content should reach the
2111        // real emitter (each buffer was discarded on failure).
2112        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, &reg, &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        // The react_loop emits an error event on the real emitter.
2130        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        // But there should be no delta/seg_content from retries.
2136        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    /// A tool that never returns (like a long blocking `run_shell`).
2152    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    /// Cancel must interrupt a blocking tool promptly (the loop must not wait
2165    /// for it), which also lets the session guard be released.
2166    #[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, &reg, &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
2212/// Returns a future that completes when the cancel flag is set.
2213/// Uses Notify for immediate wakeup, with 1s polling fallback.
2214async 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
2227/// Short one-line error description for UI status lines.
2228fn 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}