AstroResearch/src/agent/subagent.rs

// src/agent/subagent.rs
//
// 子代理运行器 — 上下文隔离子代理（参考 Claude Code s04 Subagents）。
//
// 父代理通过 subagent 工具将子任务委托给子代理执行。
// 子代理拥有:
//   - 全新的 messages 上下文（不包含父代理的中间工具调用）
//   - 完整的工具访问权限（与父代理共享 ToolRegistry）
//   - 独立的 ReAct 循环
//   - 完整的 Hook 管道（PreToolUse/PostToolUse）和权限检查
//
// 子代理只返回最终文本摘要给父代理，中间工具调用不污染父上下文。

use std::sync::Arc;
use tokio::sync::mpsc::UnboundedSender;
use tracing::{info, warn};

use super::compact;
use super::hooks::{
    HookRegistry, PostToolUseContext, PreToolUseContext, SubagentStartContext, SubagentStopContext,
};
use super::runtime::permission::{PermissionChecker, PermissionResult};
use super::runtime::{AgentConfig, AgentStreamEvent};
use super::tools::{ToolContext, ToolOutput, ToolRegistry};
use crate::api::AppState;
use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent, ToolDefinition};

/// 子代理运行器
pub struct SubAgentRunner {
    app_state: Arc<AppState>,
    config: AgentConfig,
    tool_registry: ToolRegistry,
    /// 可选的 Hook 注册表（用于 PreToolUse/PostToolUse 生命周期事件）
    hook_registry: Option<Arc<HookRegistry>>,
    /// 权限检查器
    permission_checker: Arc<PermissionChecker>,
    /// 可选的进度发送器（用于向父代理报告中间步骤）
    progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
    /// 父代理的会话 ID（用于子代理消息的数据库持久化）
    parent_session_id: String,
}

impl SubAgentRunner {
    /// 创建新的子代理运行器（无 hook/permission/progress）。
    pub fn new(app_state: Arc<AppState>) -> Self {
        let skill_registry = app_state.skill_registry.clone();
        SubAgentRunner {
            app_state,
            config: AgentConfig::default(),
            tool_registry: ToolRegistry::new(skill_registry),
            hook_registry: None,
            permission_checker: Arc::new(PermissionChecker::new()),
            progress_tx: None,
            parent_session_id: String::new(),
        }
    }

    /// 创建带完整 hooks/permissions/progress 的子代理运行器。
    pub fn new_with_hooks(
        app_state: Arc<AppState>,
        hook_registry: Option<Arc<HookRegistry>>,
        permission_checker: Arc<PermissionChecker>,
        progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
    ) -> Self {
        let skill_registry = app_state.skill_registry.clone();
        SubAgentRunner {
            app_state,
            config: AgentConfig::default(),
            tool_registry: ToolRegistry::new(skill_registry),
            hook_registry,
            permission_checker,
            progress_tx,
            parent_session_id: String::new(),
        }
    }

    /// 设置父代理会话 ID（调用者应在 run 之前设置）
    pub fn with_parent_session(mut self, session_id: String) -> Self {
        self.parent_session_id = session_id;
        self
    }

    /// 设置是否启用 LLM 思考模式
    pub fn with_thinking(mut self, enable: bool) -> Self {
        self.config.enable_thinking = enable;
        self
    }

    /// 使用自定义 ToolRegistry 创建子代理运行器。
    /// 用于受限场景（如记忆提取子代理仅需只读 + save_memory）。
    pub fn new_with_registry(app_state: Arc<AppState>, tool_registry: ToolRegistry) -> Self {
        SubAgentRunner {
            app_state,
            config: AgentConfig::default(),
            tool_registry,
            hook_registry: None,
            permission_checker: Arc::new(PermissionChecker::new()),
            progress_tx: None,
            parent_session_id: String::new(),
        }
    }

    /// 创建带完整 hooks/permissions + 自定义 ToolRegistry 的子代理运行器。
    pub fn new_with_registry_and_hooks(
        app_state: Arc<AppState>,
        tool_registry: ToolRegistry,
        hook_registry: Option<Arc<HookRegistry>>,
        permission_checker: Arc<PermissionChecker>,
        progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
    ) -> Self {
        SubAgentRunner {
            app_state,
            config: AgentConfig::default(),
            tool_registry,
            hook_registry,
            permission_checker,
            progress_tx,
            parent_session_id: String::new(),
        }
    }

    /// 运行子代理的 ReAct 循环，返回最终文本摘要。
    ///
    /// # Arguments
    /// * `system_prompt` - 子代理的系统提示词
    /// * `research_prompt` - 要执行的研究任务描述
    /// * `max_steps` - 子代理最大推理步数（默认 5）
    pub async fn run(
        &self,
        system_prompt: &str,
        research_prompt: &str,
        max_steps: usize,
    ) -> ToolOutput {
        let subagent_name = format!("sub_{}", &uuid::Uuid::new_v4().to_string()[..8]);

        // OnSubagentStart hook
        if let Some(ref registry) = self.hook_registry {
            registry
                .run_on_subagent_start(&SubagentStartContext {
                    parent_session_id: self.parent_session_id.clone(),
                    subagent_name: subagent_name.clone(),
                    prompt: research_prompt.to_string(),
                })
                .await;
        }

        // 保存子代理的 system prompt + user prompt 到数据库
        self.save_subagent_message(
            &subagent_name,
            0,
            0,
            "system",
            system_prompt,
            None,
            None,
            None,
        );

        self.save_subagent_message(
            &subagent_name,
            0,
            1,
            "user",
            research_prompt,
            None,
            None,
            None,
        );

        // 执行实际工作并捕获结果，以便触发 OnSubagentStop hook
        let result = self
            .run_inner(system_prompt, research_prompt, max_steps, &subagent_name)
            .await;
        let (is_error, result_summary) = if result.is_error {
            (true, result.content.clone())
        } else {
            (false, result.content.chars().take(200).collect())
        };

        // 保存最终结果
        self.save_subagent_message(
            &subagent_name,
            0,
            max_steps as i32 + 1,
            "assistant",
            &result.content,
            None,
            None,
            None,
        );

        if let Some(ref registry) = self.hook_registry {
            registry
                .run_on_subagent_stop(&SubagentStopContext {
                    parent_session_id: self.parent_session_id.clone(),
                    subagent_name: subagent_name.clone(),
                    result_summary,
                    steps: max_steps,
                    is_error,
                })
                .await;
        }

        result
    }

    /// 保存子代理消息到 agent_messages 表
    #[allow(clippy::too_many_arguments)]
    fn save_subagent_message(
        &self,
        agent_name: &str,
        turn_index: i32,
        step_index: i32,
        role: &str,
        content: &str,
        thought: Option<&str>,
        tool_calls_json: Option<&str>,
        tool_call_id: Option<&str>,
    ) {
        if self.parent_session_id.is_empty() {
            return;
        }
        let db = self.app_state.db.clone();
        let session_id = self.parent_session_id.clone();
        let agent = agent_name.to_string();
        let role_owned = role.to_string();
        let content_owned = content.to_string();
        let thought_owned = thought.map(|s| s.to_string());
        let tc_json = tool_calls_json.map(|s| s.to_string());
        let tc_id = tool_call_id.map(|s| s.to_string());
        let metadata = serde_json::json!({ "agent": agent, "is_subagent": true });
        let metadata_str = serde_json::to_string(&metadata).unwrap_or_default();

        tokio::spawn(async move {
            let token_count = content_owned.len() as i32 / 4;
            let _ = sqlx::query(
                "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, agent_name) \
                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            )
            .bind(&session_id)
            .bind(turn_index)
            .bind(step_index)
            .bind(&role_owned)
            .bind(&content_owned)
            .bind(&thought_owned)
            .bind(&tc_json)
            .bind(&tc_id)
            .bind(token_count)
            .bind(&metadata_str)
            .bind(&agent)
            .execute(&db)
            .await;
        });
    }

    /// 实际执行逻辑（提取为内部方法以便 hook 包装）
    async fn run_inner(
        &self,
        system_prompt: &str,
        research_prompt: &str,
        max_steps: usize,
        subagent_name: &str,
    ) -> ToolOutput {
        let llm = &self.app_state.llm;
        let tool_defs = self.tool_registry.definitions();

        // 全新上下文
        let mut messages = vec![
            ChatMessage::system(system_prompt),
            ChatMessage::user(research_prompt),
        ];

        // 跟踪工具调用防止死循环
        let mut last_call: Option<(String, String)> = None;
        let mut consecutive_count: usize = 0;
        let duplicate_threshold: usize = 3;

        for step in 1..=max_steps {
            // 上下文压缩检查
            let est_tokens: usize = messages
                .iter()
                .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4)
                .sum();
            if est_tokens > self.config.context_char_limit * 3 / 2 {
                info!(
                    "[SubAgent] 上下文超限 (est. {} tokens)，触发压缩",
                    est_tokens
                );
                compact::compress_context_with_hooks(
                    &mut messages,
                    llm,
                    self.config.context_char_limit,
                    subagent_name,
                    self.hook_registry.as_ref().map(|a| a.as_ref()),
                )
                .await;
            }

            // LLM 流式调用
            let mut stream_rx = match llm
                .chat_stream(&messages, &tool_defs, self.config.enable_thinking)
                .await
            {
                Ok(rx) => rx,
                Err(e) => {
                    warn!("[SubAgent] LLM stream 失败: {}", e);
                    return ToolOutput::error(format!("子代理 LLM 调用失败: {}", e));
                }
            };

            let mut accumulated_content = String::new();
            let mut accumulated_reasoning = String::new();
            let mut accumulated_tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;
            let mut activity_log: Vec<String> = Vec::new();

            while let Some(event) = stream_rx.recv().await {
                match event {
                    StreamEvent::ReasoningDelta(delta) => {
                        accumulated_reasoning.push_str(&delta);
                        // 转发子代理思考过程到父代理
                        if let Some(ref tx) = self.progress_tx {
                            let _ = tx.send(AgentStreamEvent::Thought {
                                content: format!("[子代理] {}", accumulated_reasoning),
                                step,
                            });
                        }
                    }
                    StreamEvent::TextDelta(delta) => {
                        accumulated_content.push_str(&delta);
                    }
                    StreamEvent::ToolCallsComplete(tool_calls) => {
                        // 确保每个工具调用有唯一 ID（LLM 可能不返回 id）
                        let fixed_tool_calls: Vec<crate::clients::llm::ToolCall> = tool_calls
                            .into_iter()
                            .map(|tc| {
                                let id = if tc.id.is_empty() {
                                    format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8])
                                } else {
                                    tc.id
                                };
                                activity_log.push(format!("🔧 调用工具: {}", tc.function.name));
                                crate::clients::llm::ToolCall { id, ..tc }
                            })
                            .collect();
                        accumulated_tool_calls = Some(fixed_tool_calls);
                    }
                    StreamEvent::Done => break,
                    StreamEvent::Error(e) => {
                        warn!("[SubAgent] 流式错误: {}", e);
                        return ToolOutput::error(format!("子代理流式错误: {}", e));
                    }
                    _ => {}
                }
            }

            // 记录思考过程
            if !accumulated_reasoning.is_empty() {
                activity_log.push(format!(
                    "💭 思考: {}",
                    accumulated_reasoning.chars().take(300).collect::<String>()
                ));
            }

            // 无工具调用 = 最终回答
            let tool_calls = match accumulated_tool_calls {
                Some(ref tc) if !tc.is_empty() => tc.clone(),
                _ => {
                    // 转发子代理结论到父代理（作为 thought 显示在时间线，不污染 finalAnswer）
                    if let Some(ref tx) = self.progress_tx {
                        let _ = tx.send(AgentStreamEvent::Thought {
                            content: accumulated_content.clone(),
                            step,
                        });
                    }
                    let content_len = accumulated_content.len();
                    info!("[SubAgent] 子代理完成，返回 {} 字符摘要", content_len);
                    let summary = format!(
                        "[子代理活动记录]\n\n{}\n\n[子代理结论]\n\n{}",
                        activity_log.join("\n"),
                        accumulated_content
                    );
                    return ToolOutput::success(
                        summary,
                        serde_json::json!({
                            "steps": step,
                            "content_length": content_len,
                            "tool_calls": activity_log.iter().filter(|e| e.starts_with("🔧")).count(),
                        }),
                    );
                }
            };

            // 构建 assistant 消息
            let assistant_msg = ChatMessage::assistant_with_reasoning(
                if accumulated_content.is_empty() {
                    None
                } else {
                    Some(accumulated_content.clone())
                },
                None,
                Some(tool_calls.clone()),
            );
            // 持久化 assistant 消息
            let tc_json = serde_json::to_string(&tool_calls).unwrap_or_default();
            self.save_subagent_message(
                subagent_name,
                0,
                step as i32,
                "assistant",
                &accumulated_content,
                None,
                Some(&tc_json),
                None,
            );
            messages.push(assistant_msg);

            // 执行工具调用
            for tool_call in &tool_calls {
                let tool_name = &tool_call.function.name;
                let tool_args_str = &tool_call.function.arguments;

                // 死循环检测
                let call_key = (tool_name.clone(), tool_args_str.clone());
                if last_call.as_ref() == Some(&call_key) {
                    consecutive_count += 1;
                    if consecutive_count >= duplicate_threshold {
                        warn!("[SubAgent] 检测到死循环：{}", tool_name);
                        let error_msg = ChatMessage::tool_result(
                            &tool_call.id,
                            format!(
                                "工具 {} 被连续重复调用。请停止并给出当前收集到的答案。",
                                tool_name
                            ),
                        );
                        messages.push(error_msg);
                        continue;
                    }
                } else {
                    last_call = Some(call_key);
                    consecutive_count = 1;
                }

                // 解析参数
                let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
                    Ok(v) => v,
                    Err(e) => {
                        let error_msg =
                            ChatMessage::tool_result(&tool_call.id, format!("参数解析失败: {}", e));
                        messages.push(error_msg);
                        continue;
                    }
                };

                // ── 向父代理发送进度事件 ──
                if let Some(ref tx) = self.progress_tx {
                    let _ = tx.send(AgentStreamEvent::ToolCall {
                        id: tool_call.id.clone(),
                        name: format!("[sub] {}", tool_name),
                        arguments: args.clone(),
                        step,
                    });
                }

                // ── PreToolUse hooks + Permission check ──
                let tool_ctx = ToolContext::silent(self.app_state.clone());

                let final_args = if let Some(ref registry) = self.hook_registry {
                    let pre_ctx = PreToolUseContext {
                        session_id: "subagent".to_string(),
                        tool_name: tool_name.clone(),
                        tool_args: args.clone(),
                        step,
                    };
                    let pre_result = registry.run_pre_tool_use(&pre_ctx).await;

                    // Block check
                    if pre_result.action.is_blocked() {
                        let reason = pre_result
                            .action
                            .block_reason()
                            .unwrap_or("tool blocked by hook");
                        warn!(
                            "[SubAgent] PreToolUse hook 阻止了工具: {} ({})",
                            tool_name, reason
                        );
                        let tool_msg = ChatMessage::tool_result(
                            &tool_call.id,
                            format!("工具 {} 被阻止: {}", tool_name, reason),
                        );
                        self.save_subagent_message(
                            subagent_name,
                            0,
                            step as i32,
                            "tool",
                            &format!("工具 {} 被阻止: {}", tool_name, reason),
                            None,
                            None,
                            Some(&tool_call.id),
                        );
                        messages.push(tool_msg);
                        continue;
                    }

                    pre_result.final_args
                } else {
                    args.clone()
                };

                // Permission check — 完整三态检查（包含内容级匹配）
                let perm_result = self.permission_checker.check(tool_name, Some(&final_args));
                match perm_result {
                    PermissionResult::Denied { reason } => {
                        warn!("[SubAgent] 权限检查拒绝工具 {}: {}", tool_name, reason);
                        let err_msg =
                            format!("工具 {} 在子代理上下文中不可用: {}", tool_name, reason);
                        let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg);
                        self.save_subagent_message(
                            subagent_name,
                            0,
                            step as i32,
                            "tool",
                            &err_msg,
                            None,
                            None,
                            Some(&tool_call.id),
                        );
                        messages.push(tool_msg);
                        continue;
                    }
                    PermissionResult::AskUser { .. } => {
                        // 子代理上下文中无用户可询问，自动拒绝
                        warn!(
                            "[SubAgent] 工具 {} 需要用户确认，子代理上下文中自动拒绝",
                            tool_name
                        );
                        let err_msg =
                            format!("工具 {} 需要用户确认但在子代理上下文中不可用", tool_name);
                        let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg);
                        self.save_subagent_message(
                            subagent_name,
                            0,
                            step as i32,
                            "tool",
                            &err_msg,
                            None,
                            None,
                            Some(&tool_call.id),
                        );
                        messages.push(tool_msg);
                        continue;
                    }
                    PermissionResult::Allowed => {
                        // 继续执行
                    }
                }

                // 执行工具
                let output = match self.tool_registry.get(tool_name) {
                    Some(tool) => {
                        match tokio::time::timeout(
                            std::time::Duration::from_secs(self.config.tool_timeout_secs),
                            tool.execute(final_args.clone(), &tool_ctx),
                        )
                        .await
                        {
                            Ok(output) => output,
                            Err(_) => ToolOutput::error(format!("工具 {} 执行超时", tool_name)),
                        }
                    }
                    None => ToolOutput::error(format!("未知工具: {}", tool_name)),
                };

                // ── PostToolUse hooks ──
                let final_output_content = if let Some(ref registry) = self.hook_registry {
                    let post_ctx = PostToolUseContext {
                        session_id: "subagent".to_string(),
                        agent_name: "subagent".to_string(),
                        tool_name: tool_name.clone(),
                        tool_args: final_args,
                        output_content: output.content.clone(),
                        is_error: output.is_error,
                        step,
                        elapsed_ms: 0,
                    };
                    let post_result = registry.run_post_tool_use(&post_ctx).await;
                    post_result.final_content
                } else {
                    output.content.clone()
                };

                // 向父代理发送工具结果进度
                if let Some(ref tx) = self.progress_tx {
                    let preview: String = final_output_content.chars().take(200).collect();
                    let _ = tx.send(AgentStreamEvent::ToolResult {
                        tool_call_id: tool_call.id.clone(),
                        name: format!("[sub] {}", tool_name),
                        output: preview,
                        is_error: output.is_error,
                        metadata: serde_json::json!({}),
                        step,
                    });
                }

                // 截断输出（使用 post-hook 处理后的内容）
                let truncated = if final_output_content.len() > self.config.max_tool_output_chars {
                    let t: String = final_output_content
                        .chars()
                        .take(self.config.max_tool_output_chars)
                        .collect();
                    format!(
                        "{}...\n[已截断，原始 {} 字符]",
                        t,
                        final_output_content.len()
                    )
                } else {
                    final_output_content.clone()
                };

                let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated);
                self.save_subagent_message(
                    subagent_name,
                    0,
                    step as i32,
                    "tool",
                    &truncated,
                    None,
                    None,
                    Some(&tool_call.id),
                );
                messages.push(tool_msg);
            }
        }

        // 达到最大步数，强制生成最终答案
        info!("[SubAgent] 达到最大步数 ({}), 生成最终答案", max_steps);
        self.force_final_answer(llm, &messages).await
    }

    /// 强制 LLM 生成最终答案（不带工具调用）
    async fn force_final_answer(&self, llm: &LlmClient, messages: &[ChatMessage]) -> ToolOutput {
        let mut final_messages = messages.to_vec();
        final_messages.push(ChatMessage::user(
            "请根据已收集的信息直接给出最终答案，不要再调用工具。",
        ));

        let empty_tools: Vec<ToolDefinition> = Vec::new();
        let mut stream_rx = match llm
            .chat_stream(&final_messages, &empty_tools, self.config.enable_thinking)
            .await
        {
            Ok(rx) => rx,
            Err(e) => {
                return ToolOutput::error(format!("子代理最终答案生成失败: {}", e));
            }
        };

        let mut accumulated = String::new();
        while let Some(event) = stream_rx.recv().await {
            match event {
                StreamEvent::TextDelta(delta) => {
                    accumulated.push_str(&delta);
                }
                StreamEvent::Done => break,
                StreamEvent::Error(e) => {
                    warn!("[SubAgent] 最终答案流式错误: {}", e);
                    break;
                }
                _ => {}
            }
        }

        if accumulated.is_empty() {
            ToolOutput::error("子代理无法生成最终答案")
        } else {
            ToolOutput::success(accumulated, serde_json::json!({ "forced": true }))
        }
    }
}