use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use tracing::{debug, info, warn}; use zesdex_domain::agent::{AgentTurnParams, TurnEvent}; use zesdex_domain::core::{ChatMessage, StreamEvent, ToolDef}; use zesdex_domain::main_agent_prompt; use super::ToolExecutor; use crate::ports::ProviderService; /// Maximum tool-call iterations per agent turn before forcing termination. const MAX_TURN_ITERATIONS: u32 = 50; /// Maximum number of consecutive identical tool errors before the loop /// injects a recovery note and forces a different approach. const MAX_CONSECUTIVE_TOOL_ERRORS: usize = 3; /// Total tool-call errors tolerated per turn before the loop is stopped. const MAX_TOTAL_TOOL_ERRORS: usize = 8; /// Ceiling for a single tool-result message inserted into context. /// /// Tool outputs can be huge (read / semantic_search). Truncating keeps the /// context window from exploding while preserving the important head. const TOOL_OUTPUT_MAX_CHARS: usize = 12_000; /// Total conversation characters that trigger auto-compaction before the /// next LLM call. const AUTO_COMPACT_CHARS: usize = 60_000; // --------------------------------------------------------------------------- // Helper: push a TurnEvent onto the shared queue. // --------------------------------------------------------------------------- fn push_event(queue: &Arc>>, event: TurnEvent) { if let Ok(mut q) = queue.lock() { q.push_back(event); } } // --------------------------------------------------------------------------- // Helper: stream-event callback that forwards tokens to the turn-event queue // and checks the abort flag on each emission. // --------------------------------------------------------------------------- fn make_stream_callback( abort: &Arc, turn_events: &Arc>>, ) -> Box bool + Send> { let abort_clone = Arc::clone(abort); let events_clone = Arc::clone(turn_events); Box::new(move |event: &StreamEvent| -> bool { if abort_clone.load(Ordering::SeqCst) { return false; } match event { StreamEvent::Token(s) => { push_event(&events_clone, TurnEvent::StreamToken(s.clone())); } StreamEvent::Reasoning(s) => { push_event(&events_clone, TurnEvent::StreamReasoning(s.clone())); } _ => {} } true }) } // --------------------------------------------------------------------------- // Helper: truncate a long tool output before it enters the conversation // context. Preserves the head and appends a clear truncation marker. // --------------------------------------------------------------------------- fn truncate_tool_output(output: String) -> String { if output.len() <= TOOL_OUTPUT_MAX_CHARS { return output; } let mut result: String = output.chars().take(TOOL_OUTPUT_MAX_CHARS).collect(); result.push_str(&format!( "\n...[truncated {} chars]", output.len() - TOOL_OUTPUT_MAX_CHARS )); result } // --------------------------------------------------------------------------- // Helper: adaptive generation parameters. // --------------------------------------------------------------------------- /// Pick a `max_tokens` budget for the turn's next LLM call based on the /// length of the user's request. Short requests need far fewer tokens than /// the current hardcoded 4096 — big savings on small tasks. fn adaptive_max_tokens(request_len: usize) -> u32 { if request_len <= 80 { 800 } else if request_len <= 400 { 1600 } else { 4096 } } /// Sum the character length of the conversation (user + assistant + /// tool content) as a cheap proxy for context size. fn conversation_chars(messages: &[ChatMessage]) -> usize { messages .iter() .map(|m| m.content.as_deref().map(str::len).unwrap_or(0)) .sum() } /// Track repeated tool-call errors so the loop can recover instead of /// burning iterations retrying the same failing tool. #[derive(Default)] struct ErrorTracker { consecutive: usize, total: usize, last_tool: String, last_error: String, } impl ErrorTracker { fn record(&mut self, tool_name: &str, error: &str, messages: &mut Vec) { if self.last_tool == tool_name { self.consecutive += 1; } else { self.consecutive = 1; } self.last_tool = tool_name.to_string(); self.last_error = error.to_string(); self.total += 1; // Inject a recovery note once the same tool keeps failing. if self.consecutive >= MAX_CONSECUTIVE_TOOL_ERRORS && !messages.iter().any(|m| { m.content .as_deref() .is_some_and(|c| c.contains("[System note]")) }) { messages.push(ChatMessage::system( zesdex_domain::agent::prompt::error_recovery_note(tool_name, error), )); } } fn should_stop(&self) -> bool { self.consecutive >= MAX_CONSECUTIVE_TOOL_ERRORS * 2 || self.total >= MAX_TOTAL_TOOL_ERRORS } } // --------------------------------------------------------------------------- // Helper: execute a single tool call, push events, return the result string. // --------------------------------------------------------------------------- async fn execute_tool_call( tool_executor: &T, turn_events: &Arc>>, tc: &zesdex_domain::core::ToolCall, ) -> String { let name = &tc.function.name; let args = zesdex_domain::core::tool_call::sanitize_tool_arguments(&tc.function.arguments); debug!("executing tool: {name}"); let output = match tool_executor.execute(name, &args).await { Ok(o) => o, Err(e) => format!("Error: {e}"), }; let is_error = output.starts_with("Error:"); let output = truncate_tool_output(output); push_event( turn_events, TurnEvent::ToolResult { tool_call_id: tc.id.clone(), tool_name: name.clone(), output: output.clone(), is_error, path: None, }, ); output } // --------------------------------------------------------------------------- // Helper: emit usage event from optional LLM response metadata. // --------------------------------------------------------------------------- fn emit_usage(turn_events: &Arc>>, usage: Option<(u64, u64)>) { if let Some((tokens_in, tokens_out)) = usage { push_event( turn_events, TurnEvent::Usage { tokens_in, tokens_out, }, ); } } // --------------------------------------------------------------------------- // Service implementation // --------------------------------------------------------------------------- /// Service implementation for executing an agent turn asynchronously. /// /// The turn loop is adaptive and token-aware: /// - No mandatory explore phase — the *agent* decides when to call the /// `explore_codebase` tool (see the main prompt), so simple queries skip /// exploration entirely. /// - `max_tokens` / `temperature` adapt to the request length and phase. /// - Repeated tool errors trigger a system recovery note and eventually /// stop the loop instead of burning iterations. /// - Tool outputs are truncated before entering context. /// - Oversized histories are auto-compacted before the next LLM call. pub struct AgentTurnServiceImpl { provider: Arc

, tool_executor: Arc, tool_defs: Vec, } impl AgentTurnServiceImpl { pub fn new(provider: Arc

, tool_executor: Arc, tool_defs: Vec) -> Self { Self { provider, tool_executor, tool_defs, } } /// Execute a single LLM call with the current message list, handling /// streaming events and error reporting. async fn call_llm( &self, messages: &[ChatMessage], abort: &Arc, turn_events: &Arc>>, max_tokens: u32, temperature: f32, ) -> Result<(ChatMessage, Option<(u64, u64)>), String> { let on_event = make_stream_callback(abort, turn_events); self.provider .chat_stream( messages, Some(self.tool_defs.clone()), Some(max_tokens), Some(temperature), on_event, ) .await .map_err(|e| format!("LLM error: {e}")) } /// Auto-compact the history in place if it exceeds the threshold. /// /// Runs at most once per turn. Skips the synthetic system prompt that /// this service inserts at index 0. async fn auto_compact_if_needed(&self, messages: &mut Vec) { if conversation_chars(messages) <= AUTO_COMPACT_CHARS { return; } // Keep the system prompt (index 0) out of compaction. let sys = messages[0].clone(); let mut rest: Vec = messages.drain(1..).collect(); let before = rest.len(); if let Err(e) = super::compact_messages_with_ai(&mut rest, self.provider.as_ref()).await { warn!("auto-compact failed (non-fatal): {e}"); } info!( "auto-compacted history: {} messages -> {}", before, rest.len() ); let mut rebuilt = Vec::with_capacity(rest.len() + 1); rebuilt.push(sys); rebuilt.extend(rest); *messages = rebuilt; } } impl super::AgentTurnService for AgentTurnServiceImpl { async fn run_turn(&self, mut params: AgentTurnParams) -> anyhow::Result<()> { info!( "Starting async agent turn with {} messages (model: {})", params.messages.len(), params.model ); // Insert system prompt at position 0 once and keep it there for the // entire turn, avoiding per-iteration clones of the full message list. params .messages .insert(0, ChatMessage::system(main_agent_prompt())); let original_count = params.messages.len(); // Estimate request complexity from the last user message. let request_len = params .messages .last() .and_then(|m| m.content.as_deref()) .map(str::len) .unwrap_or(0); let mut errors = ErrorTracker::default(); // Track whether the previous call produced tool calls — used to // lower temperature once the agent starts producing a final answer. let mut saw_tool_calls = false; for iteration in 0..MAX_TURN_ITERATIONS { // ── Check abort flag ──────────────────────────────────────── if params.abort.load(Ordering::SeqCst) { params.abort.store(false, Ordering::SeqCst); push_event( ¶ms.turn_events, TurnEvent::SystemNote { kind: "info".into(), message: "Turn aborted by user".into(), }, ); break; } if errors.should_stop() { push_event( ¶ms.turn_events, TurnEvent::SystemNote { kind: "warn".into(), message: "Stopping: repeated tool errors without progress".into(), }, ); break; } debug!("agent turn iteration {iteration}"); // ── Auto-compact oversized history before the LLM call ───── self.auto_compact_if_needed(&mut params.messages).await; // ── Adaptive generation parameters ───────────────────────── let max_tokens = adaptive_max_tokens(request_len); // Lower temperature while the agent is still choosing tools to // keep tool selection deterministic; raise it for the final // free-form answer. let temperature = if saw_tool_calls { 0.2 } else { 0.7 }; // ── Stream start + call LLM ───────────────────────────────── push_event(¶ms.turn_events, TurnEvent::StreamStart); let result = self .call_llm( ¶ms.messages, ¶ms.abort, ¶ms.turn_events, max_tokens, temperature, ) .await; match result { Ok((assistant_msg, usage)) => { let content = assistant_msg.content.clone().unwrap_or_default(); let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default(); push_event( ¶ms.turn_events, TurnEvent::StreamDone(assistant_msg.clone()), ); emit_usage(¶ms.turn_events, usage); // ── No tool calls → assistant is done ────────────── if tool_calls.is_empty() { params.messages.push(ChatMessage::assistant(Some(content))); break; } saw_tool_calls = true; params.messages.push(assistant_msg); // ── Execute each tool call ────────────────────────── for tc in &tool_calls { let output = execute_tool_call(self.tool_executor.as_ref(), ¶ms.turn_events, tc) .await; if output.starts_with("Error:") { errors.record(&tc.function.name, &output, &mut params.messages); } params .messages .push(ChatMessage::tool(tc.id.clone(), output)); } } Err(e) => { warn!("{e}"); push_event(¶ms.turn_events, TurnEvent::Error(e)); break; } } } // Remove the synthetic sys_msg before shipping events to the TUI // so the transcript shows only the actual user/assistant/tool exchange. let compacted: Vec = params.messages.drain(original_count - 1..).collect(); push_event(¶ms.turn_events, TurnEvent::Compacted(compacted)); push_event(¶ms.turn_events, TurnEvent::Done); params.in_flight.store(false, Ordering::SeqCst); Ok(()) } } // --------------------------------------------------------------------------- // Conversation compaction // --------------------------------------------------------------------------- /// Maximum number of recent messages to preserve during compaction. const COMPACT_KEEP_TAIL: usize = 6; /// Compacts conversation history using AI summarisation. /// /// Flow: if the message count exceeds `KEEP_TAIL + 2`, the oldest messages /// are drained and summarised by the LLM. The summary is inserted as a /// system message at the head of the remaining history. pub async fn compact_messages_with_ai( messages: &mut Vec, provider: &P, ) -> anyhow::Result<()> { if messages.len() <= COMPACT_KEEP_TAIL + 2 { return Ok(()); // Not enough messages to compact } let split_idx = messages.len() - COMPACT_KEEP_TAIL; let evicted: Vec<_> = messages.drain(..split_idx).collect(); let mut summary_prompt = vec![ChatMessage::system(zesdex_domain::compaction_prompt())]; summary_prompt.extend(evicted); summary_prompt.push(ChatMessage::user( "Please summarise our previous conversation above for context continuity.".to_string(), )); match provider .chat(&summary_prompt, None, Some(1024), Some(0.3)) .await { Ok((summary_msg, _)) => { let summary_text = summary_msg .content .unwrap_or_else(|| "Previous context summarised.".to_string()); let summary_node = ChatMessage::system(format!( "[AI Summary of Previous Conversation]\n{}", summary_text.trim() )); messages.insert(0, summary_node); Ok(()) } Err(e) => { warn!("AI summarisation failed during compact, falling back to simple notice: {e}"); messages.insert( 0, ChatMessage::system( "[Earlier conversation messages compacted to save context window]".to_string(), ), ); Ok(()) } } } // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- #[cfg(test)] mod tests { use super::*; #[test] fn truncate_short_output_is_unchanged() { let out = "short".to_string(); assert_eq!(truncate_tool_output(out.clone()), out); } #[test] fn truncate_long_output_preserves_head_and_marks_cut() { let long = "x".repeat(TOOL_OUTPUT_MAX_CHARS + 500); let truncated = truncate_tool_output(long.clone()); assert!(truncated.len() < long.len()); assert!(truncated.contains("...[truncated")); assert!(truncated.starts_with("xxx")); } #[test] fn adaptive_max_tokens_scales_with_request_len() { assert_eq!(adaptive_max_tokens(10), 800); assert_eq!(adaptive_max_tokens(200), 1600); assert_eq!(adaptive_max_tokens(5000), 4096); } #[test] fn error_tracker_injects_recovery_note_after_repeats() { let mut tracker = ErrorTracker::default(); let mut messages: Vec = Vec::new(); tracker.record("read", "Error: File not found", &mut messages); tracker.record("read", "Error: File not found", &mut messages); assert!(!tracker.should_stop()); // Third consecutive failure → recovery note injected. tracker.record("read", "Error: File not found", &mut messages); assert!(messages.iter().any(|m| m .content .as_deref() .is_some_and(|c| c.contains("[System note]")))); } #[test] fn error_tracker_stops_after_too_many_errors() { let mut tracker = ErrorTracker::default(); let mut messages: Vec = Vec::new(); for i in 0..MAX_TOTAL_TOOL_ERRORS { tracker.record("bash", &format!("Error: boom {i}"), &mut messages); } assert!(tracker.should_stop()); } #[test] fn conversation_chars_sums_content_only() { let messages = vec![ ChatMessage::system("sys".to_string()), ChatMessage::user("hello world".to_string()), ChatMessage::tool("id".to_string(), "output".to_string()), ]; assert_eq!(conversation_chars(&messages), 3 + 11 + 6); } }