use crate::agent::compact::estimate_messages_tokens;
use crate::agent::retrieval::RetrievalState;
use crate::config::AgentConfig;
use crate::error::{AacodeError, Result};
use crate::llm::types::{ChatMessage, ToolCall};
use crate::llm::LlmClient;
use crate::session::{SessionManager, SessionMessage};
use crate::stream::{CollectingSink, EventSink};
use crate::tools::ToolRegistry;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::Notify;
#[derive(Debug, Clone, PartialEq)]
pub enum RunStatus {
Completed,
MaxIterations,
Cancelled,
Error(String),
}
impl RunStatus {
pub fn as_str(&self) -> &'static str {
match self {
RunStatus::Completed => "completed",
RunStatus::MaxIterations => "max_iterations",
RunStatus::Cancelled => "cancelled",
RunStatus::Error(_) => "error",
}
}
}
pub struct RunResult {
pub status: RunStatus,
pub iterations: u32,
pub final_text: String,
}
impl RunResult {
pub fn to_result_json(&self) -> serde_json::Value {
serde_json::json!({
"status": self.status.as_str(),
"iterations": self.iterations,
"final_text": self.final_text,
})
}
}
pub struct ReactLoop<'a> {
pub llm: &'a dyn LlmClient,
pub tools: &'a ToolRegistry,
pub config: &'a AgentConfig,
pub native_tools: Vec<Value>,
}
impl<'a> ReactLoop<'a> {
pub fn new(
llm: &'a dyn LlmClient,
tools: &'a ToolRegistry,
config: &'a AgentConfig,
native_tools: Vec<Value>,
) -> Self {
ReactLoop {
llm,
tools,
config,
native_tools,
}
}
pub async fn run(
&self,
mut messages: Vec<ChatMessage>,
session: &mut SessionManager,
emitter: &dyn EventSink,
cancel: &AtomicBool,
) -> Result<RunResult> {
let cancel_notify = Arc::new(Notify::new());
let mut stale = StaleTracker::default();
let project_path = session.project_path();
let mut retrieval = RetrievalState::from_value(session.context_checkpoint());
let session_id = session.current_session_id.clone().unwrap_or_default();
let usage_tracker = crate::usage::UsageTracker::new(&session.project_path());
let run_uuid = uuid::Uuid::new_v4().simple().to_string();
let run_id = &run_uuid[..12.min(run_uuid.len())];
let mut last_text = String::new();
for iteration in 0..self.config.max_iterations {
if cancel.load(Ordering::SeqCst) {
emitter.error("cancelled");
return Ok(RunResult {
status: RunStatus::Cancelled,
iterations: iteration,
final_text: String::new(),
});
}
let (mut view, compacted, tokens, checkpoint_changed) =
retrieval.model_view(&messages, &self.config.context, &project_path, &session_id);
if checkpoint_changed {
let _ = session.set_context_checkpoint(retrieval.to_value());
}
if tokens > self.config.context.max_context_tokens {
return Err(AacodeError::Api(format!(
"context too large: {} tokens exceeds limit of {} ({} messages). Try reducing history or restarting the session.",
tokens,
self.config.context.max_context_tokens,
messages.len(),
)));
}
if compacted {
crate::agent::sanitize::sanitize_history(&mut view);
}
let resp = match self
.chat_with_retry(&view, emitter, cancel, &cancel_notify)
.await
{
Ok(r) => r,
Err(AacodeError::Cancelled) => {
emitter.error("cancelled");
return Ok(RunResult {
status: RunStatus::Cancelled,
iterations: iteration,
final_text: String::new(),
});
}
Err(e) => {
let classified = classify_api_error(&e);
emitter.error(&classified);
return Ok(RunResult {
status: RunStatus::Error(classified),
iterations: iteration,
final_text: String::new(),
});
}
};
{
let gateway = match self.config.model.gateway {
crate::config::Gateway::Anthropic => "anthropic",
crate::config::Gateway::Openai => "openai",
};
let mut event = crate::usage::usage_event(
&session_id,
run_id,
iteration + 1,
gateway,
&self.config.model.name,
tokens,
view.len(),
compacted,
resp.usage.as_ref(),
);
let tools: Vec<String> = self
.native_tools
.iter()
.filter_map(|t| {
t.get("function")
.and_then(|f| f.get("name"))
.and_then(|n| n.as_str())
.map(str::to_string)
})
.collect();
event["tools"] = serde_json::json!(tools);
if let Some(cp) = retrieval.checkpoint.as_ref() {
event["retrieval_cutoff"] = serde_json::json!(cp.cutoff_index);
event["retrieval_stub_chars"] = serde_json::json!(cp.stub.chars().count());
}
if usage_tracker.record(&event) && usage_tracker.verbose {
if let Some(u) = &resp.usage {
eprintln!(
"📊 usage iter {}: input={} (cache {:.1}%) output={} reasoning={}",
iteration + 1,
u.input_tokens,
u.cache_hit_rate() * 100.0,
u.output_tokens,
u.reasoning_tokens
);
}
}
}
last_text = resp.text.clone();
if resp.tool_calls.is_empty() {
if resp.is_truncated() {
let assistant = ChatMessage {
role: "assistant".into(),
content: last_text.clone(),
tool_calls: None,
tool_call_id: None,
reasoning_content: resp.reasoning_content.clone(),
};
messages.push(assistant.clone());
let _ = session.add_message(SessionMessage::from_chat(&assistant));
let cont = ChatMessage::user("continue");
messages.push(cont.clone());
let _ = session.add_message(SessionMessage::from_chat(&cont));
continue;
}
let assistant = ChatMessage {
role: "assistant".into(),
content: last_text.clone(),
tool_calls: None,
tool_call_id: None,
reasoning_content: resp.reasoning_content.clone(),
};
messages.push(assistant.clone());
let _ = session.add_message(SessionMessage::from_chat(&assistant));
let _ = session.flush();
emitter.done_result(&session_id, "completed", iteration + 1, &last_text);
return Ok(RunResult {
status: RunStatus::Completed,
iterations: iteration + 1,
final_text: last_text,
});
}
if let Some(w) = stale.detect(&resp.tool_calls) {
let warn_msg = ChatMessage::system(format!("[SYSTEM WARNING]: {w}"));
let _ = session.add_message(SessionMessage::from_chat(&warn_msg));
messages.push(warn_msg);
}
let assistant =
ChatMessage::assistant_with_tools(resp.text.clone(), resp.tool_calls.clone());
let mut assistant_msg = assistant.clone();
assistant_msg.reasoning_content = resp.reasoning_content.clone();
messages.push(assistant_msg.clone());
let _ = session.add_message(SessionMessage::from_chat(&assistant_msg));
let tool_count = resp.tool_calls.len();
let has_lane = resp
.tool_calls
.iter()
.any(|tc| self.tools.lane_of(&tc.name).is_some());
if tool_count > 1 && !has_lane {
use futures::future::join_all;
let cancel_arc = Arc::new(AtomicBool::new(false));
let cancel_ref = cancel as &AtomicBool;
let futures: Vec<_> = resp
.tool_calls
.iter()
.map(|tc| {
let name = tc.name.clone();
let args = tc.parsed_args();
let cancel = cancel_arc.clone();
async move {
let started = std::time::Instant::now();
if cancel_ref.load(Ordering::SeqCst) || cancel.load(Ordering::SeqCst) {
return (
tc.id.clone(),
tc.name.clone(),
"cancelled".to_string(),
0u64,
);
}
let obs = self
.execute_with_retry(&name, args, emitter, cancel.as_ref())
.await;
(
tc.id.clone(),
tc.name.clone(),
obs,
started.elapsed().as_millis() as u64,
)
}
})
.collect();
let results = join_all(futures).await;
for (tc_id, tc_name, observation, duration_ms) in results {
if cancel.load(Ordering::SeqCst) {
emitter.error("cancelled");
return Ok(RunResult {
status: RunStatus::Cancelled,
iterations: iteration + 1,
final_text: String::new(),
});
}
if tc_name == "fetch_url" {
let url = resp
.tool_calls
.iter()
.find(|tc| tc.id == tc_id)
.and_then(|tc| {
tc.parsed_args()
.get("url")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
})
.unwrap_or_default();
stale.record_fetch(&url, &observation);
}
let (status, code) = classify_tool_status(&observation);
emitter.tool_status(
&tc_id,
&tc_name,
status,
duration_ms,
code,
if code.is_some() {
Some(observation.as_str())
} else {
None
},
);
emitter.seg_observation(&observation, self.config.limits.display_preview_chars);
let tool_msg = ChatMessage::tool_result(tc_id, observation);
messages.push(tool_msg.clone());
let _ = session.add_message(SessionMessage::from_chat(&tool_msg));
}
} else {
for tc in &resp.tool_calls {
if cancel.load(Ordering::SeqCst) {
emitter.error("cancelled");
return Ok(RunResult {
status: RunStatus::Cancelled,
iterations: iteration + 1,
final_text: String::new(),
});
}
let args = tc.parsed_args();
let started = std::time::Instant::now();
let observation = self
.execute_with_retry(&tc.name, args, emitter, cancel)
.await;
let duration_ms = started.elapsed().as_millis() as u64;
if tc.name == "fetch_url" {
let url = tc
.parsed_args()
.get("url")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
stale.record_fetch(&url, &observation);
}
let (status, code) = classify_tool_status(&observation);
emitter.tool_status(
&tc.id,
&tc.name,
status,
duration_ms,
code,
if code.is_some() {
Some(observation.as_str())
} else {
None
},
);
emitter.seg_observation(&observation, self.config.limits.display_preview_chars);
let tool_msg = ChatMessage::tool_result(tc.id.clone(), observation);
messages.push(tool_msg.clone());
let _ = session.add_message(SessionMessage::from_chat(&tool_msg));
}
}
let _ = estimate_messages_tokens(&messages);
}
let _ = session.flush();
emitter.done_result(
&session_id,
"max_iterations",
self.config.max_iterations,
&last_text,
);
Ok(RunResult {
status: RunStatus::MaxIterations,
iterations: self.config.max_iterations,
final_text: last_text,
})
}
async fn chat_with_retry(
&self,
view: &[ChatMessage],
emitter: &dyn EventSink,
cancel: &AtomicBool,
cancel_notify: &Notify,
) -> Result<crate::llm::types::LlmResponse> {
let max = self.config.limits.max_retries.max(1);
let mut last: Option<AacodeError> = None;
for attempt in 0..max {
if cancel.load(Ordering::SeqCst) {
return Err(AacodeError::Cancelled);
}
let use_buffered = attempt > 0 && !emitter.is_tty();
let (result, buffered): (_, Option<CollectingSink>) = if use_buffered {
let b = CollectingSink::new(false);
let fut = self.llm.chat_stream(view, &self.native_tools, &b, cancel);
let r = tokio::select! {
r = fut => r,
_ = cancel_checker(cancel, cancel_notify) => Err(AacodeError::Cancelled),
};
(r, Some(b))
} else {
let fut = self
.llm
.chat_stream(view, &self.native_tools, emitter, cancel);
let r = tokio::select! {
r = fut => r,
_ = cancel_checker(cancel, cancel_notify) => Err(AacodeError::Cancelled),
};
(r, None)
};
match result {
Ok(r) => {
if let Some(collected) = buffered {
for line in collected.lines() {
emitter.emit_line(&line);
}
}
return Ok(r);
}
Err(AacodeError::Cancelled) => return Err(AacodeError::Cancelled),
Err(e) => {
let retryable = e.is_retryable();
if retryable && attempt + 1 < max {
emitter.tool_progress(
"building",
&format!("llm retry {}/{} ({})", attempt + 2, max, brief(&e)),
0,
);
}
last = Some(e);
if !retryable || attempt + 1 >= max {
break;
}
let delay = 1u64 << attempt.min(4);
let jitter = {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.subsec_nanos() as u64;
let half_range = (delay * 250).max(1); nanos % (half_range * 2).max(1)
};
let actual = if jitter > 0 {
(delay * 1000 + jitter - delay * 250).max(100)
} else {
delay * 1000
};
tokio::time::sleep(std::time::Duration::from_millis(actual)).await;
}
}
}
Err(last.unwrap_or_else(|| AacodeError::Api("model call failed".into())))
}
async fn execute_with_retry(
&self,
name: &str,
args: Value,
_emitter: &dyn EventSink,
cancel: &AtomicBool,
) -> String {
if self.config.safety.permission_mode == crate::config::PermissionMode::ReadOnly
&& !crate::tools::registry::is_read_only(name)
{
return serde_json::json!({
"success": false,
"error": format!(
"read-only mode: tool '{name}' is not allowed (switch to full access to run it)"
),
"tool": name,
})
.to_string();
}
let max_retries = self.config.limits.max_retries.max(1);
let mut last = String::new();
for attempt in 0..max_retries {
if cancel.load(Ordering::SeqCst) {
return "cancelled".to_string();
}
let notify = Notify::new();
let obs = tokio::select! {
obs = self.tools.execute(name, args.clone(), cancel) => obs,
_ = cancel_checker(cancel, ¬ify) => return "cancelled".to_string(),
};
let lower = obs.to_lowercase();
let retryable = lower.contains("timeout")
|| lower.contains("timed out")
|| lower.contains("connection")
|| lower.contains("temporary")
|| lower.contains("network");
last = obs;
if !retryable || attempt + 1 >= max_retries {
break;
}
}
last
}
}
fn classify_tool_status(observation: &str) -> (&'static str, Option<&'static str>) {
let l = observation.to_lowercase();
let failed = l.contains("parameter validation failed")
|| l.contains("execution error")
|| l.contains("execution cancelled")
|| l.contains("unknown tool")
|| observation.starts_with('❌')
|| l.contains("\"success\":false")
|| l.contains("\"success\": false")
|| l.contains("command rejected by safety guard")
|| l.contains("traceback");
if !failed {
return ("success", None);
}
if l.contains("safety guard") || l.contains("dangerous operation") {
("error", Some("safety_guard"))
} else if l.contains("timed out") || l.contains("timeout") {
("error", Some("timeout"))
} else if l.contains("parameter validation")
|| l.contains("parameter parsing")
|| l.contains("missing required")
{
("error", Some("validation"))
} else if l.contains("unknown tool") || l.contains("not found") {
("error", Some("not_found"))
} else if l.contains("cancel") {
("error", Some("cancelled"))
} else {
("error", Some("tool_error"))
}
}
fn classify_api_error(e: &AacodeError) -> String {
let s = e.to_string();
let l = s.to_lowercase();
if l.contains("401") || l.contains("authentication") || l.contains("invalid api key") {
format!("API authentication failed (check LLM_API_KEY / endpoint): {s}")
} else if l.contains("quota") || l.contains("rate limit") || l.contains("429") {
format!("API quota/rate limit reached (check usage or wait): {s}")
} else if l.contains("connection") || l.contains("timeout") || l.contains("network") {
format!("Network error reaching the model (check connectivity/endpoint): {s}")
} else if l.contains("config") || l.contains("api key not configured") {
format!("Configuration error: {s}")
} else {
s
}
}
#[derive(Default)]
struct StaleTracker {
by_domain: HashMap<String, Vec<bool>>, warned: std::collections::HashSet<String>,
}
impl StaleTracker {
fn detect(&mut self, tool_calls: &[ToolCall]) -> Option<String> {
if !tool_calls.iter().any(|tc| tc.name == "fetch_url") {
return None;
}
let mut candidate: Option<String> = None;
for (domain, entries) in &self.by_domain {
if entries.len() < 3 {
continue;
}
let last3 = &entries[entries.len() - 3..];
if last3.iter().all(|has| !has) {
let key = format!("stale_{domain}");
if !self.warned.contains(&key) {
candidate = Some(domain.clone());
break;
}
}
}
let domain = candidate?;
self.warned.insert(format!("stale_{domain}"));
Some(format!(
"Last 3 fetch_url calls to '{domain}' returned unreadable content. Try search_web or proceed with existing knowledge."
))
}
fn record_fetch(&mut self, url: &str, observation: &str) {
let domain = domain_of(url);
let text: String = observation
.chars()
.filter(|c| *c != '<' && *c != '>')
.collect();
let has_content = text.trim().len() >= 200;
let e = self.by_domain.entry(domain).or_default();
e.push(has_content);
if e.len() > 5 {
let drain = e.len() - 5;
e.drain(0..drain);
}
}
}
fn domain_of(url: &str) -> String {
let after_scheme = url.split("://").nth(1).unwrap_or(url);
after_scheme
.split('/')
.next()
.unwrap_or(after_scheme)
.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::AgentConfig;
use crate::llm::types::LlmResponse;
use crate::stream::CollectingSink;
use crate::tools::registry::Tool;
use crate::tools::schema::{ParamType, ToolParameter, ToolSchema};
use std::sync::Mutex;
struct ScriptedLlm {
responses: Mutex<Vec<LlmResponse>>,
}
#[async_trait::async_trait]
impl LlmClient for ScriptedLlm {
async fn chat_stream(
&self,
_m: &[ChatMessage],
_t: &[Value],
emitter: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
let mut q = self.responses.lock().unwrap();
if q.is_empty() {
return Err(AacodeError::Api("no more scripted responses".into()));
}
let r = q.remove(0);
emitter.seg("thought", &r.text);
Ok(r)
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
struct RecordingLlm {
responses: Mutex<Vec<LlmResponse>>,
seen: Mutex<Vec<Vec<ChatMessage>>>,
}
#[async_trait::async_trait]
impl LlmClient for RecordingLlm {
async fn chat_stream(
&self,
m: &[ChatMessage],
_t: &[Value],
emitter: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
self.seen.lock().unwrap().push(m.to_vec());
let mut q = self.responses.lock().unwrap();
if q.is_empty() {
return Err(AacodeError::Api("no more scripted responses".into()));
}
let r = q.remove(0);
emitter.seg("thought", &r.text);
Ok(r)
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
struct EchoTool;
#[async_trait::async_trait]
impl Tool for EchoTool {
fn schema(&self) -> ToolSchema {
ToolSchema::new(
"echo",
"echo",
vec![ToolParameter::new(
"text",
ParamType::String,
true,
"t",
&[],
)],
)
}
async fn call(&self, args: &Value, _c: &AtomicBool) -> Result<String> {
Ok(format!("echoed: {}", args["text"].as_str().unwrap_or("")))
}
}
struct SleepyTool {
name: &'static str,
delay_ms: u64,
lane: Option<&'static str>,
log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
}
#[async_trait::async_trait]
impl Tool for SleepyTool {
fn schema(&self) -> ToolSchema {
ToolSchema::new(self.name, "sleep", vec![])
}
fn lane(&self) -> Option<&'static str> {
self.lane
}
async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
self.log
.lock()
.unwrap()
.push(format!("start:{}", self.name));
tokio::time::sleep(std::time::Duration::from_millis(self.delay_ms)).await;
self.log.lock().unwrap().push(format!("end:{}", self.name));
Ok(format!("{} done", self.name))
}
}
fn tmp_proj() -> std::path::PathBuf {
let d = std::env::temp_dir().join(format!(
"aacode_react_{}_{}",
std::process::id(),
uuid::Uuid::new_v4().simple()
));
std::fs::create_dir_all(&d).unwrap();
d
}
fn tc(name: &str, args: &str) -> ToolCall {
ToolCall {
id: "c1".into(),
name: name.into(),
arguments: args.into(),
}
}
fn named_call(id: &str, name: &str) -> ToolCall {
ToolCall {
id: id.into(),
name: name.into(),
arguments: "{}".into(),
}
}
#[tokio::test]
async fn lane_batch_runs_sequentially() {
let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let mut reg = ToolRegistry::new();
reg.register(Box::new(SleepyTool {
name: "a",
delay_ms: 30,
lane: Some("x"),
log: log.clone(),
}));
reg.register(Box::new(SleepyTool {
name: "b",
delay_ms: 30,
lane: Some("x"),
log: log.clone(),
}));
let llm = ScriptedLlm {
responses: Mutex::new(vec![
LlmResponse {
tool_calls: vec![named_call("c1", "a"), named_call("c2", "b")],
..Default::default()
},
LlmResponse {
text: "done".into(),
..Default::default()
},
]),
};
let mut cfg = AgentConfig::default();
cfg.max_iterations = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
loop_
.run(
vec![ChatMessage::system("s"), ChatMessage::user("u")],
&mut sm,
&sink,
&cancel,
)
.await
.unwrap();
let events = log.lock().unwrap().clone();
assert_eq!(events, vec!["start:a", "end:a", "start:b", "end:b"]);
let ids: Vec<String> = sm
.messages
.iter()
.filter(|m| m.role == "tool")
.map(|m| m.tool_call_id.clone().unwrap_or_default())
.collect();
assert_eq!(ids, vec!["c1".to_string(), "c2".to_string()]);
}
#[tokio::test]
async fn concurrent_batch_submits_results_in_call_order() {
let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let mut reg = ToolRegistry::new();
reg.register(Box::new(SleepyTool {
name: "slow",
delay_ms: 80,
lane: None,
log: log.clone(),
}));
reg.register(Box::new(SleepyTool {
name: "fast",
delay_ms: 5,
lane: None,
log: log.clone(),
}));
let llm = ScriptedLlm {
responses: Mutex::new(vec![
LlmResponse {
tool_calls: vec![named_call("c1", "slow"), named_call("c2", "fast")],
..Default::default()
},
LlmResponse {
text: "done".into(),
..Default::default()
},
]),
};
let mut cfg = AgentConfig::default();
cfg.max_iterations = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
loop_
.run(
vec![ChatMessage::system("s"), ChatMessage::user("u")],
&mut sm,
&sink,
&cancel,
)
.await
.unwrap();
let ids: Vec<String> = sm
.messages
.iter()
.filter(|m| m.role == "tool")
.map(|m| m.tool_call_id.clone().unwrap_or_default())
.collect();
assert_eq!(ids, vec!["c1".to_string(), "c2".to_string()]);
let events = log.lock().unwrap().clone();
let end_fast = events.iter().position(|e| e == "end:fast");
let end_slow = events.iter().position(|e| e == "end:slow");
assert!(
end_fast < end_slow,
"independent tools should overlap: {events:?}"
);
}
#[tokio::test]
async fn read_only_mode_blocks_mutating_tools() {
let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let mut reg = ToolRegistry::new();
reg.register(Box::new(SleepyTool {
name: "run_shell",
delay_ms: 1,
lane: None,
log: log.clone(),
}));
reg.register(Box::new(SleepyTool {
name: "get_page_title",
delay_ms: 1,
lane: None,
log: log.clone(),
}));
let llm = ScriptedLlm {
responses: Mutex::new(vec![
LlmResponse {
tool_calls: vec![
named_call("c1", "run_shell"),
named_call("c2", "get_page_title"),
],
..Default::default()
},
LlmResponse {
text: "done".into(),
..Default::default()
},
]),
};
let mut cfg = AgentConfig::default();
cfg.max_iterations = 3;
cfg.safety.permission_mode = crate::config::PermissionMode::ReadOnly;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
loop_
.run(
vec![ChatMessage::system("s"), ChatMessage::user("u")],
&mut sm,
&sink,
&cancel,
)
.await
.unwrap();
let events = log.lock().unwrap().clone();
assert!(
!events.iter().any(|e| e == "start:run_shell"),
"mutating tool ran in read-only mode: {events:?}"
);
assert!(events.iter().any(|e| e == "start:get_page_title"));
let blocked = sm
.messages
.iter()
.find(|m| m.role == "tool" && m.tool_call_id.as_deref() == Some("c1"))
.expect("blocked tool result");
assert!(
blocked.content.contains("read-only mode"),
"{}",
blocked.content
);
}
#[test]
fn classifies_tool_status() {
assert_eq!(super::classify_tool_status("all good"), ("success", None));
assert_eq!(
super::classify_tool_status(
"{\"success\":false,\"hint\":\"command timed out — split it\"}"
),
("error", Some("timeout"))
);
assert_eq!(
super::classify_tool_status("Parameter validation failed\n\nmissing required"),
("error", Some("validation"))
);
assert_eq!(
super::classify_tool_status("Execution error: boom"),
("error", Some("tool_error"))
);
assert_eq!(
super::classify_tool_status("Error: Unknown tool 'x'"),
("error", Some("not_found"))
);
assert_eq!(
super::classify_tool_status("Execution cancelled"),
("error", Some("cancelled"))
);
assert_eq!(
super::classify_tool_status(
"{\"success\":false,\"error\":\"Command rejected by safety guard\"}"
),
("error", Some("safety_guard"))
);
}
#[tokio::test]
async fn retrieval_compacts_view_and_persists_checkpoint() {
let mut cfg = AgentConfig::default();
cfg.max_iterations = 2;
cfg.context.max_context_tokens = 2_000_000;
cfg.context.retrieval_trigger_tokens = 50;
cfg.context.retrieval_keep_tokens = 40;
cfg.context.retrieval_keep_min_rounds = 1;
cfg.context.retrieval_keep_recent_user_messages = 1;
let llm = RecordingLlm {
responses: Mutex::new(vec![LlmResponse {
text: "done".into(),
..Default::default()
}]),
seen: Mutex::new(vec![]),
};
let reg = ToolRegistry::new();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("seed", None).unwrap();
let mut msgs = vec![ChatMessage::system("SYS")];
for i in 0..20 {
msgs.push(ChatMessage::user(format!("task {i} {}", "x".repeat(120))));
msgs.push(ChatMessage::assistant("ok"));
}
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(msgs.clone(), &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let seen = llm.seen.lock().unwrap();
let first = &seen[0];
assert!(first.len() < msgs.len(), "view should be compacted");
assert_eq!(first[0].role, "system");
assert!(
first[1].content.starts_with("<conversation-checkpoint>"),
"checkpoint stub missing: {}",
first[1].content
);
let cp = sm.context_checkpoint().expect("checkpoint persisted");
assert!(cp["cutoff_index"].as_u64().unwrap() > 0);
}
#[tokio::test]
async fn completes_when_no_tool_calls() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![LlmResponse {
text: "all done".into(),
..Default::default()
}]),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("hi")];
let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
assert_eq!(res.status, RunStatus::Completed);
assert_eq!(res.final_text, "all done");
assert!(sink.lines().iter().any(|l| l.contains(r#""type":"done""#)));
}
#[tokio::test]
async fn executes_tool_then_completes() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![
LlmResponse {
text: "let me echo".into(),
tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
..Default::default()
},
LlmResponse {
text: "done".into(),
..Default::default()
},
]),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
assert_eq!(res.status, RunStatus::Completed);
assert!(sink.lines().iter().any(|l| l.contains("echoed: hey")));
assert!(sm.messages.iter().any(|m| m.tool_calls.is_some()));
assert!(sm.messages.iter().any(|m| m.role == "tool"));
}
#[tokio::test]
async fn persists_turn_messages_without_segment_duplication() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![
LlmResponse {
text: "let me echo".into(),
tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
reasoning_content: Some("I need to echo".into()),
..Default::default()
},
LlmResponse {
text: "all done".into(),
reasoning_content: Some("summarize".into()),
..Default::default()
},
]),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let tool_asst = sm
.messages
.iter()
.find(|m| m.tool_calls.is_some())
.expect("assistant with tool_calls present");
assert_eq!(
tool_asst.reasoning_content.as_deref(),
Some("I need to echo")
);
assert_eq!(tool_asst.content, "let me echo");
assert_eq!(tool_asst.tool_calls.as_ref().unwrap()[0].name, "echo");
let tool_msg = sm
.messages
.iter()
.find(|m| m.role == "tool")
.expect("tool result present");
assert!(tool_msg.content.contains("echoed: hey"));
let final_asst = sm
.messages
.iter()
.rev()
.find(|m| m.role == "assistant")
.expect("final assistant present");
assert_eq!(final_asst.reasoning_content.as_deref(), Some("summarize"));
assert_eq!(final_asst.content, "all done");
}
#[tokio::test]
async fn llm_request_has_exactly_one_copy_of_each_tool_message() {
let llm = RecordingLlm {
responses: Mutex::new(vec![
LlmResponse {
text: "let me echo".into(),
tool_calls: vec![tc("echo", "{\"text\":\"hey\"}")],
reasoning_content: Some("I need to echo".into()),
..Default::default()
},
LlmResponse {
text: "all done".into(),
reasoning_content: Some("summarize".into()),
..Default::default()
},
]),
seen: Mutex::new(vec![]),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let msgs = vec![ChatMessage::system("sys"), ChatMessage::user("echo hey")];
let res = loop_.run(msgs, &mut sm, &sink, &cancel).await.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let seen = llm.seen.lock().unwrap();
assert_eq!(seen.len(), 2, "one request per model turn");
let second = &seen[1];
let with_calls: Vec<&ChatMessage> =
second.iter().filter(|m| m.tool_calls.is_some()).collect();
assert_eq!(with_calls.len(), 1, "exactly one assistant tool_calls msg");
let tool_msgs: Vec<&ChatMessage> = second.iter().filter(|m| m.role == "tool").collect();
assert_eq!(
tool_msgs.len(),
1,
"exactly one tool result (no duplication)"
);
assert_eq!(tool_msgs[0].tool_call_id.as_deref(), Some("c1"));
assert_eq!(tool_msgs[0].content, "echoed: hey");
assert_eq!(with_calls[0].content, "let me echo");
assert_eq!(
with_calls[0].reasoning_content.as_deref(),
Some("I need to echo")
);
let ai = second.iter().position(|m| m.tool_calls.is_some()).unwrap();
assert_eq!(second[ai + 1].role, "tool", "result adjacent to its call");
assert_eq!(second.iter().filter(|m| m.role == "user").count(), 1);
assert_eq!(second.iter().filter(|m| m.role == "system").count(), 1);
assert_eq!(sm.history_chat().len(), sm.messages.len());
}
#[tokio::test]
async fn respects_cancel_before_start() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![]),
};
let reg = ToolRegistry::new();
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(true);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Cancelled);
}
#[tokio::test]
async fn max_iterations_reached() {
let responses: Vec<LlmResponse> = (0..5)
.map(|_| LlmResponse {
text: "again".into(),
tool_calls: vec![tc("echo", "{\"text\":\"x\"}")],
..Default::default()
})
.collect();
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let mut cfg = AgentConfig::default();
cfg.max_iterations = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::MaxIterations);
assert_eq!(res.iterations, 3);
}
#[tokio::test]
async fn domain_parse() {
assert_eq!(domain_of("https://example.com/path"), "example.com");
assert_eq!(domain_of("http://a.b.c/x/y"), "a.b.c");
}
#[tokio::test]
async fn classify_errors() {
assert!(
classify_api_error(&AacodeError::Api("HTTP 401 no".into())).contains("authentication")
);
assert!(classify_api_error(&AacodeError::Api("rate limit".into())).contains("quota"));
assert!(classify_api_error(&AacodeError::Network("reset".into())).contains("Network"));
assert!(classify_api_error(&AacodeError::Config("no key".into())).contains("Configuration"));
}
struct FlakyLlm {
remaining_failures: Mutex<u32>,
}
#[async_trait::async_trait]
impl LlmClient for FlakyLlm {
async fn chat_stream(
&self,
_m: &[ChatMessage],
_t: &[Value],
_e: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
let mut n = self.remaining_failures.lock().unwrap();
if *n > 0 {
*n -= 1;
return Err(AacodeError::Network("temporary reset".into()));
}
Ok(LlmResponse {
text: "recovered".into(),
..Default::default()
})
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
#[tokio::test]
async fn retries_transient_llm_errors() {
let llm = FlakyLlm {
remaining_failures: Mutex::new(1),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
assert_eq!(res.final_text, "recovered");
}
#[tokio::test]
async fn truncated_response_not_completed() {
let responses: Vec<LlmResponse> = (0..5)
.map(|_| LlmResponse {
text: "partial".into(),
tool_calls: vec![],
finish_reason: Some("length".to_string()),
..Default::default()
})
.collect();
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let mut cfg = AgentConfig::default();
cfg.max_iterations = 5;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(
res.status,
RunStatus::MaxIterations,
"truncated response must not be treated as completed"
);
}
#[tokio::test]
async fn max_iterations_returns_last_text() {
let responses: Vec<LlmResponse> = (0..5)
.map(|i| LlmResponse {
text: format!("iter {}", i),
tool_calls: vec![tc("echo", "{\"text\":\"x\"}")],
..Default::default()
})
.collect();
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(EchoTool));
let mut cfg = AgentConfig::default();
cfg.max_iterations = 2;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::MaxIterations);
assert!(
!res.final_text.is_empty(),
"max_iterations must return last response text"
);
}
struct FailThenStreamLlm {
remaining_failures: Mutex<u32>,
}
#[async_trait::async_trait]
impl LlmClient for FailThenStreamLlm {
async fn chat_stream(
&self,
_m: &[ChatMessage],
_t: &[Value],
emitter: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
let mut n = self.remaining_failures.lock().unwrap();
if *n > 0 {
*n -= 1;
return Err(AacodeError::Network("temporary reset".into()));
}
emitter.delta("thought", "streaming");
emitter.seg("thinking", "reasoning done");
emitter.seg("thought", "complete response");
Ok(LlmResponse {
text: "complete response".into(),
..Default::default()
})
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
struct EmitThenFailLlm {
call: Mutex<u32>,
}
#[async_trait::async_trait]
impl LlmClient for EmitThenFailLlm {
async fn chat_stream(
&self,
_m: &[ChatMessage],
_t: &[Value],
emitter: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
let mut call = self.call.lock().unwrap();
*call += 1;
if *call == 1 {
emitter.delta("thought", "partial tok");
emitter.seg("thinking", "partial reasoning");
emitter.seg("thought", "partial text");
return Err(AacodeError::Network("mid-stream disconnect".into()));
}
emitter.delta("thought", "complete");
emitter.seg("thinking", "full reasoning");
emitter.seg("thought", "final complete text");
Ok(LlmResponse {
text: "final complete text".into(),
..Default::default()
})
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
struct AlwaysFailLlm;
#[async_trait::async_trait]
impl LlmClient for AlwaysFailLlm {
async fn chat_stream(
&self,
_m: &[ChatMessage],
_t: &[Value],
_emitter: &dyn EventSink,
_c: &AtomicBool,
) -> Result<LlmResponse> {
Err(AacodeError::Network("always fails".into()))
}
async fn validate(&self) -> Result<()> {
Ok(())
}
}
struct TimeoutTool;
#[async_trait::async_trait]
impl Tool for TimeoutTool {
fn schema(&self) -> ToolSchema {
ToolSchema::new(
"timeout_cmd",
"timeout",
vec![ToolParameter::new("x", ParamType::String, true, "x", &[])],
)
}
async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
Ok("command timed out after 10s".to_string())
}
}
struct CountingTool {
count: std::sync::Arc<Mutex<u32>>,
tool_name: &'static str,
response: String,
}
impl CountingTool {
fn new(name: &'static str, response: String) -> (Self, std::sync::Arc<Mutex<u32>>) {
let count = std::sync::Arc::new(Mutex::new(0u32));
(
CountingTool {
count: count.clone(),
tool_name: name,
response,
},
count,
)
}
}
#[async_trait::async_trait]
impl Tool for CountingTool {
fn schema(&self) -> ToolSchema {
ToolSchema::new(self.tool_name, self.tool_name, vec![])
}
async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
*self.count.lock().unwrap() += 1;
Ok(self.response.clone())
}
}
#[tokio::test]
async fn execute_with_retry_retries_on_network() {
let mut reg = ToolRegistry::new();
let (tool, counter) = CountingTool::new(
"understand_image",
r#"{"success":false,"error":"network error: vision: error sending request for url (https://api.moonshot.cn/v1/chat/completions)"}"#.to_string(),
);
reg.register(Box::new(tool));
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
cfg.max_iterations = 1;
let responses = vec![LlmResponse {
text: "try".into(),
tool_calls: vec![tc(
"understand_image",
r#"{"image_path":"a.jpg","prompt":"desc"}"#,
)],
..Default::default()
}];
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let _ = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await;
let invocations = *counter.lock().unwrap();
assert!(
invocations >= 2,
"network error must trigger retry, got {invocations} invocations"
);
}
#[tokio::test]
async fn execute_with_retry_does_not_retry_on_config_error() {
let mut reg = ToolRegistry::new();
let (tool, counter) = CountingTool::new(
"understand_image",
r#"{"success":false,"error":"config error: no multimodal model configured"}"#
.to_string(),
);
reg.register(Box::new(tool));
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
cfg.max_iterations = 1;
let responses = vec![LlmResponse {
text: "try".into(),
tool_calls: vec![tc(
"understand_image",
r#"{"image_path":"a.jpg","prompt":"desc"}"#,
)],
..Default::default()
}];
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let _ = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await;
assert_eq!(
*counter.lock().unwrap(),
1,
"config error must not trigger retry"
);
}
#[tokio::test]
async fn execute_with_retry_timeout_retried() {
let mut reg = ToolRegistry::new();
let (tool, counter) =
CountingTool::new("timeout_cmd", "command timed out after 10s".to_string());
reg.register(Box::new(tool));
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
cfg.max_iterations = 1;
let responses = vec![LlmResponse {
text: "try".into(),
tool_calls: vec![tc("timeout_cmd", r#"{"x":"1"}"#)],
..Default::default()
}];
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let _ = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await;
assert!(*counter.lock().unwrap() >= 2, "timeout must trigger retry");
}
#[tokio::test]
async fn execute_with_retry_connection_retried() {
let mut reg = ToolRegistry::new();
let (tool, counter) = CountingTool::new("fetch_url", "connection reset".to_string());
reg.register(Box::new(tool));
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
cfg.max_iterations = 1;
let responses = vec![LlmResponse {
text: "try".into(),
tool_calls: vec![tc("fetch_url", r#"{"url":"http://x"}"#)],
..Default::default()
}];
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let _ = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await;
assert!(
*counter.lock().unwrap() >= 2,
"connection must trigger retry"
);
}
#[tokio::test]
async fn execute_with_retry_respects_cancel() {
let responses = vec![LlmResponse {
text: "try".into(),
tool_calls: vec![tc("timeout_cmd", "{\"x\":\"1\"}")],
..Default::default()
}];
let llm = ScriptedLlm {
responses: Mutex::new(responses),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(TimeoutTool));
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 10;
cfg.max_iterations = 5;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(true); let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Cancelled);
}
#[tokio::test]
async fn max_context_tokens_enforced_and_errors() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![LlmResponse {
text: "done".into(),
..Default::default()
}]),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.max_iterations = 5;
cfg.context.max_context_tokens = 10;
cfg.context.compact_trigger_tokens = 5;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let msgs = vec![
ChatMessage::system("you are a helpful assistant who provides detailed responses"),
ChatMessage::user("do something"),
];
let res = loop_.run(msgs, &mut sm, &sink, &cancel).await;
assert!(
res.is_err(),
"context over max must error, not silently proceed"
);
let msg = format!("{}", res.err().unwrap());
assert!(
msg.contains("context too large") || msg.contains("exceeds limit"),
"error must mention context size: {msg}"
);
}
#[tokio::test]
async fn pipe_mode_retry_buffers_and_replays() {
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(1),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
assert_eq!(res.final_text, "complete response");
let lines = sink.lines();
assert!(
lines
.iter()
.any(|l| l.contains(r#""type":"delta""#) && l.contains("streaming")),
"delta event must be replayed to real sink"
);
assert!(
lines
.iter()
.any(|l| l.contains(r#""type":"seg_content""#) && l.contains("complete response")),
"seg_content event must be replayed to real sink"
);
}
#[tokio::test]
async fn tty_mode_never_buffers_retries() {
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(1),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(true);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let lines = sink.lines();
assert!(
lines.iter().any(|l| l == "streaming"),
"TTY delta must emit raw text: {lines:?}"
);
}
#[tokio::test]
async fn multiple_retries_only_last_success_replayed() {
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(2),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 4;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let lines = sink.lines();
let seg_count = lines
.iter()
.filter(|l| l.contains(r#""type":"seg_content""#) && l.contains("complete response"))
.count();
assert_eq!(
seg_count, 1,
"only one seg_content for the final thought; found {seg_count}. lines: {lines:?}"
);
}
#[tokio::test]
async fn first_attempt_streams_directly_leak_is_allowed() {
let llm = EmitThenFailLlm {
call: Mutex::new(0),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
assert_eq!(res.final_text, "final complete text");
let lines = sink.lines();
assert!(
lines.iter().any(|l| l.contains("partial text")),
"first attempt leaked: {lines:?}"
);
assert!(
lines.iter().any(|l| l.contains("final complete text")),
"retry must replay: {lines:?}"
);
}
#[tokio::test]
async fn buffered_retry_replays_all_event_types() {
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(1),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Completed);
let lines = sink.lines();
assert!(
lines
.iter()
.any(|l| l.contains(r#""type":"delta""#) && l.contains("streaming")),
"delta missing: {lines:?}"
);
assert!(
lines.iter().any(|l| l.contains(r#""type":"seg_content""#) && l.contains(r#""seg":"thinking""#)),
"thinking seg missing: {lines:?}"
);
assert!(
lines
.iter()
.any(|l| l.contains(r#""type":"seg_content""#) && l.contains(r#""seg":"thought""#)),
"thought seg missing: {lines:?}"
);
assert!(
lines.iter().any(|l| l.contains(r#""type":"done""#)),
"done event missing: {lines:?}"
);
}
#[tokio::test]
async fn buffered_retry_respects_cancel() {
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(0),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(true);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert_eq!(res.status, RunStatus::Cancelled);
}
#[tokio::test]
async fn cancel_during_retry_backoff_stops_cleanly() {
use std::sync::Arc;
let llm = FailThenStreamLlm {
remaining_failures: Mutex::new(2),
};
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 5;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = Arc::new(AtomicBool::new(false));
let cancel2 = Arc::clone(&cancel);
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(500));
cancel2.store(true, Ordering::SeqCst);
});
let res = loop_
.run(
vec![ChatMessage::user("x")],
&mut sm,
&sink,
cancel.as_ref(),
)
.await
.unwrap();
assert_eq!(
res.status,
RunStatus::Cancelled,
"cancel during retry backoff must return Cancelled"
);
}
#[tokio::test]
async fn exhausted_retries_buffers_discarded_no_events() {
let llm = AlwaysFailLlm;
let reg = ToolRegistry::new();
let mut cfg = AgentConfig::default();
cfg.limits.max_retries = 3;
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = AtomicBool::new(false);
let res = loop_
.run(vec![ChatMessage::user("x")], &mut sm, &sink, &cancel)
.await
.unwrap();
assert!(matches!(res.status, RunStatus::Error(_)));
let lines = sink.lines();
assert!(
lines.iter().any(|l| l.contains(r#""type":"error""#)),
"error event must be emitted: {lines:?}"
);
let deltas = lines
.iter()
.filter(|l| l.contains(r#""type":"delta""#))
.count();
let segs = lines
.iter()
.filter(|l| l.contains(r#""type":"seg_content""#))
.count();
assert_eq!(deltas, 0, "no delta events from failed retries: {lines:?}");
assert_eq!(
segs, 0,
"no seg_content events from failed retries: {lines:?}"
);
}
struct BlockingTool;
#[async_trait::async_trait]
impl Tool for BlockingTool {
fn schema(&self) -> ToolSchema {
ToolSchema::new("blocking_tool", "never returns", vec![])
}
async fn call(&self, _args: &Value, _c: &AtomicBool) -> Result<String> {
futures::future::pending::<()>().await;
Ok(String::new())
}
}
#[tokio::test]
async fn cancel_interrupts_a_blocking_tool() {
let llm = ScriptedLlm {
responses: Mutex::new(vec![LlmResponse {
text: "work".into(),
tool_calls: vec![tc("blocking_tool", "{}")],
..Default::default()
}]),
};
let mut reg = ToolRegistry::new();
reg.register(Box::new(BlockingTool));
let cfg = AgentConfig::default();
let loop_ = ReactLoop::new(&llm, ®, &cfg, vec![]);
let proj = tmp_proj();
let mut sm = SessionManager::new(&proj);
sm.create_session("t", None).unwrap();
let sink = CollectingSink::new(false);
let cancel = Arc::new(AtomicBool::new(false));
let c2 = cancel.clone();
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(200));
c2.store(true, Ordering::SeqCst);
});
let start = std::time::Instant::now();
let res = loop_
.run(
vec![ChatMessage::user("x")],
&mut sm,
&sink,
cancel.as_ref(),
)
.await
.unwrap();
assert!(
matches!(res.status, RunStatus::Cancelled),
"expected Cancelled, got {:?}",
res.status
);
assert!(
start.elapsed() < std::time::Duration::from_secs(3),
"cancel should be prompt, took {:?}",
start.elapsed()
);
}
}
async fn cancel_checker(cancel: &AtomicBool, notify: &Notify) {
use std::time::Duration;
loop {
tokio::select! {
_ = notify.notified() => {},
_ = tokio::time::sleep(Duration::from_secs(1)) => {},
}
if cancel.load(Ordering::SeqCst) {
return;
}
}
}
fn brief(e: &crate::error::AacodeError) -> String {
let s = e.to_string();
let first = s.lines().next().unwrap_or(&s);
let mut out: String = first.chars().take(60).collect();
if first.chars().count() > 60 {
out.push('…');
}
out
}