AstroResearch/src/agent/runtime/executor/helpers.rs

// src/agent/runtime/executor/helpers.rs
//
// 执行器辅助类型与函数：PreparedCall, ToolResultMessage, ToolExecutionResult,
// execute_single_tool, process_single_result, save_tool_message_sync。

use sqlx::SqlitePool;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::warn;

use crate::clients::llm::ChatMessage;

use super::AgentStreamEvent;
use crate::agent::hooks::{
    event_label, HookRegistry, PostToolUseContext, PostToolUseFailureContext,
};

#[derive(Debug, Clone)]
pub struct PreparedCall {
    pub tool_call_id: String,
    pub tool_name: String,
    pub args: serde_json::Value,
}

/// 单次工具执行后的消息 + 元数据
pub struct ToolResultMessage {
    pub chat_message: ChatMessage,
    pub was_error: bool,
}

/// 工具执行结果摘要
pub struct ToolExecutionResult {
    /// 每条工具调用对应的 tool_result 消息（供调用方 push 到 messages）
    pub tool_messages: Vec<ToolResultMessage>,
    pub was_cancelled: bool,
    pub had_duplicate: bool,
    /// Hook 注入的附加上下文（PreToolUse + PostToolUse），需注入 LLM 消息列表
    pub hook_contexts: Vec<String>,
    /// Hook 的阻塞错误详情（用于日志和诊断）
    pub blocking_errors: Vec<String>,
}

/// 执行单个工具调用（含超时和取消检测）。
///
/// 从原 `execute_parallel` 的闭包提取，供分区后的批次执行复用。
pub(super) async fn execute_single_tool(
    tool_opt: Option<&dyn crate::agent::tools::AgentTool>,
    args: serde_json::Value,
    tool_ctx: &crate::agent::tools::ToolContext,
    cancelled: &Arc<AtomicBool>,
    timeout_dur: std::time::Duration,
    tool_name: &str,
) -> crate::agent::tools::ToolOutput {
    let tool = match tool_opt {
        Some(t) => t,
        None => return crate::agent::tools::ToolOutput::error(format!("未知工具: {}", tool_name)),
    };

    let interrupt_behavior = tool.interrupt_behavior();
    let is_blocking = interrupt_behavior == crate::agent::tools::InterruptBehavior::Block;

    let tool_fut = tool.execute(args, tool_ctx);
    let cancelled = cancelled.clone();

    let cancel_fut = async {
        loop {
            tokio::time::sleep(std::time::Duration::from_millis(250)).await;
            if !is_blocking && cancelled.load(Ordering::SeqCst) {
                return;
            }
        }
    };

    tokio::select! {
        res = tokio::time::timeout(timeout_dur, tool_fut) => {
            match res {
                Ok(output) => output,
                Err(_) => crate::agent::tools::ToolOutput::error(format!(
                    "工具 {} 执行超时（{}秒）",
                    tool_name,
                    timeout_dur.as_secs()
                )),
            }
        }
        _ = cancel_fut => {
            crate::agent::tools::ToolOutput::error("执行已被用户取消")
        }
    }
}

/// 处理单个工具执行结果（SSE 事件、PostToolUse hooks、持久化）。
///
/// 从原 `execute_parallel` 的结果处理循环提取。
#[allow(clippy::too_many_arguments)]
pub(super) async fn process_single_result(
    tool_call_id: &str,
    tool_name: &str,
    tool_args: &serde_json::Value,
    output: &crate::agent::tools::ToolOutput,
    cancelled_flag: bool,
    exec_start: std::time::Instant,
    tx: &mpsc::UnboundedSender<AgentStreamEvent>,
    hook_registry: &HookRegistry,
    library_dir: &std::path::Path,
    sid: &str,
    agent_name: &str,
    step: usize,
    max_output_chars: usize,
    tool_messages: &mut Vec<ToolResultMessage>,
    additional_contexts: &mut Vec<String>,
    db: &SqlitePool,
    turn_index: i32,
    tool_registry: &crate::agent::tools::ToolRegistry,
) {
    use crate::agent::tools::persist::maybe_persist_tool_result;

    let elapsed_ms = exec_start.elapsed().as_millis() as u64;

    let (is_internal, display_name) = if let Some(tool) = tool_registry.get(tool_name) {
        (tool.is_internal(), tool.display_name().to_string())
    } else {
        (false, tool_name.to_string())
    };

    // SSE 事件 — 立即推送到前端
    let _ = tx.send(AgentStreamEvent::ToolResult {
        tool_call_id: tool_call_id.to_string(),
        name: tool_name.to_string(),
        display_name,
        output: output.content.clone(),
        is_error: output.is_error,
        metadata: output.metadata.clone(),
        step,
        is_internal,
    });

    // 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub
    // 但对于已从磁盘读取内容的工具（如 read_file），跳过持久化以防止级联
    let tool_results_dir = library_dir.join(".agent").join("tool-results");
    let (processed_content, _persisted_path) = if output.skip_persist {
        (output.content.clone(), None)
    } else {
        maybe_persist_tool_result(
            &output.content,
            tool_call_id,
            max_output_chars,
            &tool_results_dir,
        )
        .await
    };

    // PostToolUse hook
    let post_ctx = PostToolUseContext {
        session_id: sid.to_string(),
        agent_name: agent_name.to_string(),
        tool_name: tool_name.to_string(),
        tool_args: tool_args.clone(),
        output_content: processed_content.clone(),
        is_error: output.is_error,
        step,
        elapsed_ms,
    };
    let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
    let final_content = post_result.final_content;

    // 非可信内容包裹（间接 prompt 注入防御）
    let llm_content =
        crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content);

    // 收集 PostToolUse hook 注入的上下文
    if !post_result.tagged_contexts.is_empty() {
        for tc in &post_result.tagged_contexts {
            additional_contexts.push(format!(
                "[Hook: {} | {}] {}",
                tc.hook_name,
                event_label(tc.source_event),
                tc.content,
            ));
        }
    } else {
        for ctx in &post_result.additional_contexts {
            additional_contexts.push(ctx.clone());
        }
    }

    // 收集 PostToolUse 的警告
    for warning in &post_result.warnings {
        additional_contexts.push(format!("[Hook Warning] {}", warning));
    }

    // 事后权限请求（audit trail）
    for (perm_tool, perm) in &post_result.post_permission_requests {
        warn!(
            "[Executor] Hook 事后请求权限: tool={} permission={}",
            perm_tool, perm
        );
    }

    // PostToolUseFailure hook
    if output.is_error {
        let failure_ctx = PostToolUseFailureContext {
            session_id: sid.to_string(),
            agent_name: agent_name.to_string(),
            tool_name: tool_name.to_string(),
            tool_args: tool_args.clone(),
            error_message: output.content.clone(),
            is_interrupt: cancelled_flag,
            step,
            elapsed_ms,
        };
        hook_registry
            .run_on_post_tool_use_failure(&failure_ctx)
            .await;
    }

    // 发送给 LLM 使用包裹后的内容（安全防御）
    let chat_message = ChatMessage::tool_result(tool_call_id, &llm_content);

    // 持久化到数据库（fire-and-forget）
    save_tool_message_sync(db, sid, turn_index, step, &chat_message);

    tool_messages.push(ToolResultMessage {
        chat_message,
        was_error: output.is_error,
    });
}

/// 同步保存 tool 角色消息到数据库。
pub(super) fn save_tool_message_sync(
    db: &SqlitePool,
    session_id: &str,
    turn_index: i32,
    step_index: usize,
    msg: &ChatMessage,
) {
    let db_clone = db.clone();
    let session_id = session_id.to_string();
    let content = msg.text().unwrap_or("").to_string();
    let tool_call_id = msg.tool_call_id.clone();
    // 提前序列化，避免闭包内的生命周期问题
    let metadata_str =
        serde_json::to_string(&serde_json::json!({ "role": "tool" })).unwrap_or_default();
    let raw_json = serde_json::to_string(&msg).unwrap_or_default();
    // fire-and-forget: tool 消息保存失败不影响主流程
    tokio::spawn(async move {
        let token_count = content.len() as i32 / 4;
        if let Err(e) = sqlx::query(
            "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, token_count, metadata, raw_json, agent_name) \
             VALUES (?, ?, ?, 'tool', ?, ?, ?, ?, ?, ?)",
        )
        .bind(&session_id)
        .bind(turn_index)
        .bind(step_index as i32)
        .bind(&content)
        .bind(&tool_call_id)
        .bind(token_count)
        .bind(&metadata_str)
        .bind(&raw_json)
        .bind("lead")
        .execute(&db_clone)
        .await
        {
            warn!("[Executor] 保存 tool 消息失败（非致命）: {}", e);
        }
    });
}