// src/agent/tools/background.rs — 后台任务工具 (bg_task_run / bg_task_check) // // 参考 Claude Code s08 Background Tasks 设计。 // 慢速操作可在后台异步执行,结果在下一轮 LLM 调用前注入上下文。 use async_trait::async_trait; use serde_json::json; use std::sync::Arc; use tracing::info; use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; use crate::agent::background::{self, BgNotificationQueue}; /// 支持后台执行的工具列表 const BG_SUPPORTED_TOOLS: &[&str] = &["process_paper"]; /// 后台任务启动工具 pub struct BgTaskRunTool { queue: Arc, } impl BgTaskRunTool { pub fn new(queue: Arc) -> Self { BgTaskRunTool { queue } } } #[async_trait] impl AgentTool for BgTaskRunTool { fn name(&self) -> &str { "bg_task_run" } fn display_name(&self) -> &str { "后台任务" } fn description(&self) -> &str { "在后台异步执行慢速工具(process_paper)。\ 返回任务ID后立即让 LLM 继续思考,后台完成的结果会在下一轮对话中自动通知。\ 适用于:文献下载、解析、向量化等耗时操作。使用 bg_task_check 查询任务状态。" } fn parameters(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "tool_name": { "type": "string", "description": "要在后台执行的工具名称", "enum": ["process_paper"] }, "bibcode": { "type": "string", "description": "文献的唯一标识符(ADS bibcode)" }, "tasks": { "type": "array", "items": { "type": "string", "enum": ["download", "parse", "embed"] }, "description": "需执行的任务列表。不指定时默认执行 ['download', 'parse']" } }, "required": ["tool_name", "bibcode"] }) } /// 后台任务启动有副作用(spawn tokio task),中断时应阻塞以完成 fn interrupt_behavior(&self) -> InterruptBehavior { InterruptBehavior::Block } async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { let tool_name = match args.get("tool_name").and_then(|v| v.as_str()) { Some(s) => s.to_string(), None => return ToolOutput::error("缺少必需参数 'tool_name'"), }; let bibcode = match args.get("bibcode").and_then(|v| v.as_str()) { Some(s) => s.to_string(), None => return ToolOutput::error("缺少必需参数 'bibcode'"), }; if !BG_SUPPORTED_TOOLS.contains(&tool_name.as_str()) { return ToolOutput::error(format!( "工具 '{}' 不支持后台执行。支持的工具: {}", tool_name, BG_SUPPORTED_TOOLS.join(", ") )); } // 提取 tasks 参数(用于 process_paper) let tasks: Vec = args .get("tasks") .and_then(|t| t.as_array()) .map(|arr| { arr.iter() .filter_map(|v| v.as_str().map(String::from)) .collect() }) .unwrap_or_else(|| vec!["download".into(), "parse".into()]); info!( "[BgTaskRun] 启动后台任务: tool={}, bibcode={}, tasks={:?}", tool_name, bibcode, tasks ); let handle = background::spawn_background_task( ctx.app_state.clone(), self.queue.clone(), tool_name.clone(), bibcode.clone(), tasks, ) .await; ToolOutput::success( format!( "✅ 后台任务已启动。\n\ 任务ID: {}\n\ 工具: {}\n\ 文献: {}\n\ 状态: 运行中\n\n\ 使用 bg_task_check 查询任务状态。完成后结果会自动通知。", handle.task_id, handle.tool_name, handle.bibcode, ), json!({ "task_id": handle.task_id, "tool_name": handle.tool_name, "bibcode": handle.bibcode, "status": "running" }), ) } fn defer_loading(&self) -> bool { true } } /// 后台任务查询工具 pub struct BgTaskCheckTool { queue: Arc, } impl BgTaskCheckTool { pub fn new(queue: Arc) -> Self { BgTaskCheckTool { queue } } } #[async_trait] impl AgentTool for BgTaskCheckTool { fn name(&self) -> &str { "bg_task_check" } fn display_name(&self) -> &str { "检查后台" } fn is_internal(&self) -> bool { true } fn description(&self) -> &str { "查询后台任务状态。不指定 task_id 时返回所有任务。" } fn parameters(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "task_id": { "type": "string", "description": "可选:要查询的任务ID。不指定则返回所有任务。" } }, "required": [] }) } /// 纯读取内存队列状态,并发安全 fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { true } async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { let task_id = args.get("task_id").and_then(|v| v.as_str()); match task_id { Some(tid) => match self.queue.get_task(tid).await { Some(task) => { let status_icon = match task.status { background::BgTaskStatus::Running => "🔄", background::BgTaskStatus::Completed => "✅", background::BgTaskStatus::Failed => "❌", }; ToolOutput::success( format!( "{} 任务 {}: {} ({})\n文献: {}", status_icon, task.task_id, task.tool_name, match task.status { background::BgTaskStatus::Running => "运行中", background::BgTaskStatus::Completed => "已完成", background::BgTaskStatus::Failed => "失败", }, task.bibcode, ), serde_json::to_value(&task).unwrap_or_default(), ) } None => ToolOutput::error(format!("任务 '{}' 未找到", tid)), }, None => { let tasks = self.queue.get_all_tasks().await; if tasks.is_empty() { return ToolOutput::success("当前没有后台任务。", json!({ "tasks": [] })); } let mut lines = vec!["📊 后台任务状态:\n".to_string()]; for task in &tasks { let icon = match task.status { background::BgTaskStatus::Running => "🔄", background::BgTaskStatus::Completed => "✅", background::BgTaskStatus::Failed => "❌", }; lines.push(format!( "{} [{}] {} — {}", icon, task.task_id, task.tool_name, task.bibcode )); } ToolOutput::success(lines.join("\n"), json!({ "tasks": tasks })) } } } fn defer_loading(&self) -> bool { true } }