lellm-agent 0.4.9

//! 工具执行器 — 注册、分派、批量执行、并行安全分级。
//!
//! 通过 `ToolCatalog` 消费工具快照，不持有工具所有权。

use std::sync::Arc;

use lellm_core::{
    Message, ParallelSafety, ToolCall, ToolCategory, ToolError, ToolErrorKind, ToolResult,
};

use super::super::event::AgentEvent;
use super::super::retry::RetryPolicy;
use super::{ExecutableTool, ToolCatalog, ToolFn, ToolSnapshot};
use tokio::sync::mpsc::Sender;

/// 批量执行结果 — 长度、顺序、完整性三重保证。
///
/// **不变量：**
/// 1. `results.len() == calls.len()` 永远成立
/// 2. `results[i]` 对应 `calls[i]` 的执行结果（原始顺序）
/// 3. panic 永远被转换成 `ToolResult(is_error: true)`，不会丢失
/// 4. `panicked` 仅作为观测信号，不改变结果完整性
#[derive(Debug)]
pub struct BatchExecutionResult {
    /// 按原始调用顺序排列的工具结果，长度等于输入 calls 长度。
    pub results: Vec<Message>,
    /// 是否有任意 spawned task panic（仅作为观测信号）
    pub panicked: bool,
}

/// 工具执行器 — 按名称分派 ToolCall 到实际工具函数。
///
/// 内部持有 `ToolCatalog`，通过 `snapshot()` 获取冻结工具快照。
/// Clone 为 O(1)（Arc 浅拷贝）。
#[derive(Clone)]
pub struct ToolExecutor {
    catalog: Arc<dyn ToolCatalog>,
    retry_policy: RetryPolicy,
}

impl ToolExecutor {
    /// 绑定工具目录。
    pub fn new(catalog: Arc<dyn ToolCatalog>) -> Self {
        Self {
            catalog,
            retry_policy: RetryPolicy::default(),
        }
    }

    /// 绑定工具目录，使用默认重试策略。
    pub fn with_catalog(catalog: Arc<dyn ToolCatalog>) -> Self {
        Self::new(catalog)
    }

    /// 构造时绑定全局重试策略。
    pub fn with_retry_policy(catalog: Arc<dyn ToolCatalog>, policy: RetryPolicy) -> Self {
        Self {
            catalog,
            retry_policy: policy,
        }
    }

    /// 设置/替换重试策略。
    pub fn set_retry_policy(&mut self, policy: RetryPolicy) {
        self.retry_policy = policy;
    }

    /// 获取当前重试策略的克隆。
    pub fn retry_policy(&self) -> RetryPolicy {
        self.retry_policy.clone()
    }

    /// 获取冻结工具快照。
    ///
    /// 每轮迭代调用一次，固定本轮工具集。
    pub async fn snapshot(&self) -> Arc<ToolSnapshot> {
        self.catalog.snapshot().await
    }

    /// 执行单个工具调用，自带重试。
    ///
    /// 使用预解析的快照执行。
    pub async fn execute_with_snapshot(
        &self,
        call: &ToolCall,
        snapshot: &ToolSnapshot,
    ) -> ToolResult {
        match snapshot.get(&call.name) {
            Some(entry) => {
                let tool_fn = tool_fn_from_reg(entry);
                self.retry_policy
                    .execute_with_retry(&tool_fn, &call.arguments)
                    .await
            }
            None => Err(ToolError::not_found(format!("unknown tool: {}", call.name))),
        }
    }

    /// 执行单个工具调用，自带重试 + Retry 事件发射。
    pub async fn execute_with_emission(
        &self,
        call: &ToolCall,
        snapshot: &ToolSnapshot,
        tx: &Sender<AgentEvent>,
    ) -> ToolResult {
        match snapshot.get(&call.name) {
            Some(entry) => {
                let tool_fn = tool_fn_from_reg(entry);
                self.retry_policy
                    .execute_with_retry_and_emission(&tool_fn, &call.arguments, tx, &call.id)
                    .await
            }
            None => Err(ToolError::not_found(format!("unknown tool: {}", call.name))),
        }
    }
}

/// 使用预解析的快照批量执行 tool_calls。
///
/// 这是动态目录模式的核心执行函数。
///
/// # 用法
///
/// ```ignore
/// let snapshot = executor.snapshot().await;
/// let result = execute_batch_with(&tool_calls, &snapshot, &executor.retry_policy()).await;
/// ```
///
/// # ParallelSafety 契约
///
/// - `Safe`: 全并发（每个 tool 独立 spawn）
/// - `CategoryExclusive(cat)`: 组内串行，组间并发
/// - `Exclusive`: 全串行
///
/// # 一致性保证
///
/// `snapshot` 快照在函数执行期间固定不变，不会因目录刷新而漂移。
pub async fn execute_batch_with(
    calls: &[ToolCall],
    snapshot: &ToolSnapshot,
    retry_policy: &RetryPolicy,
) -> BatchExecutionResult {
    if calls.is_empty() {
        return BatchExecutionResult {
            results: Vec::new(),
            panicked: false,
        };
    }

    // 分组时保留原始索引
    let mut safe_calls: Vec<(usize, ToolCall)> = Vec::new();
    let mut category_calls: std::collections::HashMap<ToolCategory, Vec<(usize, ToolCall)>> =
        std::collections::HashMap::new();
    let mut exclusive_calls: Vec<(usize, ToolCall)> = Vec::new();

    for (idx, call) in calls.iter().enumerate() {
        let safety = snapshot
            .get(&call.name)
            .map(|t| t.safety.clone())
            .unwrap_or(ParallelSafety::Exclusive);

        match safety {
            ParallelSafety::Safe => safe_calls.push((idx, call.clone())),
            ParallelSafety::CategoryExclusive => {
                if let Some(cat) = snapshot.get(&call.name).and_then(|t| t.category.clone()) {
                    category_calls
                        .entry(cat)
                        .or_default()
                        .push((idx, call.clone()));
                } else {
                    exclusive_calls.push((idx, call.clone()));
                }
            }
            ParallelSafety::Exclusive => exclusive_calls.push((idx, call.clone())),
        }
    }

    // 构建 group handles
    let mut group_handles: Vec<tokio::task::JoinHandle<Vec<(usize, Message)>>> = Vec::new();
    let mut group_indices: Vec<Vec<usize>> = Vec::new();

    let snapshot = Arc::new(snapshot.clone_for_spawn());
    let retry_policy = retry_policy.clone();

    // Safe: 每个 tool 独立 spawn（全并发）
    if !safe_calls.is_empty() {
        let s = Arc::clone(&snapshot);
        let rp = retry_policy.clone();
        let indices: Vec<usize> = safe_calls.iter().map(|(i, _)| *i).collect();
        group_handles.push(tokio::spawn(async move {
            run_parallel_indexed_with(&s, &rp, safe_calls).await
        }));
        group_indices.push(indices);
    }

    // CategoryExclusive: 按 category 分组，组内串行、组间并发
    for group_calls in category_calls.into_values() {
        let s = Arc::clone(&snapshot);
        let rp = retry_policy.clone();
        let indices: Vec<usize> = group_calls.iter().map(|(i, _)| *i).collect();
        group_handles.push(tokio::spawn(async move {
            run_serial_indexed_with(&s, &rp, group_calls).await
        }));
        group_indices.push(indices);
    }

    // Exclusive: 全部串行，一个 task
    if !exclusive_calls.is_empty() {
        let s = Arc::clone(&snapshot);
        let rp = retry_policy.clone();
        let indices: Vec<usize> = exclusive_calls.iter().map(|(i, _)| *i).collect();
        group_handles.push(tokio::spawn(async move {
            run_serial_indexed_with(&s, &rp, exclusive_calls).await
        }));
        group_indices.push(indices);
    }

    // 按原始索引回填结果；panic 的 group 按索引列表精准回填错误
    let mut results: Vec<Option<Message>> = vec![None; calls.len()];
    let mut panicked = false;
    let all_handles = futures_util::future::join_all(group_handles).await;

    for (handle_result, indices) in all_handles.into_iter().zip(group_indices.into_iter()) {
        match handle_result {
            Ok(indexed_messages) => {
                for (idx, msg) in indexed_messages {
                    results[idx] = Some(msg);
                }
            }
            Err(join_err) => {
                panicked = true;
                for idx in indices {
                    let call = &calls[idx];
                    results[idx] = Some(Message::tool_result(
                        call,
                        &Err(ToolError {
                            kind: ToolErrorKind::Internal,
                            message: format!("tool group task panicked: {join_err}"),
                        }),
                    ));
                }
            }
        }
    }

    BatchExecutionResult {
        results: results.into_iter().flatten().collect(),
        panicked,
    }
}

// ─── 内部辅助：快照克隆 ──────────────────────────────────────────

impl ToolSnapshot {
    /// 克隆内部工具映射，供 spawn 使用。
    pub fn clone_for_spawn(&self) -> Arc<indexmap::IndexMap<String, ExecutableTool>> {
        self.tools.clone()
    }
}

// ─── Safe group: 全并发 ──────────────────────────────────────────

async fn run_parallel_indexed_with(
    tools: &Arc<indexmap::IndexMap<String, ExecutableTool>>,
    retry_policy: &RetryPolicy,
    calls: Vec<(usize, ToolCall)>,
) -> Vec<(usize, Message)> {
    let handles: Vec<_> = calls
        .iter()
        .map(|(idx, call)| {
            let tools = Arc::clone(tools);
            let rp = retry_policy.clone();
            let call = call.clone();
            let idx = *idx;
            tokio::spawn(async move {
                let result = match tools.get(&call.name) {
                    Some(entry) => {
                        let tool_fn = tool_fn_from_reg(entry);
                        rp.execute_with_retry(&tool_fn, &call.arguments).await
                    }
                    None => Err(ToolError::not_found(format!("unknown tool: {}", call.name))),
                };
                (idx, Message::tool_result(&call, &result))
            })
        })
        .collect();

    let raw = futures_util::future::join_all(handles).await;
    raw.into_iter()
        .zip(calls.into_iter())
        .map(|(h, (idx, call))| match h {
            Ok((_, msg)) => (idx, msg),
            Err(join_err) => (
                idx,
                Message::tool_result(
                    &call,
                    &Err(ToolError {
                        kind: ToolErrorKind::Internal,
                        message: format!("tool '{}' task panicked: {join_err}", call.name),
                    }),
                ),
            ),
        })
        .collect()
}

// ─── 组内串行 ────────────────────────────────────────────────────

async fn run_serial_indexed_with(
    tools: &Arc<indexmap::IndexMap<String, ExecutableTool>>,
    retry_policy: &RetryPolicy,
    calls: Vec<(usize, ToolCall)>,
) -> Vec<(usize, Message)> {
    let mut results = Vec::with_capacity(calls.len());
    for (idx, call) in calls {
        let exec_result = match tools.get(&call.name) {
            Some(entry) => {
                let tool_fn = tool_fn_from_reg(entry);
                retry_policy
                    .execute_with_retry(&tool_fn, &call.arguments)
                    .await
            }
            None => Err(ToolError::not_found(format!("unknown tool: {}", call.name))),
        };
        results.push((idx, Message::tool_result(&call, &exec_result)));
    }
    results
}

// ─── Helper: 从 ExecutableTool 获取 ToolFn ──────────────────────

fn tool_fn_from_reg(entry: &ExecutableTool) -> ToolFn {
    let entry = entry.clone();
    Arc::new(move |args: &serde_json::Value| entry.execute(args))
}