opencrabs 0.3.62

//! Concurrent execution of an auto-approved tool batch (#361).
//!
//! When one assistant response carries several tool calls and none of them
//! needs interactive approval, they run concurrently, capped by
//! `[agent] max_concurrent`. Result blocks are reassembled in the original
//! tool_use order (providers require results to match the request order);
//! per-tool bookkeeping runs sequentially after the join so no shared
//! mutable state crosses the parallel region. Batches with any
//! approval-gated tool, single-tool batches, and `max_concurrent = 1` all
//! take the existing sequential path untouched.

use super::tool_loop::{
    build_tool_result_content, extract_path_for_recent_buffer, strip_ansi_output,
};
use super::types::{ProgressCallback, ProgressEvent};
use crate::brain::provider::ContentBlock;
use crate::brain::tools::r#trait::ToolExecutionContext;
use futures::StreamExt;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

/// Outcome of one concurrently executed tool, carried back for in-order
/// assembly and bookkeeping.
struct ToolOutcome {
    tool_id: String,
    tool_name: String,
    tool_input: Value,
    success: bool,
    content: String,
    images: Vec<(String, String)>,
    /// Tool never ran (unknown name / invalid args): a model miss, bucketed
    /// as discovery_miss instead of a tool failure.
    pre_execution_miss: bool,
    /// True when the execution returned Err (vs a failed ToolResult).
    exec_error: bool,
}

/// What the batch produced, in original order.
pub(crate) struct ParallelBatchOutcome {
    pub results: Vec<ContentBlock>,
    pub descriptions: Vec<String>,
    pub outputs: Vec<(bool, String)>,
    /// Count of successful tool runs (drives the phantom-detection counters).
    pub successes: usize,
    /// The cancel token fired mid-batch; later tools did not run.
    pub cancelled: bool,
}

impl super::AgentService {
    /// True when every tool in the batch is auto-approved under the current
    /// flags, so the whole batch can run concurrently. Any approval-gated
    /// tool keeps the batch on the sequential path (interactive prompts
    /// cannot be parallelized sensibly).
    pub(crate) fn batch_is_parallel_eligible(
        &self,
        tool_uses: &[(String, String, Value)],
        tool_context: &ToolExecutionContext,
        has_override_approval: bool,
    ) -> bool {
        if tool_uses.len() < 2 || self.max_concurrent < 2 {
            return false;
        }
        tool_uses.iter().all(|(_, name, input)| {
            match self.tool_registry.get(name) {
                Some(tool) => {
                    let needs_approval = tool.requires_approval_for_input(input)
                        && (!self.auto_approve_tools || has_override_approval)
                        && !tool_context.auto_approve;
                    !needs_approval
                }
                // Unknown tool: the sequential path has dedicated handling.
                None => false,
            }
        })
    }

    /// Run the batch concurrently (capped at `max_concurrent`), preserving
    /// result order. Mirrors the sequential non-approval branch: same
    /// progress events, feedback recording, dashboard records, recent-path
    /// capture, and result/output shapes.
    pub(crate) async fn execute_tools_parallel(
        &self,
        session_id: Uuid,
        tool_uses: Vec<(String, String, Value)>,
        tool_context: &ToolExecutionContext,
        cancel_token: Option<&CancellationToken>,
        progress_callback: Option<&ProgressCallback>,
        assistant_msg_id: Uuid,
    ) -> ParallelBatchOutcome {
        let mut approved_context = tool_context.clone();
        approved_context.auto_approve = true;
        let total = tool_uses.len();
        tracing::info!(
            "[TOOL_EXEC] ⚡ Executing {total} tools concurrently (max_concurrent={})",
            self.max_concurrent
        );

        let mut stream = futures::stream::iter(tool_uses.into_iter().map(
            |(tool_id, tool_name, tool_input)| {
                let ctx = approved_context.clone();
                let cb = progress_callback.cloned();
                async move {
                    if let Some(ref cb) = cb {
                        cb(
                            session_id,
                            ProgressEvent::ToolStarted {
                                tool_name: tool_name.clone(),
                                tool_input: tool_input.clone(),
                            },
                        );
                    }
                    let exec = self
                        .tool_registry
                        .execute(&tool_name, tool_input.clone(), &ctx)
                        .await;
                    let outcome = match exec {
                        Ok(result) => ToolOutcome {
                            tool_id,
                            tool_name,
                            tool_input,
                            success: result.success,
                            content: build_tool_result_content(
                                result.success,
                                result.error,
                                &result.output,
                            ),
                            images: result.images,
                            pre_execution_miss: false,
                            exec_error: false,
                        },
                        Err(e) => ToolOutcome {
                            tool_id,
                            tool_name,
                            tool_input,
                            success: false,
                            content: format!("Tool execution error: {}", e),
                            images: Vec::new(),
                            pre_execution_miss: e.is_pre_execution_miss(),
                            exec_error: true,
                        },
                    };
                    // Completion event fires as each tool finishes (real-time
                    // display), even though bookkeeping waits for the join.
                    if let Some(ref cb) = cb {
                        let summary: String = strip_ansi_output(&outcome.content)
                            .chars()
                            .take(2000)
                            .collect();
                        cb(
                            session_id,
                            ProgressEvent::ToolCompleted {
                                tool_name: outcome.tool_name.clone(),
                                tool_input: outcome.tool_input.clone(),
                                success: outcome.success,
                                summary,
                            },
                        );
                    }
                    outcome
                }
            },
        ))
        // `buffered` polls up to N futures concurrently while yielding
        // outcomes in ORIGINAL order — exactly the ordering contract
        // providers need for tool_result blocks.
        .buffered(self.max_concurrent);

        let mut outcomes: Vec<ToolOutcome> = Vec::with_capacity(total);
        let mut cancelled = false;
        loop {
            let next = tokio::select! {
                biased;
                _ = async {
                    match cancel_token {
                        Some(t) => t.cancelled().await,
                        None => std::future::pending().await,
                    }
                } => {
                    tracing::warn!(
                        "🛑 Parallel tool batch cancelled after {}/{total} outcomes",
                        outcomes.len()
                    );
                    cancelled = true;
                    break;
                }
                n = stream.next() => n,
            };
            match next {
                Some(outcome) => outcomes.push(outcome),
                None => break,
            }
        }
        drop(stream); // aborts any still-running futures on cancellation

        // ── In-order bookkeeping and assembly (sequential, single-threaded) ──
        let mut out = ParallelBatchOutcome {
            results: Vec::new(),
            descriptions: Vec::new(),
            outputs: Vec::new(),
            successes: 0,
            cancelled,
        };
        for o in outcomes {
            out.descriptions
                .push(Self::format_tool_summary(&o.tool_name, &o.tool_input));
            if o.success {
                tracing::info!(
                    "[TOOL_EXEC] ✅ Tool '{}' executed successfully, output_len={}",
                    o.tool_name,
                    o.content.len()
                );
                out.successes += 1;
                if let Some(p) = extract_path_for_recent_buffer(
                    &o.tool_name,
                    &o.tool_input,
                    &approved_context.working_directory,
                ) {
                    self.record_recent_path(&approved_context.working_directory, &p);
                }
            } else {
                tracing::error!(
                    "[TOOL_EXEC] ❌ Tool '{}' failed: {}",
                    o.tool_name,
                    o.content.chars().take(200).collect::<String>()
                );
            }
            if o.pre_execution_miss {
                self.record_tool_discovery_miss(
                    session_id,
                    &o.tool_name,
                    Some(&o.tool_input),
                    Some(&o.content),
                );
            } else {
                self.record_tool_feedback(
                    session_id,
                    &o.tool_name,
                    Some(&o.tool_input),
                    o.success,
                    if o.success { None } else { Some(&o.content) },
                );
            }
            if let Some(pool) = crate::db::global_pool() {
                let tool_repo = crate::db::repository::ToolExecutionRepository::new(pool.clone());
                let exec_id = Uuid::new_v4().to_string();
                let mid = assistant_msg_id.to_string();
                let sid = session_id.to_string();
                let tname = o.tool_name.clone();
                let status = if o.exec_error {
                    "error"
                } else if o.success {
                    "success"
                } else {
                    "error"
                };
                tokio::spawn(async move {
                    if let Err(e) = tool_repo.record(&exec_id, &mid, &sid, &tname, status).await {
                        tracing::error!("[TOOL_EXEC] Failed to record tool execution: {}", e);
                    }
                });
            }
            let output_summary: String = strip_ansi_output(&o.content).chars().take(2000).collect();
            out.outputs.push((o.success, output_summary));
            out.results.push(ContentBlock::ToolResult {
                tool_use_id: o.tool_id,
                content: o.content,
                is_error: Some(!o.success),
            });
            for (media_type, data) in o.images {
                out.results.push(ContentBlock::Image {
                    source: crate::brain::provider::ImageSource::Base64 { media_type, data },
                });
            }
        }
        out
    }
}