AstroResearch/src/agent/tools/mod.rs

// src/agent/tools/mod.rs
//
// 科研智能体工具集定义与实现。
// 每个工具遵循 AgentTool trait，向大模型声明 JSON Schema 参数定义，
// 并在 execute 中调用已有的服务层完成实际业务操作。
//
// 按功能域拆分为子模块：
//   filesystem/ — 文件 I/O（read、grep、glob、bash、write、edit）
//   astro/      — 天文科研（文献搜索/下载/解析、RAG、天体查询、笔记）
//   team.rs     — 团队协作
//   todo.rs     — 任务规划
//   compress.rs — 手动上下文压缩

use async_trait::async_trait;
use serde_json::json;
use std::sync::{Arc, RwLock};

use crate::agent::runtime::file_cache::FileStateCache;
use crate::agent::skills::SkillRegistry;
use crate::api::AppState;
use crate::clients::llm::ToolDefinition;

pub mod ask_user;
pub mod astro;
mod background;
mod compress;
mod filesystem;
pub mod memory;
pub mod persist;
mod skill;
pub mod subagent;
mod team;
mod todo;

pub use ask_user::AskUserTool;
pub use astro::note::SaveNoteTool;
pub use astro::paper::{DownloadPaperTool, GetPaperContentTool, ParsePaperTool};
pub use astro::rag::RagSearchTool;
pub use astro::search::{GetPaperMetadataTool, SearchPapersTool};
pub use astro::target::QueryTargetTool;
pub use background::{BgTaskCheckTool, BgTaskRunTool};
pub use compress::CompressTool;
pub use filesystem::{
    FileEditTool, FileWriteTool, GlobFilesTool, GrepFilesTool, ReadFileTool, RunBashTool,
};
pub use skill::LoadSkillTool;
pub use subagent::SubAgentTool;
pub use team::{CheckTeamInboxTool, SendTeammateMessageTool, SpawnTeammateTool, TeamBroadcastTool};
pub use todo::persist_tasks;
pub use todo::TodoWriteTool;

/// 工具执行上下文，封装全局共享状态
#[derive(Clone)]
pub struct ToolContext {
    pub app_state: Arc<AppState>,
    /// 当前会话 ID（用于工具将数据关联到正确的 session）
    pub session_id: String,
    /// 静默模式：子代理运行时为 true，跳过用户权限提示
    pub silent: bool,
    /// 文件状态缓存（跨工具调用共享，用于 Read 去重）
    pub read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
    /// 当前 SSE 发送通道（工具可通过此通道向用户推送中间事件）
    pub sse_tx: Option<tokio::sync::mpsc::UnboundedSender<crate::agent::runtime::AgentStreamEvent>>,
    /// 是否启用 LLM 思考模式（继承自父代理配置）
    pub enable_thinking: bool,
    /// 附加允许目录（扩展文件沙箱范围）
    pub additional_allowed_dirs: Vec<String>,
}

impl ToolContext {
    /// 创建标准上下文
    pub fn new(app_state: Arc<AppState>) -> Self {
        ToolContext {
            app_state,
            session_id: String::new(),
            silent: false,
            read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())),
            sse_tx: None,
            enable_thinking: false,
            additional_allowed_dirs: Vec::new(),
        }
    }

    /// 创建带共享文件缓存的上下文（用于 AgentRuntime 保持同一 cache 实例）
    pub fn with_file_cache(
        app_state: Arc<AppState>,
        read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
    ) -> Self {
        ToolContext {
            app_state,
            session_id: String::new(),
            silent: false,
            read_file_state,
            sse_tx: None,
            enable_thinking: false,
            additional_allowed_dirs: Vec::new(),
        }
    }

    /// 设置 SSE 通道
    pub fn with_sse_tx(
        mut self,
        tx: tokio::sync::mpsc::UnboundedSender<crate::agent::runtime::AgentStreamEvent>,
    ) -> Self {
        self.sse_tx = Some(tx);
        self
    }

    /// 设置会话 ID
    pub fn with_session_id(mut self, id: String) -> Self {
        self.session_id = id;
        self
    }

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

    /// 设置附加允许目录
    pub fn with_additional_dirs(mut self, dirs: Vec<String>) -> Self {
        self.additional_allowed_dirs = dirs;
        self
    }

    /// 创建静默上下文（子代理使用）
    pub fn silent(app_state: Arc<AppState>) -> Self {
        ToolContext {
            app_state,
            session_id: String::new(),
            silent: true,
            read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())),
            sse_tx: None,
            enable_thinking: false,
            additional_allowed_dirs: Vec::new(),
        }
    }
}

/// 工具执行结果
#[derive(Debug, Clone)]
pub struct ToolOutput {
    /// 给大模型阅读的截断文本
    pub content: String,
    /// 是否为错误
    pub is_error: bool,
    /// 结构化元数据（给前端 Timeline 直接渲染）
    pub metadata: serde_json::Value,
}

impl ToolOutput {
    /// 创建成功结果
    pub fn success(content: impl Into<String>, metadata: serde_json::Value) -> Self {
        ToolOutput {
            content: content.into(),
            is_error: false,
            metadata,
        }
    }

    /// 创建错误结果
    pub fn error(msg: impl Into<String>) -> Self {
        ToolOutput {
            content: msg.into(),
            is_error: true,
            metadata: json!({}),
        }
    }
}

/// 工具被中断时的行为策略
#[derive(Debug, Clone, PartialEq)]
pub enum InterruptBehavior {
    /// 取消执行并返回错误（默认，适用于只读工具）
    Cancel,
    /// 阻塞中断信号直到执行完成（适用于有副作用的写入工具）
    Block,
}

/// 权限规则来源
#[derive(Debug, Clone, PartialEq)]
pub enum PermissionRuleSource {
    /// 环境变量加载（AGENT_PERMISSIONS_*）
    Env,
    /// 会话内动态添加（API / Always Allow）
    Session,
}

impl std::fmt::Display for PermissionRuleSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PermissionRuleSource::Env => write!(f, "env"),
            PermissionRuleSource::Session => write!(f, "session"),
        }
    }
}

/// 权限规则 — 工具自定义的权限限制
#[derive(Debug, Clone)]
pub enum PermissionRule {
    /// 不可覆盖的拒绝
    Deny {
        tool_name: String,
        reason: String,
        source: PermissionRuleSource,
    },
    /// 允许
    Allow {
        tool_name: String,
        source: PermissionRuleSource,
    },
    /// 需要用户确认
    Ask {
        tool_name: String,
        message: String,
        source: PermissionRuleSource,
    },
}

/// 智能体工具 trait
#[async_trait]
pub trait AgentTool: Send + Sync {
    /// 工具名称（与 LLM function calling 的 name 保持一致）
    fn name(&self) -> &str;
    /// 工具描述（告知 LLM 何时应该调用该工具）
    fn description(&self) -> &str;
    /// JSON Schema 格式的参数定义
    fn parameters(&self) -> serde_json::Value;
    /// 执行工具逻辑
    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput;

    // ── P1 优化：新增默认方法 ──

    /// 中断行为策略。默认 Cancel — 可以安全中断。
    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Cancel
    }

    /// 该工具是否支持并发安全执行。
    /// 默认 false（保守策略），只读工具应覆写为 true。
    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        false
    }

    /// 工具自定义权限检查。默认无额外限制。
    fn check_permissions(&self, _args: &serde_json::Value) -> Vec<PermissionRule> {
        Vec::new()
    }

    /// 该工具错误时是否应中止兄弟并行执行。
    /// 默认 false（只读工具不触发）。下载/解析类工具可覆写为 true。
    fn causes_sibling_abort(&self) -> bool {
        false
    }

    /// 带进度流式执行。默认委托给 execute()。
    /// 长时间操作的工具可覆写以发送进度更新。
    async fn execute_with_progress(
        &self,
        args: serde_json::Value,
        ctx: &ToolContext,
        _progress_tx: Option<&tokio::sync::mpsc::UnboundedSender<String>>,
    ) -> ToolOutput {
        self.execute(args, ctx).await
    }
}

/// 工具注册表，管理所有可用工具。
/// 内部使用 HashMap 实现 O(1) 按名查找，同时保留插入顺序供 definitions() 使用。
///
/// `definition_filter` 用于子代理最小权限：设置后 `definitions()` 仅返回白名单工具，
/// 但 `get()` 仍可查找所有工具（以便对未授权调用返回友好错误消息）。
pub struct ToolRegistry {
    tools: std::collections::HashMap<String, Box<dyn AgentTool>>,
    ordered_names: Vec<String>,
    /// 可选的工具白名单（子代理最小权限）
    definition_filter: Option<std::collections::HashSet<String>>,
    /// 缓存的工具 schema（session 生命周期内有效，工具注册变更时失效）
    schema_cache: Option<Vec<ToolDefinition>>,
    /// schema 缓存版本号（递增使缓存失效）
    schema_generation: u64,
}

// ── 工具注册辅助函数（消除重复代码） ──

/// 注册所有基础研究工具（文件、文献、RAG、笔记等 19 个工具）。
fn add_base_tools(registry: &mut ToolRegistry, skill_registry: Arc<RwLock<SkillRegistry>>) {
    let tools: Vec<Box<dyn AgentTool>> = vec![
        Box::new(ReadFileTool),
        Box::new(GrepFilesTool),
        Box::new(GlobFilesTool),
        Box::new(RunBashTool),
        Box::new(FileWriteTool),
        Box::new(FileEditTool),
        Box::new(SearchPapersTool),
        Box::new(GetPaperMetadataTool),
        Box::new(DownloadPaperTool),
        Box::new(ParsePaperTool),
        Box::new(GetPaperContentTool),
        Box::new(RagSearchTool),
        Box::new(QueryTargetTool),
        Box::new(SaveNoteTool),
        Box::new(TodoWriteTool),
        Box::new(CompressTool),
        Box::new(AskUserTool),
        Box::new(LoadSkillTool::new(skill_registry)),
        Box::new(SubAgentTool::new()),
    ];
    for tool in tools {
        registry.ordered_names.push(tool.name().to_string());
        registry.tools.insert(tool.name().to_string(), tool);
    }
}

/// 注册后台任务工具（bg_task_run, bg_task_check）。
fn add_background_tools(
    registry: &mut ToolRegistry,
    queue: Arc<crate::agent::background::BgNotificationQueue>,
) {
    let run_tool = Box::new(BgTaskRunTool::new(queue.clone()));
    let check_tool = Box::new(BgTaskCheckTool::new(queue));
    registry.ordered_names.push(run_tool.name().to_string());
    registry.tools.insert(run_tool.name().to_string(), run_tool);
    registry.ordered_names.push(check_tool.name().to_string());
    registry
        .tools
        .insert(check_tool.name().to_string(), check_tool);
}

/// 注册团队协作工具（spawn_teammate, send_teammate_message, team_broadcast, check_team_inbox）。
fn add_team_tools(
    registry: &mut ToolRegistry,
    team_manager: Arc<tokio::sync::Mutex<Option<crate::agent::team::manager::TeamManager>>>,
) {
    let team_tools: Vec<Box<dyn AgentTool>> = vec![
        Box::new(SpawnTeammateTool::new(team_manager.clone())),
        Box::new(SendTeammateMessageTool::new(team_manager.clone())),
        Box::new(TeamBroadcastTool::new(team_manager.clone())),
        Box::new(CheckTeamInboxTool::new(team_manager)),
    ];
    for tool in team_tools {
        registry.ordered_names.push(tool.name().to_string());
        registry.tools.insert(tool.name().to_string(), tool);
    }
}

impl ToolRegistry {
    /// 创建空工具注册表（调用者通过 add_tool 手动添加工具）。
    /// 用于受限场景（如记忆提取子代理只需要只读 + save_memory）。
    pub fn empty() -> Self {
        ToolRegistry {
            tools: std::collections::HashMap::new(),
            ordered_names: Vec::new(),
            definition_filter: None,
            schema_cache: None,
            schema_generation: 0,
        }
    }

    /// 创建默认工具注册表（包含全部科研工具，不含后台工具）
    pub fn new(skill_registry: Arc<RwLock<SkillRegistry>>) -> Self {
        Self::new_with_queue(None, skill_registry)
    }

    /// 创建工具注册表，可选注入后台通知队列以启用 bg_task_run/bg_task_check
    pub fn new_with_queue(
        queue: Option<Arc<crate::agent::background::BgNotificationQueue>>,
        skill_registry: Arc<RwLock<SkillRegistry>>,
    ) -> Self {
        let mut registry = ToolRegistry {
            tools: std::collections::HashMap::new(),
            ordered_names: Vec::new(),
            definition_filter: None,
            schema_cache: None,
            schema_generation: 0,
        };
        add_base_tools(&mut registry, skill_registry);
        if let Some(q) = queue {
            add_background_tools(&mut registry, q);
        }
        registry.precompute_definitions();
        registry
    }

    /// 创建包含团队工具的注册表
    pub fn new_with_team(
        queue: Option<Arc<crate::agent::background::BgNotificationQueue>>,
        team_manager: Arc<tokio::sync::Mutex<Option<crate::agent::team::manager::TeamManager>>>,
        skill_registry: Arc<RwLock<SkillRegistry>>,
    ) -> Self {
        // 使用 new_with_queue 获取基础 + 后台工具，再添加团队工具
        let mut registry = Self::new_with_queue(queue, skill_registry);
        add_team_tools(&mut registry, team_manager);
        registry.precompute_definitions();
        registry
    }

    /// 动态添加工具（用于需要共享状态的工具，如 MemoryManager）
    pub fn add_tool(&mut self, tool: Box<dyn AgentTool>) {
        let name = tool.name().to_string();
        self.ordered_names.push(name.clone());
        self.tools.insert(name, tool);
        // 工具变更 → 使 schema 缓存失效
        self.schema_cache = None;
        self.schema_generation += 1;
    }

    /// 替换已存在的工具（保持名称在 ordered_names 中的位置不变）。
    /// 如果工具不存在，行为等同于 add_tool。
    pub fn replace_tool(&mut self, tool: Box<dyn AgentTool>) {
        let name = tool.name().to_string();
        if !self.tools.contains_key(&name) {
            self.ordered_names.push(name.clone());
        }
        self.tools.insert(name, tool);
        // 工具变更 → 使 schema 缓存失效
        self.schema_cache = None;
        self.schema_generation += 1;
    }

    /// 根据名称查找工具 (O(1))
    pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
        self.tools.get(name).map(|t| t.as_ref())
    }

    /// 设置工具定义白名单过滤器（子代理最小权限）。
    /// 设置后 `definitions()` 仅返回白名单中的工具。
    pub fn set_definition_filter(&mut self, allowed: Vec<String>) {
        self.definition_filter = Some(allowed.into_iter().collect());
        // 过滤器变更 → 使 schema 缓存失效
        self.schema_cache = None;
        self.schema_generation += 1;
    }

    /// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM）。
    /// 按名称字母序排序以保证跨调用的稳定性，提升 prompt cache 命中率。
    /// 若设置了 definition_filter，仅返回白名单中的工具。
    ///
    /// 使用内部 schema 缓存：工具注册表不变时复用上次计算结果。
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        // 缓存命中：直接返回
        if let Some(ref cached) = self.schema_cache {
            return cached.clone();
        }

        // 缓存未命中：重新计算
        let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
            self.tools
                .iter()
                .filter(|(name, _)| filter.contains(*name))
                .map(|(_, tool)| tool)
                .collect()
        } else {
            self.tools.values().collect()
        };
        let mut defs: Vec<_> = values
            .iter()
            .map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
            .collect();
        defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));

        // 此方法在 &self 上调用，无法通过 &mut self 修改 self.schema_cache。
        // 使用 unsafe 实现内部可变性，仅在缓存空时写入（逻辑上安全）。
        // 如果 schema_cache 为 None，说明工具集未变更，写入是安全的。
        // Note: 实际上这里需要用内部可变性。最简单的方式是用 RefCell 或
        // 将 schema_cache 和 schema_generation 移到 Cell/RefCell 中。
        // 但为了最小改动，这里依赖一个事实：definitions() 在主线程中调用时
        // 没有并发写入问题，且 add_tool/replace_tool 只在初始化时调用。
        defs
    }

    /// 获取当前 schema 缓存版本号（用于外部检测工具是否变更）。
    pub fn schema_generation(&self) -> u64 {
        self.schema_generation
    }

    /// 预计算并缓存 tool definitions（在 AgentRuntime 初始化时调用，
    /// 此时拥有 &mut self，可以安全写入缓存）。
    pub fn precompute_definitions(&mut self) {
        if self.schema_cache.is_some() {
            return;
        }
        let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
            self.tools
                .iter()
                .filter(|(name, _)| filter.contains(*name))
                .map(|(_, tool)| tool)
                .collect()
        } else {
            self.tools.values().collect()
        };
        let mut defs: Vec<_> = values
            .iter()
            .map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
            .collect();
        defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
        self.schema_cache = Some(defs);
    }

    /// 设置工具白名单（子代理最小权限）。
    /// 设置后 `definitions()` 仅暴露白名单中的工具给 LLM，
    /// 但 `get()` 仍可访问所有工具以便对未授权调用返回友好错误消息。
    pub fn with_filter(mut self, allowed: &[&str]) -> Self {
        self.definition_filter = Some(allowed.iter().map(|s| s.to_string()).collect());
        self
    }

    /// 返回当前所有工具名称列表
    pub fn tool_names(&self) -> Vec<String> {
        self.ordered_names.clone()
    }
}

/// 截断文本到指定最大字符数
pub fn truncate_content(s: &str, max_chars: usize) -> String {
    if s.len() <= max_chars {
        s.to_string()
    } else {
        let truncated: String = s.chars().take(max_chars).collect();
        format!("{}\n\n[... 内容已截断，共 {} 字符 ...]", truncated, s.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_truncate_content_short() {
        let text = "Hello, world!";
        assert_eq!(truncate_content(text, 100), text);
    }

    #[test]
    fn test_truncate_content_long() {
        let text = "a".repeat(5000);
        let result = truncate_content(&text, 100);
        assert!(result.contains("内容已截断"));
        assert!(result.contains("5000"));
    }

    #[test]
    fn test_tool_output_success() {
        let output = ToolOutput::success("ok", json!({"key": "value"}));
        assert!(!output.is_error);
        assert_eq!(output.content, "ok");
    }

    #[test]
    fn test_tool_output_error() {
        let output = ToolOutput::error("something went wrong");
        assert!(output.is_error);
    }

    #[test]
    fn test_tool_registry_definitions() {
        let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
            "./skills",
        )))));
        let defs = registry.definitions();
        assert_eq!(defs.len(), 19);
        assert!(defs.iter().any(|d| d.function.name == "read_file"));
        assert!(defs.iter().any(|d| d.function.name == "grep_files"));
        assert!(defs.iter().any(|d| d.function.name == "glob_files"));
        assert!(defs.iter().any(|d| d.function.name == "run_bash"));
        assert!(defs.iter().any(|d| d.function.name == "file_write"));
        assert!(defs.iter().any(|d| d.function.name == "file_edit"));
        assert!(defs.iter().any(|d| d.function.name == "search_papers"));
        assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata"));
        assert!(defs.iter().any(|d| d.function.name == "download_paper"));
        assert!(defs.iter().any(|d| d.function.name == "parse_paper"));
        assert!(defs.iter().any(|d| d.function.name == "get_paper_content"));
        assert!(defs.iter().any(|d| d.function.name == "rag_search"));
        assert!(defs.iter().any(|d| d.function.name == "query_target"));
        assert!(defs.iter().any(|d| d.function.name == "save_note"));
        assert!(defs.iter().any(|d| d.function.name == "todo_write"));
        assert!(defs.iter().any(|d| d.function.name == "compress_context"));
        assert!(defs.iter().any(|d| d.function.name == "load_skill"));
        assert!(defs.iter().any(|d| d.function.name == "subagent"));
    }

    #[test]
    fn test_tool_registry_get() {
        let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
            "./skills",
        )))));
        assert!(registry.get("search_papers").is_some());
        assert!(registry.get("nonexistent").is_none());
    }
}