fix(agent-loop): preserve output on exhaustion, interruption, and cancellation

crates/dbx-core/src/agent_loop.rs

@@ -1,5 +1,5 @@
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use futures::future::join_all;
 use futures::FutureExt;
@@ -13,15 +13,38 @@ use crate::ai_cli_agent::CliAgentCommandSpec;
 use crate::connection::AppState;
 use crate::models::connection::DatabaseType;
 use crate::token_usage::TokenUsage;
-use tokio::sync::Mutex;
 
 /// Maximum number of agent loop turns to prevent infinite loops.
-const MAX_AGENT_TURNS: u32 = 10;
+const MAX_AGENT_TURNS: u32 = 30;
+const AGENT_CANCELLED_ERROR: &str = "Agent loop cancelled";
 const MAX_TOOL_RESULT_CONTEXT_CHARS: usize = 12_000;
 const TOOL_RESULT_HEAD_CHARS: usize = 4_000;
 const TOOL_RESULT_TAIL_CHARS: usize = 4_000;
 const TOOL_RESULT_SAMPLE_ITEMS: usize = 5;
 
+fn take_text(m: &std::sync::Mutex<String>) -> String {
+    m.lock().unwrap_or_else(|e| e.into_inner()).clone()
+}
+
+enum LoopExit {
+    Completed,
+    Cancelled,
+    Interrupted(String),
+    Exhausted,
+}
+
+impl LoopExit {
+    fn should_break_turns(&self) -> bool {
+        matches!(self, LoopExit::Cancelled | LoopExit::Interrupted(_))
+    }
+}
+
+enum CompactResult {
+    Skipped,
+    Compacted,
+    Cancelled,
+}
+
 /// Context for an agent loop run.
 pub struct AgentLoopContext {
     pub state: Arc<AppState>,
@@ -103,17 +126,34 @@ pub async fn run_agent_loop(
     let tools = if is_agent_mode { agent_tools::all_tools(agent_ctx.db_type) } else { agent_tools::read_only_tools() };
     let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
     let mut final_text = String::new();
+    let mut loop_exit = LoopExit::Exhausted;
     let mut total_usage = TokenUsage::default();
 
     for turn in 0..MAX_AGENT_TURNS {
         // Check for cancellation before each turn
         if cancelled.notified().now_or_never().is_some() {
-            on_event(AgentEvent::Error { message: "Agent loop cancelled".to_string() });
+            loop_exit = LoopExit::Cancelled;
             break;
         }
 
         // Check and maybe compact context
-        maybe_compact(config, system_prompt, &tools, &mut conversation_messages, max_tokens, &on_event, false).await;
+        if matches!(
+            maybe_compact(
+                config,
+                system_prompt,
+                &tools,
+                &mut conversation_messages,
+                max_tokens,
+                &on_event,
+                cancelled,
+                false,
+            )
+            .await,
+            CompactResult::Cancelled
+        ) {
+            loop_exit = LoopExit::Cancelled;
+            break;
+        }
 
         on_event(AgentEvent::TurnStart { turn });
 
@@ -138,9 +178,7 @@ pub async fn run_agent_loop(
             let on_chunk = move |chunk: AiStreamChunk| {
                 if !chunk.delta.is_empty() {
                     emitted.store(true, Ordering::Relaxed);
-                    if let Ok(mut text) = acc.try_lock() {
-                        text.push_str(&chunk.delta);
-                    }
+                    acc.lock().unwrap_or_else(|e| e.into_inner()).push_str(&chunk.delta);
                     on_event2(AgentEvent::TextDelta { delta: chunk.delta.clone() });
                 }
                 if let Some(ref reasoning) = chunk.reasoning_delta {
@@ -151,7 +189,7 @@ pub async fn run_agent_loop(
 
             match stream_with_tools(config, &request, &session_id, &tools, cancelled, on_chunk).await {
                 Ok((tool_calls, usage)) => {
-                    let accumulated_text = accumulated_text.lock().await.clone();
+                    let accumulated_text = take_text(&accumulated_text);
                     stream_result = Some((tool_calls, usage, accumulated_text));
                     break;
                 }
@@ -166,18 +204,43 @@ pub async fn run_agent_loop(
                         &mut conversation_messages,
                         max_tokens,
                         &on_event,
+                        cancelled,
                         true,
                     )
                     .await;
-                    if compacted.is_some() {
-                        continue;
+                    match compacted {
+                        CompactResult::Compacted => continue,
+                        CompactResult::Cancelled => {
+                            loop_exit = LoopExit::Cancelled;
+                            break;
+                        }
+                        CompactResult::Skipped => {
+                            final_text = take_text(&accumulated_text);
+                            loop_exit =
+                                LoopExit::Interrupted(last_stream_error.take().unwrap_or_else(|| {
+                                    "LLM request failed after context compaction retry".to_string()
+                                }));
+                        }
                     }
                     break;
                 }
-                Err(err) => return Err(err),
+                Err(err) if err == AGENT_CANCELLED_ERROR => {
+                    final_text = take_text(&accumulated_text);
+                    loop_exit = LoopExit::Cancelled;
+                    break;
+                }
+                Err(err) => {
+                    final_text = take_text(&accumulated_text);
+                    loop_exit = LoopExit::Interrupted(err);
+                    break;
+                }
             }
         }
 
+        if loop_exit.should_break_turns() {
+            break;
+        }
+
         let Some((collected_tool_calls, turn_usage, accumulated_text)) = stream_result else {
             return Err(
                 last_stream_error.unwrap_or_else(|| "LLM request failed after context compaction retry".to_string())
@@ -204,6 +267,7 @@ pub async fn run_agent_loop(
         if collected_tool_calls.is_empty() {
             // No tool calls -- we're done
             final_text = accumulated_text;
+            loop_exit = LoopExit::Completed;
             break;
         }
 
@@ -284,6 +348,39 @@ pub async fn run_agent_loop(
         final_text = accumulated_text;
     }
 
+    match loop_exit {
+        LoopExit::Completed => {}
+        LoopExit::Cancelled => {
+            let message = if final_text.trim().is_empty() {
+                "Agent run was cancelled before producing output.".to_string()
+            } else {
+                "\n\nAgent run was cancelled. Partial output above was preserved.".to_string()
+            };
+            on_event(AgentEvent::TextDelta { delta: message.clone() });
+            final_text.push_str(&message);
+        }
+        LoopExit::Interrupted(error) => {
+            let message = if final_text.trim().is_empty() {
+                format!("Agent stream stopped before completion: {error}.")
+            } else {
+                format!("\n\nAgent stream stopped before completion: {error}. Partial output above was preserved.")
+            };
+            on_event(AgentEvent::TextDelta { delta: message.clone() });
+            final_text.push_str(&message);
+        }
+        LoopExit::Exhausted => {
+            let message = if final_text.trim().is_empty() {
+                format!("Agent reached the {MAX_AGENT_TURNS}-turn safety limit before producing output. Send Continue to let the agent keep working.")
+            } else {
+                format!(
+                    "\n\nAgent reached the {MAX_AGENT_TURNS}-turn safety limit before a final answer. The partial output above was preserved; send Continue to let the agent keep working."
+                )
+            };
+            on_event(AgentEvent::TextDelta { delta: message.clone() });
+            final_text.push_str(&message);
+        }
+    }
+
     on_event(AgentEvent::AgentEnd {
         input_tokens: if total_usage.input_tokens > 0 { Some(total_usage.input_tokens) } else { None },
         output_tokens: if total_usage.output_tokens > 0 { Some(total_usage.output_tokens) } else { None },
@@ -325,7 +422,7 @@ async fn stream_with_tools(
 ) -> Result<(Vec<ToolCall>, Option<TokenUsage>), String> {
     // Return early if the user cancelled before the LLM call started.
     if cancelled.notified().now_or_never().is_some() {
-        return Err("Agent loop cancelled".to_string());
+        return Err(AGENT_CANCELLED_ERROR.to_string());
     }
 
     ai::stream_with_tools(config, request, session_id, tools, cancelled, on_chunk).await
@@ -521,18 +618,19 @@ async fn maybe_compact(
     messages: &mut Vec<AiMessage>,
     max_tokens: Option<u32>,
     on_event: &(impl Fn(AgentEvent) + Send + Sync),
+    cancelled: &Notify,
     force: bool,
-) -> Option<u32> {
+) -> CompactResult {
     let window = config.context_window.unwrap_or_else(|| context_window_for_model(&config.model));
     let budget = prompt_budget(window, max_tokens);
     let estimated_before = estimate_current_prompt_tokens(system_prompt, tools, messages);
 
     if !force && estimated_before <= budget {
-        return None;
+        return CompactResult::Skipped;
     }
 
     if messages.len() <= 2 {
-        return None;
+        return CompactResult::Skipped;
     }
 
     // Find cut point: keep a dynamic budget of recent messages and summarize older context.
@@ -558,7 +656,7 @@ async fn maybe_compact(
     // Always keep messages[0] (the original user question) verbatim outside the summary.
     // Only summarize messages[1..cut].
     if cut <= 1 {
-        return None;
+        return CompactResult::Skipped;
     }
     let summary_start = 1usize;
 
@@ -579,9 +677,15 @@ async fn maybe_compact(
         temperature: Some(0.1),
     };
 
-    let summary = match ai::complete(&summary_request).await {
-        Ok(s) => s,
-        Err(_) => fallback_summary(messages, cut),
+    let summary = match cancelled.notified().now_or_never() {
+        Some(_) => return CompactResult::Cancelled,
+        None => match tokio::select! {
+            result = ai::complete(&summary_request) => result,
+            _ = cancelled.notified() => return CompactResult::Cancelled,
+        } {
+            Ok(s) => s,
+            Err(_) => fallback_summary(messages, cut),
+        },
     };
 
     let summary = if validate_summary(&summary) { summary } else { fallback_summary(messages, cut) };
@@ -612,7 +716,7 @@ async fn maybe_compact(
         estimated_before,
         estimated_after,
     });
-    Some(summary_tokens)
+    CompactResult::Compacted
 }
 
 fn compact_tool_result_for_context(tool_name: &str, content: &str) -> String {

crates/dbx-core/src/agent_tools.rs

@@ -85,7 +85,9 @@ fn get_columns_tool() -> ToolDefinition {
             "required": ["table"]
         }),
         read_only: true,
-        parallel_ok: true,
+        // get_columns runs sequentially: concurrent metadata queries can exhaust
+        // single-connection drivers (e.g. DuckDB), causing cascading tool errors.
+        parallel_ok: false,
     }
 }
 /// execute_query tool definition.