AstroResearch/src/agent/team/teammate.rs

// src/agent/team/teammate.rs
//
// 队友 ReAct 循环：检查收件箱 → 执行任务 → 汇报结果。

use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{info, warn};

use crate::agent::background::BgNotificationQueue;
use crate::agent::compact;
use crate::agent::runtime::AgentConfig;
use crate::agent::tools::ToolRegistry;
use crate::api::AppState;
use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent};

use super::config::MemberStatus;
use super::inbox::{self, TeamMessageType};

/// 队友 ReAct 循环。
///
/// 生命周期：
///   SPAWN → WORKING (ReAct) → IDLE (poll inbox) → SHUTDOWN
#[allow(clippy::too_many_arguments)]
pub async fn run_teammate_loop(
    app_state: Arc<AppState>,
    team_dir: PathBuf,
    name: String,
    role: String,
    system_prompt: String,
    task_prompt: String,
    config: AgentConfig,
    status: Arc<Mutex<MemberStatus>>,
    cancelled: Arc<AtomicBool>,
) {
    let llm = &app_state.llm;
    // 队友的工具注册表排除 delegate_research（防止无限委托链）
    let queue = Arc::new(BgNotificationQueue::new());
    let tool_registry =
        ToolRegistry::new_with_queue(Some(queue.clone()), app_state.skill_registry.clone());

    let tool_defs = tool_registry.definitions();

    let mut messages = vec![
        ChatMessage::system(&system_prompt),
        ChatMessage::user(&task_prompt),
    ];

    info!("[Teammate:{}] 启动 ReAct 循环", name);

    loop {
        // ── 检查取消 ──
        if cancelled.load(Ordering::SeqCst) {
            info!("[Teammate:{}] 收到取消信号，正在停止...", name);
            *status.lock().await = MemberStatus::Shutdown;
            // 通知 lead 自己退出了
            let goodbye = super::inbox::TeamMessage::new(
                &name,
                "lead",
                &format!("队友 {} ({}) 已退出。", name, role),
                TeamMessageType::Status,
            );
            let _ = inbox::append_message(&team_dir, "lead", &goodbye);
            return;
        }

        // ── IDLE 阶段：检查收件箱 ──
        *status.lock().await = MemberStatus::Idle;

        let inbox_msgs = inbox::drain_inbox(&team_dir, &name);

        if inbox_msgs.is_empty() {
            // 等待新消息或取消，poll 间隔 5 秒，最长 60 秒
            for _ in 0..12 {
                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                if cancelled.load(Ordering::SeqCst) || inbox::has_pending(&team_dir, &name) {
                    break;
                }
            }

            if cancelled.load(Ordering::SeqCst) {
                continue;
            }

            // 再次 drain（可能因为 pending 标志被唤醒）
            let inbox_msgs = inbox::drain_inbox(&team_dir, &name);
            if inbox_msgs.is_empty() {
                continue; // 超时，没有新消息，继续 idle
            }

            // 有新消息 → 进入 WORKING 阶段
            *status.lock().await = MemberStatus::Working;

            for msg in &inbox_msgs {
                messages.push(ChatMessage::user(format!(
                    "[来自 {} 的消息 ({}):] {}",
                    msg.from, msg.timestamp, msg.content
                )));
            }

            // ── 执行 ReAct 循环 ──
            let result = run_teammate_react_turn(
                llm,
                &tool_defs,
                &tool_registry,
                &app_state,
                &mut messages,
                &config,
                &cancelled,
            )
            .await;

            // ── 汇报结果 ──
            if let Some(ref summary) = result {
                let reply =
                    super::inbox::TeamMessage::new(&name, "lead", summary, TeamMessageType::Result);
                let _ = inbox::append_message(&team_dir, "lead", &reply);
                info!(
                    "[Teammate:{}] 任务完成，已发送结果 ({}字符)",
                    name,
                    summary.len()
                );
            }
        } else {
            // 收件箱有消息 → 直接进入 WORKING
            *status.lock().await = MemberStatus::Working;

            for msg in &inbox_msgs {
                messages.push(ChatMessage::user(format!(
                    "[来自 {} 的消息: ({})] {}",
                    msg.from, msg.timestamp, msg.content
                )));
            }

            let result = run_teammate_react_turn(
                llm,
                &tool_defs,
                &tool_registry,
                &app_state,
                &mut messages,
                &config,
                &cancelled,
            )
            .await;

            if let Some(ref summary) = result {
                let reply =
                    super::inbox::TeamMessage::new(&name, "lead", summary, TeamMessageType::Result);
                let _ = inbox::append_message(&team_dir, "lead", &reply);
            }
        }
    }
}

/// 队友的单次 ReAct turn。
///
/// 一个简化的 ReAct 循环：LLM 调用 → 工具执行 → 结果注入 → 循环...
/// 与主 Agent 的循环类似但更轻量（无 SSE、无 DB 持久化、无 hooks）。
async fn run_teammate_react_turn(
    llm: &LlmClient,
    tool_defs: &[crate::clients::llm::ToolDefinition],
    tool_registry: &ToolRegistry,
    app_state: &Arc<AppState>,
    messages: &mut Vec<ChatMessage>,
    config: &AgentConfig,
    cancelled: &Arc<AtomicBool>,
) -> Option<String> {
    let max_steps = config.max_steps.min(5); // 队友步数限制更严格

    for _step in 1..=max_steps {
        // 检查取消
        if cancelled.load(Ordering::SeqCst) {
            return None;
        }

        // 上下文压缩检查
        let est_tokens: usize = messages
            .iter()
            .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4)
            .sum();
        if est_tokens > config.context_char_limit * 3 / 2 {
            compact::compress_context(messages, llm, config.context_char_limit, "teammate").await;
        }

        // LLM 流式调用
        let mut stream_rx = match llm.chat_stream(messages, tool_defs).await {
            Ok(rx) => rx,
            Err(_) => return None,
        };

        let mut accumulated = String::new();
        let mut tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;

        while let Some(event) = stream_rx.recv().await {
            match event {
                StreamEvent::TextDelta(delta) => accumulated.push_str(&delta),
                StreamEvent::ToolCallsComplete(tc) => tool_calls = Some(tc),
                StreamEvent::Done => break,
                StreamEvent::Error(_) => return None,
                _ => {}
            }
        }

        // 无工具调用 = 最终回答
        let tool_calls = match tool_calls {
            Some(ref tc) if !tc.is_empty() => tc.clone(),
            _ => {
                return if accumulated.is_empty() {
                    None
                } else {
                    Some(accumulated)
                };
            }
        };

        // 构建 assistant 消息
        messages.push(ChatMessage::assistant_with_reasoning(
            if accumulated.is_empty() {
                None
            } else {
                Some(accumulated)
            },
            None,
            Some(tool_calls.clone()),
        ));

        // 执行工具调用
        for tc in &tool_calls {
            let args: serde_json::Value = match serde_json::from_str(&tc.function.arguments) {
                Ok(a) => a,
                Err(_) => continue,
            };

            let tool_ctx = crate::agent::tools::ToolContext::silent(app_state.clone());

            let output = match tool_registry.get(&tc.function.name) {
                Some(tool) => {
                    match tokio::time::timeout(
                        std::time::Duration::from_secs(config.tool_timeout_secs),
                        tool.execute(args, &tool_ctx),
                    )
                    .await
                    {
                        Ok(o) => o,
                        Err(_) => crate::agent::tools::ToolOutput::error("执行超时"),
                    }
                }
                None => crate::agent::tools::ToolOutput::error("未知工具"),
            };

            let truncated = if output.content.len() > config.max_tool_output_chars {
                let t: String = output
                    .content
                    .chars()
                    .take(config.max_tool_output_chars)
                    .collect();
                format!("{}...\n[已截断]", t)
            } else {
                output.content
            };

            messages.push(ChatMessage::tool_result(&tc.id, &truncated));
        }
    }

    // 达到最大步数，返回 None（无结果）
    warn!("[Teammate] 达到最大步数限制 ({} steps)，无结果", max_steps);
    None
}