diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index 48d4fdb97c..ab82ad9aff 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/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 { + 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, @@ -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 = 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, Option), 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, max_tokens: Option, on_event: &(impl Fn(AgentEvent) + Send + Sync), + cancelled: &Notify, force: bool, -) -> Option { +) -> 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 { diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index fc9b289ab7..034f3a23c2 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/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.