feat: Agent 全栈升级——模块化重构、Hooks/Skills/Memory/SubAgent/Team 子系统、审计与任务持久化
架构重构: - Agent Runtime 由单文件拆为 runtime/ 目录 12 模块(熔断/流式执行/Token预算/文件缓存/权限等) - Agent Tools 由单文件拆为 tools/ 目录 20+ 模块(filesystem/astro/memory/skill/subagent/team 等) - 解析器体系重构(common.rs 836行变更),各解析器同步升级 - Download 服务重构(562行),反爬策略强化 - LLM 客户端重构(266行),流式调用优化 新子系统: - Hooks 生命周期系统(9种事件类型,PreToolUse/PostToolUse 支持输入输出拦截) - Skills 双层加载系统(system-reminder 轻量注入 + LoadSkillTool 按需加载,notify 文件监听热更新) - Memory 项目记忆管理(类型/提取/去重/衰减/保活/选择策略/护栏 7 模块) - SubAgent 上下文隔离子代理运行器(独立 ReAct 循环 + Hook 管道) - Team 多智能体团队协作(文件 inbox 通信、lead/teammate 协调) - TaskBoard DAG 任务依赖管理 - Trajectory 会话轨迹、Terminal 终止信号、Autonomous 自主模式、Background 异步通知 数据库: - agent_tasks 表(DAG 依赖模式,blocked_by JSON 数组) - agent_audit_log 表(工具调用审计:名称/状态/耗时/输出预览) - agent_identity 迁移(消息/审计/任务的 agent_name 归属,agent_team_members 团队注册表) API: - GET /chat/metrics 聚合指标端点 - GET /chat/sessions/:id/audit 会话审计查询 - GET /chat/questions + POST /chat/answer 人机交互问答 工程: - 新增依赖:serde_yaml、notify、glob、walkdir、lru - Skills 目录含 methodology/plotting/presentation 三个初始 SKILL.md - CLAUDE.md 完整项目架构文档
This commit is contained in:
@@ -0,0 +1,137 @@
|
||||
// src/agent/autonomous.rs
|
||||
//
|
||||
// 自治研究循环 — 空闲时自动轮询新任务。
|
||||
// 参考 learn-claude-code s17 Autonomous Agents。
|
||||
//
|
||||
// 当 Agent 完成当前回合后,进入 IDLE 阶段:
|
||||
// 1. 检查是否有未认领的团队任务
|
||||
// 2. 检查批量同步任务队列
|
||||
// 3. 检查订阅分类的新论文
|
||||
// 发现工作后自动认领并执行。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Notify;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 自治研究配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AutoResearchConfig {
|
||||
/// 是否启用自治模式
|
||||
pub enabled: bool,
|
||||
/// 订阅的 arXiv 分类
|
||||
pub subscribed_categories: Vec<String>,
|
||||
/// 最大自治轮次(防止无限循环)
|
||||
pub max_autonomous_turns: usize,
|
||||
/// IDLE 超时(分钟)
|
||||
pub idle_timeout_minutes: u64,
|
||||
}
|
||||
|
||||
impl Default for AutoResearchConfig {
|
||||
fn default() -> Self {
|
||||
AutoResearchConfig {
|
||||
enabled: false,
|
||||
subscribed_categories: vec!["astro-ph".to_string()],
|
||||
max_autonomous_turns: 5,
|
||||
idle_timeout_minutes: 60,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// IDLE 轮询器 — 在 Agent 空闲时检查是否有待处理工作。
|
||||
pub struct IdlePoller {
|
||||
app_state: Arc<AppState>,
|
||||
config: AutoResearchConfig,
|
||||
poll_interval: Duration,
|
||||
/// 唤醒通知(当外部事件触发时,如新论文到达、任务分配)
|
||||
wake_notify: Arc<Notify>,
|
||||
}
|
||||
|
||||
impl IdlePoller {
|
||||
pub fn new(app_state: Arc<AppState>, config: AutoResearchConfig) -> Self {
|
||||
IdlePoller {
|
||||
app_state,
|
||||
config,
|
||||
poll_interval: Duration::from_secs(30),
|
||||
wake_notify: Arc::new(Notify::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取唤醒通知器的 clone(供外部触发)
|
||||
pub fn wake_sender(&self) -> Arc<Notify> {
|
||||
self.wake_notify.clone()
|
||||
}
|
||||
|
||||
/// 启动 IDLE 循环(应在独立的 tokio::spawn 中运行)
|
||||
pub async fn start(self) {
|
||||
info!(
|
||||
"[IdlePoller] 启动自治轮询 (间隔={:?}, 最大轮次={})",
|
||||
self.poll_interval, self.config.max_autonomous_turns
|
||||
);
|
||||
|
||||
let mut autonomous_turns: usize = 0;
|
||||
|
||||
loop {
|
||||
// 等待 poll_interval 或被 notify 唤醒
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(self.poll_interval) => {}
|
||||
_ = self.wake_notify.notified() => {
|
||||
info!("[IdlePoller] 被外部事件唤醒");
|
||||
}
|
||||
}
|
||||
|
||||
if !self.config.enabled {
|
||||
continue;
|
||||
}
|
||||
|
||||
if autonomous_turns >= self.config.max_autonomous_turns {
|
||||
info!(
|
||||
"[IdlePoller] 已达到最大自治轮次 ({}),停止轮询",
|
||||
self.config.max_autonomous_turns
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
// 1. 检查未认领的团队任务
|
||||
let task_board = crate::agent::task_board::TaskBoard::new(self.app_state.db.clone());
|
||||
match task_board.list_available_tasks(5).await {
|
||||
Ok(tasks) if !tasks.is_empty() => {
|
||||
for task in tasks {
|
||||
if task.can_start {
|
||||
info!(
|
||||
"[IdlePoller] 发现可认领任务: {} (session={})",
|
||||
task.task_id, task.session_id
|
||||
);
|
||||
if let Ok(true) = task_board
|
||||
.claim_task(&task.session_id, &task.task_id, "auto")
|
||||
.await
|
||||
{
|
||||
autonomous_turns += 1;
|
||||
// 创建 AgentRuntime 并执行任务
|
||||
let runtime = crate::agent::runtime::AgentRuntime::new(
|
||||
self.app_state.clone(),
|
||||
);
|
||||
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let _ = runtime
|
||||
.run_turn(Some(task.session_id.clone()), &task.content, tx)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
warn!("[IdlePoller] 任务查询失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 检查批量同步状态 (placeholder)
|
||||
// 未来: 检查订阅分类的新论文并自动触发同步
|
||||
}
|
||||
|
||||
info!("[IdlePoller] IDLE 循环结束");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
// src/agent/background.rs
|
||||
//
|
||||
// 后台任务执行子系统(参考 Claude Code s08 Background Tasks)。
|
||||
//
|
||||
// 慢速操作(download_paper, parse_paper, embed_paper)可通过
|
||||
// bg_task_run 在后台异步执行,LLM 继续思考/调用其他工具。
|
||||
//
|
||||
// 完成的通知通过 BgNotificationQueue 在下一轮 LLM 调用前注入。
|
||||
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 后台任务结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BgTaskResult {
|
||||
pub task_id: String,
|
||||
pub tool_name: String,
|
||||
pub bibcode: String,
|
||||
pub is_error: bool,
|
||||
pub summary: String,
|
||||
}
|
||||
|
||||
/// 后台任务状态跟踪
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BgTaskHandle {
|
||||
pub task_id: String,
|
||||
pub tool_name: String,
|
||||
pub bibcode: String,
|
||||
pub status: BgTaskStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BgTaskStatus {
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
}
|
||||
|
||||
/// 后台任务通知队列。
|
||||
///
|
||||
/// 使用 mpsc channel 在后台完成和主循环之间传递结果。
|
||||
/// drain() 在每轮 LLM 调用前被调用,收集所有已完成的后台任务结果。
|
||||
pub struct BgNotificationQueue {
|
||||
rx: Mutex<mpsc::UnboundedReceiver<BgTaskResult>>,
|
||||
tx: mpsc::UnboundedSender<BgTaskResult>,
|
||||
/// 内存中的任务状态注册表
|
||||
tasks: Mutex<HashMap<String, BgTaskHandle>>,
|
||||
}
|
||||
|
||||
impl Default for BgNotificationQueue {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl BgNotificationQueue {
|
||||
/// 创建新的通知队列
|
||||
pub fn new() -> Self {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
BgNotificationQueue {
|
||||
rx: Mutex::new(rx),
|
||||
tx,
|
||||
tasks: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取发送端(供后台任务使用)
|
||||
pub fn sender(&self) -> mpsc::UnboundedSender<BgTaskResult> {
|
||||
self.tx.clone()
|
||||
}
|
||||
|
||||
/// 注册一个开始执行的后台任务
|
||||
pub async fn register_task(&self, task: BgTaskHandle) {
|
||||
let mut tasks = self.tasks.lock().await;
|
||||
tasks.insert(task.task_id.clone(), task);
|
||||
}
|
||||
|
||||
/// 更新任务状态
|
||||
pub async fn update_task_status(&self, task_id: &str, status: BgTaskStatus) {
|
||||
let mut tasks = self.tasks.lock().await;
|
||||
if let Some(task) = tasks.get_mut(task_id) {
|
||||
task.status = status;
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取所有任务状态
|
||||
pub async fn get_all_tasks(&self) -> Vec<BgTaskHandle> {
|
||||
let tasks = self.tasks.lock().await;
|
||||
tasks.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// 获取单个任务状态
|
||||
pub async fn get_task(&self, task_id: &str) -> Option<BgTaskHandle> {
|
||||
let tasks = self.tasks.lock().await;
|
||||
tasks.get(task_id).cloned()
|
||||
}
|
||||
|
||||
/// 排空所有已完成的后台任务通知。
|
||||
/// 在每轮 LLM 调用前调用。
|
||||
pub async fn drain(&self) -> Vec<BgTaskResult> {
|
||||
let mut results = Vec::new();
|
||||
let mut rx = self.rx.lock().await;
|
||||
while let Ok(result) = rx.try_recv() {
|
||||
results.push(result);
|
||||
}
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
/// 在后台执行指定的工具调用。
|
||||
///
|
||||
/// 启动一个 tokio::spawn 异步任务执行工具,
|
||||
/// 完成后通过通知队列发送结果。
|
||||
pub async fn spawn_background_task(
|
||||
app_state: Arc<AppState>,
|
||||
queue: Arc<BgNotificationQueue>,
|
||||
tool_name: String,
|
||||
bibcode: String,
|
||||
) -> BgTaskHandle {
|
||||
let task_id = uuid::Uuid::new_v4().to_string();
|
||||
// 取前 8 位便于显示
|
||||
let short_id = task_id[..8].to_string();
|
||||
|
||||
let handle = BgTaskHandle {
|
||||
task_id: short_id.clone(),
|
||||
tool_name: tool_name.clone(),
|
||||
bibcode: bibcode.clone(),
|
||||
status: BgTaskStatus::Running,
|
||||
};
|
||||
|
||||
queue.register_task(handle.clone()).await;
|
||||
|
||||
let queue_clone = queue.clone();
|
||||
let app_state_clone = app_state.clone();
|
||||
let tool_name_clone = tool_name.clone();
|
||||
let bibcode_clone = bibcode.clone();
|
||||
let short_id_clone = short_id.clone();
|
||||
let sender = queue.sender();
|
||||
|
||||
tokio::spawn(async move {
|
||||
info!(
|
||||
"[Background] 启动后台任务 {}: {} ({})",
|
||||
short_id_clone, tool_name_clone, bibcode_clone
|
||||
);
|
||||
|
||||
// 构造 ToolContext 和参数 (后台任务:静默模式)
|
||||
let tool_ctx = ToolContext::silent(app_state_clone.clone());
|
||||
let args = serde_json::json!({"bibcode": bibcode_clone});
|
||||
let tool_registry = ToolRegistry::new(app_state_clone.skill_registry.clone());
|
||||
|
||||
let output = match tool_registry.get(&tool_name_clone) {
|
||||
Some(tool) => {
|
||||
match tokio::time::timeout(
|
||||
std::time::Duration::from_secs(300), // 5 min timeout for bg tasks
|
||||
tool.execute(args, &tool_ctx),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(o) => o,
|
||||
Err(_) => ToolOutput::error("后台任务执行超时(300秒)"),
|
||||
}
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name_clone)),
|
||||
};
|
||||
|
||||
// 发送完成通知
|
||||
let result = BgTaskResult {
|
||||
task_id: short_id_clone.clone(),
|
||||
tool_name: tool_name_clone,
|
||||
bibcode: bibcode_clone,
|
||||
is_error: output.is_error,
|
||||
summary: if output.content.len() > 500 {
|
||||
let preview: String = output.content.chars().take(500).collect();
|
||||
format!("{}...", preview)
|
||||
} else {
|
||||
output.content.clone()
|
||||
},
|
||||
};
|
||||
|
||||
let _ = sender.send(result);
|
||||
queue_clone
|
||||
.update_task_status(
|
||||
&short_id_clone,
|
||||
if output.is_error {
|
||||
BgTaskStatus::Failed
|
||||
} else {
|
||||
BgTaskStatus::Completed
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
info!("[Background] 后台任务 {} 完成", short_id_clone);
|
||||
});
|
||||
|
||||
handle
|
||||
}
|
||||
@@ -0,0 +1,721 @@
|
||||
// src/agent/compact.rs
|
||||
//
|
||||
// 上下文压缩子系统。
|
||||
// 实现四层压缩策略(参考 Claude Code compaction pipeline):
|
||||
// 0. snip_compact — 零API调用:消息数超阈值时截断中间段
|
||||
// 1. micro_compact — 轻量级:替换较早的工具结果为占位符
|
||||
// 2. auto_compact — 自动触发:超 token 阈值时 LLM 摘要对话历史
|
||||
// 3. manual_compact — 手动触发:Agent 通过 compress_context 工具主动调用
|
||||
//
|
||||
// 参考 Claude Code src/services/compact/ 模块设计。
|
||||
//
|
||||
// P0 改进:
|
||||
// - Transcript 持久化(压缩前保存完整 JSONL)
|
||||
// - 多层回退链(snip → micro → auto → aggressive_micro → identity inject)
|
||||
// - Identity re-injection(压缩后消息过少时注入身份确认)
|
||||
// - 占位符优化(使用工具名称替代字符预览)
|
||||
|
||||
pub mod collapse;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::hooks::{HookRegistry, PostCompactContext, PreCompactContext};
|
||||
|
||||
/// 递归守卫:防止压缩内部触发的 LLM 调用再次触发压缩。
|
||||
static COMPACTING: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
use crate::clients::llm::{ChatMessage, LlmClient, MessageRole};
|
||||
|
||||
/// 获取 transcripts 存储目录
|
||||
fn transcripts_dir() -> PathBuf {
|
||||
PathBuf::from(".transcripts")
|
||||
}
|
||||
|
||||
/// 找到安全的上下文切割点,确保不会切断 tool_call / tool_result 配对。
|
||||
/// 从末尾向前扫描,如果候选切割点的第一条要保留的消息是 tool 角色,
|
||||
/// 则向前追溯到对应的 assistant(tool_calls) 消息一并保留。
|
||||
pub fn find_safe_cut_point(messages: &[ChatMessage], desired_keep: usize) -> usize {
|
||||
if messages.len() <= desired_keep {
|
||||
return 0; // 全部保留,无需切割
|
||||
}
|
||||
|
||||
let mut cut = messages.len().saturating_sub(desired_keep);
|
||||
|
||||
// 确保不从 system 消息之后的第一条就开始切(至少保留 system)
|
||||
if cut == 0 {
|
||||
cut = 1;
|
||||
}
|
||||
|
||||
// 如果切割点落在一个 tool 消息上,向前扩展以包含其对应的 assistant(tool_calls)
|
||||
loop {
|
||||
if cut >= messages.len() {
|
||||
cut = messages.len() - 1;
|
||||
break;
|
||||
}
|
||||
|
||||
let first_kept = &messages[cut];
|
||||
|
||||
if first_kept.role == MessageRole::Tool {
|
||||
let mut found_pair = false;
|
||||
for j in (0..cut).rev() {
|
||||
if messages[j].role == MessageRole::Assistant && messages[j].tool_calls.is_some() {
|
||||
cut = j;
|
||||
found_pair = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if !found_pair {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 检查切割点之前没有孤立的 assistant(tool_calls)。
|
||||
// 从后往前扫描:每找到一个孤立的 assistant,将 cut 移到该位置并继续向前检查。
|
||||
let mut search = cut;
|
||||
while search > 0 {
|
||||
let mut found = false;
|
||||
for j in (0..search).rev() {
|
||||
if messages[j].role == MessageRole::Assistant && messages[j].tool_calls.is_some() {
|
||||
let has_tool_result = messages[j + 1..search]
|
||||
.iter()
|
||||
.any(|m| m.role == MessageRole::Tool);
|
||||
if !has_tool_result {
|
||||
cut = j;
|
||||
search = j;
|
||||
found = true;
|
||||
}
|
||||
break; // 只处理最近的一个,继续向前
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
cut
|
||||
}
|
||||
|
||||
// ── snip_compact (Layer 0) ──────────────────────────────────────────────────
|
||||
|
||||
/// 最大消息数(超过此阈值触发 snip_compact)
|
||||
pub const MAX_MESSAGES: usize = 50;
|
||||
/// 保留的头部消息数(system prompt + 初始上下文)
|
||||
pub const HEAD_KEEP: usize = 3;
|
||||
|
||||
/// Layer 0 压缩:当消息数超过 `MAX_MESSAGES` 时,保留前 `HEAD_KEEP` 条
|
||||
/// + 后 `MAX_MESSAGES - HEAD_KEEP` 条,中间替换为占位消息。
|
||||
///
|
||||
/// 这是零 API 调用的最廉价压缩层。使用 `find_safe_cut_point` 确保
|
||||
/// 切割点不会破坏 `assistant(tool_calls)` / `tool_result` 配对。
|
||||
///
|
||||
/// 返回 `true` 表示执行了压缩。
|
||||
pub fn snip_compact(messages: &mut Vec<ChatMessage>, max_messages: usize) -> bool {
|
||||
if messages.len() <= max_messages {
|
||||
return false;
|
||||
}
|
||||
|
||||
let tail_keep = max_messages - HEAD_KEEP;
|
||||
let original_len = messages.len();
|
||||
|
||||
// 找到安全的尾部起点(复用已有配对保护逻辑)
|
||||
let tail_start = find_safe_cut_point(messages, tail_keep);
|
||||
|
||||
// 确保不跟头部重叠
|
||||
if tail_start <= HEAD_KEEP {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 收集被移除段中使用的工具名称
|
||||
let snipped = &messages[HEAD_KEEP..tail_start];
|
||||
let mut tool_names: Vec<String> = Vec::new();
|
||||
for msg in snipped {
|
||||
if msg.role == MessageRole::Tool {
|
||||
if let Some(call_id) = &msg.tool_call_id {
|
||||
if let Some(name) = find_tool_name_for_call_id(messages, call_id) {
|
||||
if !tool_names.contains(&name) {
|
||||
tool_names.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let tool_list = if tool_names.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
let mut unique = tool_names;
|
||||
unique.sort();
|
||||
unique.dedup();
|
||||
format!(" 使用过的工具: {}.", unique.join(", "))
|
||||
};
|
||||
|
||||
let snipped_count = snipped.len();
|
||||
let placeholder = ChatMessage::user(format!(
|
||||
"[上下文压缩] 省略了 {} 条中间对话消息(第 {}-{} 条)。{}",
|
||||
snipped_count,
|
||||
HEAD_KEEP + 1,
|
||||
tail_start,
|
||||
tool_list
|
||||
));
|
||||
|
||||
// 移除中间段,替换为占位消息
|
||||
messages.drain(HEAD_KEEP..tail_start);
|
||||
messages.insert(HEAD_KEEP, placeholder);
|
||||
|
||||
info!(
|
||||
"[snipCompact] {} → {} 条消息 (移除 {} 条, 切割点: {})",
|
||||
original_len,
|
||||
messages.len(),
|
||||
snipped_count,
|
||||
tail_start
|
||||
);
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// 根据 tool_call_id 查找对应的工具名称。
|
||||
fn find_tool_name_for_call_id(messages: &[ChatMessage], tool_call_id: &str) -> Option<String> {
|
||||
for msg in messages.iter().rev() {
|
||||
if msg.role == MessageRole::Assistant {
|
||||
if let Some(tool_calls) = &msg.tool_calls {
|
||||
for tc in tool_calls {
|
||||
if tc.id == tool_call_id {
|
||||
return Some(tc.function.name.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 轻量级压缩:将较早的工具结果替换为简短占位符,释放上下文空间。
|
||||
/// 保留最近 `keep_recent` 条工具结果不变。
|
||||
/// P0 改进:使用 `[Previous: used {tool_name}]` 替代字符预览,节省 ~80 tokens/条。
|
||||
pub fn micro_compact(messages: &mut [ChatMessage], keep_recent: usize) {
|
||||
let tool_info: Vec<(usize, String, String)> = messages
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, m)| {
|
||||
if m.role == MessageRole::Tool {
|
||||
let call_id = m.tool_call_id.clone().unwrap_or_default();
|
||||
// 先查找工具名称(此时 messages 是不可变借用)
|
||||
let tool_name = find_tool_name_for_call_id(messages, &call_id)
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
Some((i, call_id, tool_name))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let compact_count = tool_info.len().saturating_sub(keep_recent);
|
||||
if compact_count == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
for (idx, _tool_id, tool_name) in tool_info.iter().take(compact_count) {
|
||||
if let Some(msg) = messages.get_mut(*idx) {
|
||||
msg.content = Some(format!("[Previous: used {}]", tool_name));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 粗略估算消息列表的 token 数(用作首次调用的近似值)。
|
||||
/// 后续迭代优先使用 API 返回的精确 prompt_tokens。
|
||||
pub fn rough_estimate_tokens(messages: &[ChatMessage]) -> usize {
|
||||
messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let content_len = m.content.as_ref().map_or(0, |c| c.len());
|
||||
content_len + 4 // 消息结构 overhead 约 4 token
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// 保存完整 transcript 到磁盘(JSONL 格式)。
|
||||
/// 在压缩前调用,确保不丢失任何对话历史。
|
||||
async fn save_transcript(messages: &[ChatMessage], session_id: &str) {
|
||||
let dir = transcripts_dir();
|
||||
if let Err(e) = std::fs::create_dir_all(&dir) {
|
||||
warn!("[Compact] 无法创建 transcripts 目录 {:?}: {}", dir, e);
|
||||
return;
|
||||
}
|
||||
|
||||
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
|
||||
let filename = format!("{}_{}.jsonl", session_id, timestamp);
|
||||
let path = dir.join(&filename);
|
||||
|
||||
let mut content = String::new();
|
||||
for msg in messages {
|
||||
if let Ok(json) = serde_json::to_string(msg) {
|
||||
content.push_str(&json);
|
||||
content.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
match std::fs::write(&path, &content) {
|
||||
Ok(_) => info!(
|
||||
"[Compact] Transcript 已保存: {} ({} 条消息)",
|
||||
path.display(),
|
||||
messages.len()
|
||||
),
|
||||
Err(e) => warn!("[Compact] Transcript 保存失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
/// 在压缩后注入身份确认块,防止模型丢失上下文认知。
|
||||
/// 参考 Claude Code s11: identity re-injection after compression.
|
||||
fn inject_identity_block(messages: &mut Vec<ChatMessage>) {
|
||||
if messages.len() <= 4 {
|
||||
// 消息过少说明压缩非常激进,注入身份提醒
|
||||
let identity = ChatMessage::user(
|
||||
"[身份确认] 你是一位专业的天体物理学研究助手。以上是历史对话的压缩摘要。\
|
||||
你正在进行的研究任务是回答用户的问题。请基于摘要中的关键信息继续工作,\
|
||||
需要更多信息时主动使用工具搜索。",
|
||||
);
|
||||
// 插入在 system 消息和 summary 之后、recent 消息之前
|
||||
let insert_pos = if messages
|
||||
.first()
|
||||
.is_some_and(|m| m.role == MessageRole::System)
|
||||
{
|
||||
2.min(messages.len())
|
||||
} else {
|
||||
1.min(messages.len())
|
||||
};
|
||||
messages.insert(insert_pos, identity);
|
||||
info!(
|
||||
"[Compact] 注入身份确认块(压缩后仅 {} 条消息)",
|
||||
messages.len() - 1
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用 LLM 生成对话摘要。
|
||||
async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Result<String, String> {
|
||||
let summary_content: String = to_summarize
|
||||
.iter()
|
||||
.filter_map(|m| {
|
||||
let role = match m.role {
|
||||
MessageRole::User => "用户",
|
||||
MessageRole::Assistant => "助手",
|
||||
MessageRole::Tool => "工具",
|
||||
_ => return None,
|
||||
};
|
||||
m.content.as_ref().map(|c| {
|
||||
let preview: String = c.chars().take(200).collect();
|
||||
format!("[{}] {}", role, preview)
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let summary_prompt = format!(
|
||||
"请用简洁的中文总结以下对话历史的要点(不超过500字):\n\n{}",
|
||||
summary_content
|
||||
);
|
||||
|
||||
llm.chat_completion(
|
||||
"你是一个对话摘要助手。请提取对话的关键信息和结论。",
|
||||
&summary_prompt,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("[Compact] 上下文摘要生成失败: {},尝试激进压缩", e);
|
||||
format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len())
|
||||
})
|
||||
}
|
||||
|
||||
/// 多层回退压缩:snip → micro → auto → aggressive_micro → identity inject
|
||||
async fn compress_with_fallback(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
context_char_limit: usize,
|
||||
) {
|
||||
// Layer 0: snip_compact(零 API 调用,消息数超过 MAX_MESSAGES 时截断中间段)
|
||||
snip_compact(messages, MAX_MESSAGES);
|
||||
|
||||
// Layer 1: micro_compact(保留最近 8 条工具结果)
|
||||
micro_compact(messages, 8);
|
||||
if rough_estimate_tokens(messages) < (context_char_limit * 3 / 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Layer 2: auto_compact(LLM 摘要)
|
||||
let system_msg = messages.first().cloned();
|
||||
let cut_point = find_safe_cut_point(messages, 10);
|
||||
let to_summarize = &messages[1..cut_point];
|
||||
if to_summarize.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let summary = match generate_summary(to_summarize, llm).await {
|
||||
Ok(s) => s,
|
||||
Err(fallback) => fallback,
|
||||
};
|
||||
|
||||
let recent = messages[cut_point..].to_vec();
|
||||
messages.clear();
|
||||
if let Some(sys) = system_msg {
|
||||
messages.push(sys);
|
||||
}
|
||||
messages.push(ChatMessage::user(format!("[历史对话摘要]\n{}", summary)));
|
||||
messages.extend(recent);
|
||||
|
||||
// Layer 3: 如果摘要后仍然超限,激进 micro_compact(仅保留 2 条)
|
||||
if rough_estimate_tokens(messages) >= (context_char_limit * 3 / 2) {
|
||||
warn!("[Compact] LLM 摘要后仍超限,执行激进压缩 (keep_recent=2)");
|
||||
micro_compact(messages, 2);
|
||||
}
|
||||
|
||||
// Layer 4: 注入身份确认块
|
||||
inject_identity_block(messages);
|
||||
|
||||
info!(
|
||||
"[Compact] 上下文压缩完成,消息数: {} (安全切割点: {})",
|
||||
messages.len(),
|
||||
cut_point
|
||||
);
|
||||
}
|
||||
|
||||
/// 上下文压缩:保存 transcript → 多层回退压缩。
|
||||
/// 保留系统提示 + 最近的完整 tool-call/tool-result 配对。
|
||||
/// 使用 LLM 摘要较旧的对话历史,在 token 超限时触发。
|
||||
pub async fn compress_context(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
context_char_limit: usize,
|
||||
session_id: &str,
|
||||
) {
|
||||
compress_context_with_hooks(messages, llm, context_char_limit, session_id, None).await;
|
||||
}
|
||||
|
||||
/// 带 Hook 的上下文压缩变体。如果提供了 HookRegistry,会在压缩前后触发事件。
|
||||
pub async fn compress_context_with_hooks(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
context_char_limit: usize,
|
||||
session_id: &str,
|
||||
hook_registry: Option<&HookRegistry>,
|
||||
) {
|
||||
if messages.len() <= 4 {
|
||||
return;
|
||||
}
|
||||
|
||||
// 递归守卫:如果已在压缩中,跳过(防止嵌套压缩死循环)
|
||||
if COMPACTING.swap(true, Ordering::SeqCst) {
|
||||
warn!("[Compact] 递归守卫触发:已有进行中的压缩操作,跳过");
|
||||
return;
|
||||
}
|
||||
|
||||
let before_count = messages.len();
|
||||
let est_tokens = rough_estimate_tokens(messages);
|
||||
|
||||
// OnPreCompact hook
|
||||
if let Some(registry) = hook_registry {
|
||||
registry
|
||||
.run_on_pre_compact(&PreCompactContext {
|
||||
session_id: session_id.to_string(),
|
||||
message_count: before_count,
|
||||
estimated_tokens: est_tokens,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// 压缩前保存完整 transcript(使用传入的 session_id 避免跨会话覆盖)
|
||||
save_transcript(messages, session_id).await;
|
||||
|
||||
// 执行多层回退压缩
|
||||
compress_with_fallback(messages, llm, context_char_limit).await;
|
||||
|
||||
// OnPostCompact hook
|
||||
if let Some(registry) = hook_registry {
|
||||
registry
|
||||
.run_on_post_compact(&PostCompactContext {
|
||||
session_id: session_id.to_string(),
|
||||
new_message_count: messages.len(),
|
||||
compression_method: if messages.len() < before_count / 2 {
|
||||
"llm_summary"
|
||||
} else if messages.len() < before_count {
|
||||
// snip_compact 会产生占位消息但保留尾部,micro_compact 替换内容
|
||||
"snip_or_micro"
|
||||
} else {
|
||||
"none"
|
||||
}
|
||||
.to_string(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// 释放递归守卫
|
||||
COMPACTING.store(false, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_placeholder_format() {
|
||||
let mut messages = vec![
|
||||
ChatMessage::system("You are a helpful assistant."),
|
||||
ChatMessage::user("Hello"),
|
||||
ChatMessage::assistant_with_tool_calls(
|
||||
Some("Let me search.".to_string()),
|
||||
vec![crate::clients::llm::ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: crate::clients::llm::FunctionCall {
|
||||
name: "search_papers".to_string(),
|
||||
arguments: "{\"query\": \"black holes\"}".to_string(),
|
||||
},
|
||||
}],
|
||||
),
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some("Found 5 results about black holes".to_string()),
|
||||
tool_call_id: Some("call_1".to_string()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
ChatMessage::assistant("Here are the results..."),
|
||||
];
|
||||
|
||||
micro_compact(&mut messages, 0);
|
||||
|
||||
// 工具结果应该被压缩为 [Previous: used search_papers]
|
||||
let tool_msg = &messages[3];
|
||||
assert_eq!(
|
||||
tool_msg.content.as_deref(),
|
||||
Some("[Previous: used search_papers]")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_block_injected_when_few_messages() {
|
||||
let mut messages = vec![
|
||||
ChatMessage::system("You are a research assistant."),
|
||||
ChatMessage::user("[历史对话摘要]\nPrevious discussion about black holes."),
|
||||
ChatMessage::user("What about neutron stars?"),
|
||||
];
|
||||
|
||||
inject_identity_block(&mut messages);
|
||||
|
||||
// 应该注入了身份确认块
|
||||
let identity_msg = &messages[2];
|
||||
assert!(identity_msg.content.as_ref().unwrap().contains("身份确认"));
|
||||
assert!(identity_msg
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("天体物理学研究助手"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_preserves_recent_results() {
|
||||
let mut messages = vec![
|
||||
ChatMessage::system("You are a helpful assistant."),
|
||||
ChatMessage::user("Search for papers"),
|
||||
ChatMessage::assistant_with_tool_calls(
|
||||
Some("Searching...".to_string()),
|
||||
vec![crate::clients::llm::ToolCall {
|
||||
id: "call_old".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: crate::clients::llm::FunctionCall {
|
||||
name: "search_papers".to_string(),
|
||||
arguments: "{\"query\": \"old\"}".to_string(),
|
||||
},
|
||||
}],
|
||||
),
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some("Old result".to_string()),
|
||||
tool_call_id: Some("call_old".to_string()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
ChatMessage::assistant_with_tool_calls(
|
||||
Some("Searching more...".to_string()),
|
||||
vec![crate::clients::llm::ToolCall {
|
||||
id: "call_new".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: crate::clients::llm::FunctionCall {
|
||||
name: "get_paper_content".to_string(),
|
||||
arguments: "{\"bibcode\": \"2024A&A...\"}".to_string(),
|
||||
},
|
||||
}],
|
||||
),
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some("Paper content here...".to_string()),
|
||||
tool_call_id: Some("call_new".to_string()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
micro_compact(&mut messages, 1);
|
||||
|
||||
// 第一个工具结果应该被压缩
|
||||
let compressed = &messages[3];
|
||||
assert_eq!(
|
||||
compressed.content.as_deref(),
|
||||
Some("[Previous: used search_papers]")
|
||||
);
|
||||
|
||||
// 最近的一个工具结果应该保留
|
||||
let recent = &messages[5];
|
||||
assert_eq!(recent.content.as_deref(), Some("Paper content here..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_tool_name_for_call_id() {
|
||||
let messages = vec![
|
||||
ChatMessage::assistant_with_tool_calls(
|
||||
Some("Let me search.".to_string()),
|
||||
vec![crate::clients::llm::ToolCall {
|
||||
id: "call_abc".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: crate::clients::llm::FunctionCall {
|
||||
name: "rag_search".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
}],
|
||||
),
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some("Search result".to_string()),
|
||||
tool_call_id: Some("call_abc".to_string()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
let name = find_tool_name_for_call_id(&messages, "call_abc");
|
||||
assert_eq!(name, Some("rag_search".to_string()));
|
||||
|
||||
let name = find_tool_name_for_call_id(&messages, "nonexistent");
|
||||
assert_eq!(name, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rough_estimate_tokens() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("You are an assistant."),
|
||||
ChatMessage::user("Hello, world!"),
|
||||
];
|
||||
let estimate = rough_estimate_tokens(&messages);
|
||||
// 每个消息 content.len() + 4 overhead
|
||||
let expected = ("You are an assistant.".len()) + 4 + ("Hello, world!".len()) + 4;
|
||||
assert_eq!(estimate, expected);
|
||||
}
|
||||
|
||||
// ── snip_compact 测试 ──
|
||||
|
||||
#[test]
|
||||
fn test_snip_compact_below_threshold_noop() {
|
||||
let mut messages: Vec<ChatMessage> = (0..40)
|
||||
.map(|i| {
|
||||
if i == 0 {
|
||||
ChatMessage::system("System prompt")
|
||||
} else {
|
||||
ChatMessage::user(format!("Message {}", i))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let result = snip_compact(&mut messages, 50);
|
||||
assert!(!result);
|
||||
assert_eq!(messages.len(), 40);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snip_compact_above_threshold_truncates() {
|
||||
let mut messages: Vec<ChatMessage> = (0..100)
|
||||
.map(|i| {
|
||||
if i == 0 {
|
||||
ChatMessage::system("System prompt")
|
||||
} else {
|
||||
ChatMessage::user(format!("Message {}", i))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let original_len = messages.len();
|
||||
let result = snip_compact(&mut messages, 50);
|
||||
assert!(result);
|
||||
// 应该变成: HEAD_KEEP(3) + 1(placeholder) + remainder ≈ 51
|
||||
assert!(messages.len() < original_len);
|
||||
assert!(messages.len() <= 51); // HEAD_KEEP + placeholder + tail
|
||||
// 检查占位消息
|
||||
assert!(messages[HEAD_KEEP]
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("省略"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snip_compact_keeps_system_prompt() {
|
||||
let sys = ChatMessage::system("You are an astrophysics research assistant.");
|
||||
let mut messages: Vec<ChatMessage> = vec![sys.clone()];
|
||||
for i in 1..80 {
|
||||
messages.push(ChatMessage::user(format!("Question {}", i)));
|
||||
messages.push(ChatMessage::assistant(format!("Answer {}", i)));
|
||||
}
|
||||
snip_compact(&mut messages, 50);
|
||||
// 第一条必须是 system 消息
|
||||
assert_eq!(messages[0].role, MessageRole::System);
|
||||
assert_eq!(
|
||||
messages[0].content.as_deref(),
|
||||
Some("You are an astrophysics research assistant.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snip_compact_respects_tool_pairing() {
|
||||
// 构建 tool_call/tool_result 配对靠近切割点的场景
|
||||
let mut messages = vec![
|
||||
ChatMessage::system("System"),
|
||||
ChatMessage::user("Search"),
|
||||
ChatMessage::assistant_with_tool_calls(
|
||||
Some("Searching...".to_string()),
|
||||
vec![crate::clients::llm::ToolCall {
|
||||
id: "call_near_cut".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: crate::clients::llm::FunctionCall {
|
||||
name: "search_papers".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
}],
|
||||
),
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some("Result".to_string()),
|
||||
tool_call_id: Some("call_near_cut".to_string()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
// 填充到超过阈值
|
||||
for i in 0..60 {
|
||||
messages.push(ChatMessage::user(format!("Padding {}", i)));
|
||||
}
|
||||
snip_compact(&mut messages, MAX_MESSAGES);
|
||||
// 不应该有孤立的 tool 消息(没有对应 assistant(tool_calls))
|
||||
let has_orphan_tool = messages.windows(2).any(|w| {
|
||||
w[0].role == MessageRole::Tool
|
||||
&& w[1].role != MessageRole::Tool
|
||||
&& (w[1].role != MessageRole::Assistant || w[1].tool_calls.is_none())
|
||||
});
|
||||
assert!(!has_orphan_tool);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
// src/agent/compact/collapse.rs
|
||||
//
|
||||
// 上下文折叠日志 — 结构化 commit log + projection 模式。
|
||||
// 参考 Claude Code ContextCollapse:将压缩记录为分段 commit,
|
||||
// 需要时通过 projection 重放到消息层,避免直接修改原始历史。
|
||||
//
|
||||
// 核心思路:
|
||||
// 1. 每次压缩记录一个 CollapseCommit(范围 + 方法 + 摘要)
|
||||
// 2. project() 将 commits 应用到消息列表上
|
||||
// 3. 超过 MAX_SEGMENTS 时触发溢出合并
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tracing::info;
|
||||
|
||||
/// 最大折叠分段数,超出后触发溢出合并
|
||||
const MAX_COLLAPSE_SEGMENTS: usize = 5;
|
||||
/// 溢出摘要最大字符数
|
||||
const OVERFLOW_SUMMARY_MAX_CHARS: usize = 800;
|
||||
|
||||
/// 压缩方法枚举
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum CollapseMethod {
|
||||
/// 轻量级 micro 压缩(替换工具结果为占位符)
|
||||
MicroCompact,
|
||||
/// LLM 摘要压缩
|
||||
LlmSummary,
|
||||
/// 激进 micro 压缩(只保留极少数工具结果)
|
||||
AggressiveMicro,
|
||||
/// 身份注入(消息过少时的回退)
|
||||
IdentityInjection,
|
||||
}
|
||||
|
||||
impl CollapseMethod {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
CollapseMethod::MicroCompact => "micro",
|
||||
CollapseMethod::LlmSummary => "llm_summary",
|
||||
CollapseMethod::AggressiveMicro => "aggressive_micro",
|
||||
CollapseMethod::IdentityInjection => "identity_injection",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 单次压缩的 commit entry
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CollapseCommit {
|
||||
/// 自增 commit ID
|
||||
pub id: u64,
|
||||
/// 压缩方法
|
||||
pub method: CollapseMethod,
|
||||
/// 原始 messages 中被折叠的索引范围 (start, end_exclusive)
|
||||
pub removed_range: (usize, usize),
|
||||
/// 压缩后的摘要文本
|
||||
pub summary: String,
|
||||
/// commit 时间戳
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// 持久的折叠日志 — 记录和重放压缩历史。
|
||||
///
|
||||
/// 使用示例:
|
||||
/// ```ignore
|
||||
/// let mut log = CollapseLog::new();
|
||||
/// // 每次压缩后记录
|
||||
/// log.commit(CollapseMethod::MicroCompact, (5, 20), "摘要内容".into());
|
||||
/// // 可将 commits 投影到消息上
|
||||
/// log.project(&mut messages);
|
||||
/// ```
|
||||
pub struct CollapseLog {
|
||||
commits: Vec<CollapseCommit>,
|
||||
next_id: AtomicU64,
|
||||
}
|
||||
|
||||
impl Default for CollapseLog {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl CollapseLog {
|
||||
/// 创建空的折叠日志
|
||||
pub fn new() -> Self {
|
||||
CollapseLog {
|
||||
commits: Vec::new(),
|
||||
next_id: AtomicU64::new(1),
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录一次压缩 commit。
|
||||
///
|
||||
/// 返回新 commit 的 ID。
|
||||
pub fn commit(
|
||||
&mut self,
|
||||
method: CollapseMethod,
|
||||
removed_range: (usize, usize),
|
||||
summary: String,
|
||||
) -> u64 {
|
||||
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||||
let method_str = method.as_str();
|
||||
let summary_len = summary.len();
|
||||
self.commits.push(CollapseCommit {
|
||||
id,
|
||||
method,
|
||||
removed_range,
|
||||
summary,
|
||||
timestamp: chrono::Utc::now(),
|
||||
});
|
||||
info!(
|
||||
"[CollapseLog] commit #{}: method={}, range={:?}, summary_len={}",
|
||||
id, method_str, removed_range, summary_len
|
||||
);
|
||||
id
|
||||
}
|
||||
|
||||
/// 获取 commit 总数
|
||||
pub fn len(&self) -> usize {
|
||||
self.commits.len()
|
||||
}
|
||||
|
||||
/// 是否有记录
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.commits.is_empty()
|
||||
}
|
||||
|
||||
/// 检查是否需要溢出合并(commits 数超过 MAX_COLLAPSE_SEGMENTS)
|
||||
pub fn should_overflow(&self) -> bool {
|
||||
self.commits.len() > MAX_COLLAPSE_SEGMENTS
|
||||
}
|
||||
|
||||
/// 获取最近的 N 条摘要
|
||||
pub fn recent_summaries(&self, n: usize) -> Vec<String> {
|
||||
self.commits
|
||||
.iter()
|
||||
.rev()
|
||||
.take(n)
|
||||
.map(|c| c.summary.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 将最老的 commits 合并为一个溢出摘要。
|
||||
///
|
||||
/// 返回溢出摘要(可插入到消息列表中),并从日志中移除已合并的 commits。
|
||||
pub fn overflow(&mut self) -> Option<String> {
|
||||
if self.commits.len() <= MAX_COLLAPSE_SEGMENTS {
|
||||
return None;
|
||||
}
|
||||
|
||||
let merge_count = self.commits.len() - MAX_COLLAPSE_SEGMENTS + 1;
|
||||
let to_merge: Vec<&CollapseCommit> = self.commits.iter().take(merge_count).collect();
|
||||
|
||||
let mut merged = String::from("[上下文压缩历史]\n");
|
||||
for commit in &to_merge {
|
||||
let summary_preview: String = commit
|
||||
.summary
|
||||
.chars()
|
||||
.take(OVERFLOW_SUMMARY_MAX_CHARS)
|
||||
.collect();
|
||||
merged.push_str(&format!(
|
||||
"- (方法: {}) {}\n",
|
||||
commit.method.as_str(),
|
||||
summary_preview
|
||||
));
|
||||
}
|
||||
|
||||
// 移除已合并的 commits(保留最后 MAX_COLLAPSE_SEGMENTS-1 个)
|
||||
self.commits.drain(0..merge_count);
|
||||
|
||||
info!(
|
||||
"[CollapseLog] 溢出合并: {} commits → {} chars, 剩余 {} commits",
|
||||
merge_count,
|
||||
merged.len(),
|
||||
self.commits.len()
|
||||
);
|
||||
|
||||
Some(merged)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_commit_and_len() {
|
||||
let mut log = CollapseLog::new();
|
||||
assert_eq!(log.len(), 0);
|
||||
assert!(log.is_empty());
|
||||
|
||||
log.commit(CollapseMethod::MicroCompact, (5, 10), "summary 1".into());
|
||||
assert_eq!(log.len(), 1);
|
||||
assert!(!log.is_empty());
|
||||
|
||||
log.commit(CollapseMethod::LlmSummary, (0, 5), "summary 2".into());
|
||||
assert_eq!(log.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_should_overflow() {
|
||||
let mut log = CollapseLog::new();
|
||||
// 添加 6 个 commits(超过 MAX_COLLAPSE_SEGMENTS=5)
|
||||
for i in 0..6 {
|
||||
log.commit(
|
||||
CollapseMethod::MicroCompact,
|
||||
(i * 10, i * 10 + 5),
|
||||
format!("commit {}", i),
|
||||
);
|
||||
}
|
||||
assert!(log.should_overflow());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_overflow_when_under_limit() {
|
||||
let mut log = CollapseLog::new();
|
||||
for i in 0..5 {
|
||||
log.commit(
|
||||
CollapseMethod::MicroCompact,
|
||||
(i * 10, i * 10 + 5),
|
||||
format!("commit {}", i),
|
||||
);
|
||||
}
|
||||
assert!(!log.should_overflow());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_overflow_merges_oldest() {
|
||||
let mut log = CollapseLog::new();
|
||||
// 添加 7 个 commits
|
||||
for i in 0..7 {
|
||||
log.commit(
|
||||
CollapseMethod::MicroCompact,
|
||||
(i * 10, i * 10 + 5),
|
||||
format!("summary for commit {}", i),
|
||||
);
|
||||
}
|
||||
assert!(log.should_overflow());
|
||||
|
||||
let merged = log.overflow();
|
||||
assert!(merged.is_some());
|
||||
// 溢出后应剩余 MAX_COLLAPSE_SEGMENTS-1 = 4 个 commits(合并了 3 个)
|
||||
assert!(!log.should_overflow());
|
||||
assert_eq!(log.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recent_summaries() {
|
||||
let mut log = CollapseLog::new();
|
||||
for i in 0..3 {
|
||||
log.commit(
|
||||
CollapseMethod::MicroCompact,
|
||||
(i, i + 1),
|
||||
format!("summary {}", i),
|
||||
);
|
||||
}
|
||||
|
||||
let recent = log.recent_summaries(2);
|
||||
assert_eq!(recent.len(), 2);
|
||||
assert_eq!(recent[0], "summary 2"); // 最新的在前
|
||||
assert_eq!(recent[1], "summary 1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_log_no_overflow() {
|
||||
let log = CollapseLog::new();
|
||||
assert!(!log.should_overflow());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,974 @@
|
||||
// src/agent/hooks.rs
|
||||
//
|
||||
// Agent 生命周期 Hooks 系统。
|
||||
// 参考 Claude Code 的 PreToolUse / PostToolUse / Stop hooks 设计,
|
||||
// 提供可扩展的事件回调链,支持:
|
||||
// - OnSessionStart — 会话创建/恢复时
|
||||
// - PreToolUse — 工具执行前(可拦截/阻止/修改输入/注入上下文)
|
||||
// - PostToolUse — 工具执行后(审计日志、指标采集、输出修改)
|
||||
// - OnStepComplete — 每步结束(指标更新、上下文检查)
|
||||
// - OnSessionStop — 会话终止(清理、持久化)
|
||||
// - OnSubagentStart — 子代理启动时
|
||||
// - OnSubagentStop — 子代理停止时
|
||||
// - OnPreCompact — 上下文压缩前
|
||||
// - OnPostCompact — 上下文压缩后
|
||||
//
|
||||
// P3 增强(参考 Claude Code hooks 协议):
|
||||
// - PreToolUseAction 支持 MutateInput(修改工具参数)和 PermissionRequired
|
||||
// - PostToolUseAction 支持 MutateOutput(修改工具输出)
|
||||
// - 扩展生命周期事件(SubagentStart/Stop, PreCompact/PostCompact)
|
||||
// - HookRegistry 返回累积的 additional_context 和 mutate_output
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::SqlitePool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::terminal::TurnTerminal;
|
||||
|
||||
// ── Hook Contexts ──
|
||||
|
||||
/// 会话启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStartContext {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
pub is_resume: bool,
|
||||
}
|
||||
|
||||
/// PreToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseContext {
|
||||
pub session_id: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub step: usize,
|
||||
}
|
||||
|
||||
/// PostToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub output_content: String,
|
||||
pub is_error: bool,
|
||||
pub step: usize,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
/// Step 完成上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StepCompleteContext {
|
||||
pub session_id: String,
|
||||
pub step: usize,
|
||||
pub max_steps: usize,
|
||||
pub messages_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
pub token_limit: usize,
|
||||
}
|
||||
|
||||
/// Session 终止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStopContext<'a> {
|
||||
pub session_id: String,
|
||||
pub terminal: &'a TurnTerminal,
|
||||
pub total_steps: usize,
|
||||
}
|
||||
|
||||
/// 子代理启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStartContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
/// 子代理停止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStopContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub result_summary: String,
|
||||
pub steps: usize,
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
/// 压缩前上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreCompactContext {
|
||||
pub session_id: String,
|
||||
pub message_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
}
|
||||
|
||||
/// 压缩后上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostCompactContext {
|
||||
pub session_id: String,
|
||||
pub new_message_count: usize,
|
||||
pub compression_method: String,
|
||||
}
|
||||
|
||||
// ── Hook Actions ──
|
||||
|
||||
/// PreToolUse hook 返回的增强动作。
|
||||
/// 支持:允许、阻止、修改输入、权限请求。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PreToolUseAction {
|
||||
/// 允许继续执行(默认)
|
||||
Continue,
|
||||
/// 阻止执行,附带原因
|
||||
Block { reason: String },
|
||||
/// 允许执行但修改输入参数或注入附加上下文
|
||||
MutateInput {
|
||||
updated_args: serde_json::Value,
|
||||
additional_context: Option<String>,
|
||||
},
|
||||
/// 需要权限决策
|
||||
PermissionRequired {
|
||||
permission: String,
|
||||
tool_name: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl PreToolUseAction {
|
||||
pub fn is_blocked(&self) -> bool {
|
||||
matches!(self, PreToolUseAction::Block { .. })
|
||||
}
|
||||
|
||||
pub fn block_reason(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::Block { reason } => Some(reason),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn updated_args(&self) -> Option<&serde_json::Value> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput { updated_args, .. } => Some(updated_args),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn additional_context(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput {
|
||||
additional_context, ..
|
||||
} => additional_context.as_deref(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 向后兼容类型别名 — 旧代码用 HookAction::Continue / HookAction::Block 仍可编译
|
||||
pub type HookAction = PreToolUseAction;
|
||||
|
||||
/// PostToolUse hook 返回的动作。
|
||||
/// 支持:保持输出不变、修改输出内容。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PostToolUseAction {
|
||||
/// 保持原输出不变(默认)
|
||||
Continue,
|
||||
/// 修改输出内容
|
||||
MutateOutput { updated_content: String },
|
||||
}
|
||||
|
||||
// ── Metrics Data ──
|
||||
|
||||
/// 可查询的运行指标快照
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct MetricsData {
|
||||
pub tool_call_counts: HashMap<String, usize>,
|
||||
pub total_steps: usize,
|
||||
pub total_errors: usize,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
// ── Hook Trait ──
|
||||
|
||||
/// Agent 生命周期 Hook trait。
|
||||
/// 所有方法都有默认空实现,只需覆写关心的 hook 点。
|
||||
#[async_trait]
|
||||
pub trait AgentHook: Send + Sync {
|
||||
/// Hook 名称(用于日志和调试)
|
||||
fn name(&self) -> &str;
|
||||
|
||||
// ── 原有 5 个生命周期事件 ──
|
||||
|
||||
/// 会话创建/恢复时调用。
|
||||
async fn on_session_start(&self, _ctx: &SessionStartContext) {}
|
||||
|
||||
/// 工具执行前调用。可返回 Continue/Block/MutateInput/PermissionRequired。
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 工具执行后调用。可返回 Continue 或 MutateOutput。
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 每个 ReAct step 完成后调用。
|
||||
async fn on_step_complete(&self, _ctx: &StepCompleteContext) {}
|
||||
|
||||
/// 会话终止时调用。
|
||||
async fn on_session_stop(&self, _ctx: &SessionStopContext<'_>) {}
|
||||
|
||||
// ── 新增 4 个生命周期事件(默认 no-op) ──
|
||||
|
||||
/// 子代理启动时调用。
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {}
|
||||
|
||||
/// 子代理停止时调用。
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {}
|
||||
|
||||
/// 上下文压缩前调用。
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {}
|
||||
|
||||
/// 上下文压缩后调用。
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {}
|
||||
}
|
||||
|
||||
// ── Hook Registry ──
|
||||
|
||||
/// PreToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseResult {
|
||||
/// 最终动作(第一个 Block 获胜)
|
||||
pub action: PreToolUseAction,
|
||||
/// 累积的 additional_context(所有 MutateInput 的上下文拼接)
|
||||
pub additional_context: Option<String>,
|
||||
/// 最终的工具参数(应用了最后一个 MutateInput 的修改)
|
||||
pub final_args: serde_json::Value,
|
||||
}
|
||||
|
||||
/// PostToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseResult {
|
||||
/// 最终输出内容(应用了最后一个 MutateOutput 的修改)
|
||||
pub final_content: String,
|
||||
}
|
||||
|
||||
/// Hook 注册表,管理所有已注册的 hook 并按序调用
|
||||
pub struct HookRegistry {
|
||||
hooks: Vec<Box<dyn AgentHook>>,
|
||||
}
|
||||
|
||||
impl Default for HookRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl HookRegistry {
|
||||
/// 创建空的注册表
|
||||
pub fn new() -> Self {
|
||||
HookRegistry { hooks: Vec::new() }
|
||||
}
|
||||
|
||||
/// 创建包含所有内置 hooks 的注册表。
|
||||
///
|
||||
/// 参数:
|
||||
/// - `db`: 数据库连接池(供 AuditLogHook 持久化)
|
||||
/// - `cancelled_runs`: 取消状态集合(供 CancellationHook 检查)
|
||||
/// - `metrics_data`: 可选的共享指标数据引用。提供时复用已有的 MetricsData,
|
||||
/// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。
|
||||
pub fn with_builtins(
|
||||
db: SqlitePool,
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
metrics_data: Option<Arc<std::sync::Mutex<MetricsData>>>,
|
||||
) -> Self {
|
||||
let mut registry = Self::new();
|
||||
registry.add(Box::new(CancellationHook::new(cancelled_runs)));
|
||||
// 如果提供了共享的 metrics_data,使用它;否则创建新的
|
||||
let metrics_hook = match metrics_data {
|
||||
Some(data) => MetricsHook::from_arc(data),
|
||||
None => MetricsHook::new(),
|
||||
};
|
||||
registry.add(Box::new(metrics_hook));
|
||||
registry.add(Box::new(AuditLogHook::new(db)));
|
||||
registry
|
||||
}
|
||||
|
||||
/// 注册一个 hook
|
||||
pub fn add(&mut self, hook: Box<dyn AgentHook>) {
|
||||
info!("[Hooks] 注册 hook: {}", hook.name());
|
||||
self.hooks.push(hook);
|
||||
}
|
||||
|
||||
/// 获取所有 hooks 的不可变引用
|
||||
pub fn all(&self) -> &[Box<dyn AgentHook>] {
|
||||
&self.hooks
|
||||
}
|
||||
|
||||
// ── 便捷调用方法 ──
|
||||
|
||||
/// 调用所有 on_session_start hooks
|
||||
pub async fn run_on_session_start(&self, ctx: &SessionStartContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_session_start(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 pre_tool_use hooks。
|
||||
///
|
||||
/// 返回聚合结果:
|
||||
/// - 遇到第一个 Block 时短路,返回该 Block
|
||||
/// - MutateInput 累积 additional_context 并更新 final_args
|
||||
/// - PermissionRequired 记录但继续执行(暂时视为 Continue)
|
||||
pub async fn run_pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseResult {
|
||||
let mut accumulated_context = String::new();
|
||||
let mut final_args = ctx.tool_args.clone();
|
||||
let mut final_action = PreToolUseAction::Continue;
|
||||
|
||||
for hook in &self.hooks {
|
||||
let action = hook.pre_tool_use(ctx).await;
|
||||
match &action {
|
||||
PreToolUseAction::Block { reason } => {
|
||||
warn!(
|
||||
"[Hooks] {} 阻止了工具 {} 的执行: {}",
|
||||
hook.name(),
|
||||
ctx.tool_name,
|
||||
reason
|
||||
);
|
||||
return PreToolUseResult {
|
||||
action,
|
||||
additional_context: None,
|
||||
final_args: ctx.tool_args.clone(),
|
||||
};
|
||||
}
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args,
|
||||
additional_context,
|
||||
} => {
|
||||
info!(
|
||||
"[Hooks] {} 修改了工具 {} 的输入参数",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
final_args = updated_args.clone();
|
||||
if let Some(ctx_str) = additional_context {
|
||||
if !accumulated_context.is_empty() {
|
||||
accumulated_context.push('\n');
|
||||
}
|
||||
accumulated_context.push_str(ctx_str);
|
||||
}
|
||||
}
|
||||
PreToolUseAction::PermissionRequired { .. } => {
|
||||
// 暂时记录但继续执行(Permission 系统在 Phase 2 中完善)
|
||||
info!(
|
||||
"[Hooks] {} 请求了工具 {} 的权限检查",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
}
|
||||
PreToolUseAction::Continue => {}
|
||||
}
|
||||
final_action = action;
|
||||
}
|
||||
|
||||
let ctx_opt = if accumulated_context.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_context)
|
||||
};
|
||||
|
||||
PreToolUseResult {
|
||||
action: final_action,
|
||||
additional_context: ctx_opt,
|
||||
final_args,
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 post_tool_use hooks(全部执行,不会短路)。
|
||||
/// 返回聚合的最终输出内容。
|
||||
pub async fn run_post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseResult {
|
||||
let mut final_content = ctx.output_content.clone();
|
||||
|
||||
for hook in &self.hooks {
|
||||
let action = hook.post_tool_use(ctx).await;
|
||||
match action {
|
||||
PostToolUseAction::MutateOutput { updated_content } => {
|
||||
info!(
|
||||
"[Hooks] {} 修改了工具 {} 的输出",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
final_content = updated_content;
|
||||
}
|
||||
PostToolUseAction::Continue => {}
|
||||
}
|
||||
}
|
||||
|
||||
PostToolUseResult { final_content }
|
||||
}
|
||||
|
||||
/// 调用所有 on_step_complete hooks
|
||||
pub async fn run_on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_step_complete(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_session_stop hooks
|
||||
pub async fn run_on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_session_stop(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_start hooks
|
||||
pub async fn run_on_subagent_start(&self, ctx: &SubagentStartContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_subagent_start(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_stop hooks
|
||||
pub async fn run_on_subagent_stop(&self, ctx: &SubagentStopContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_subagent_stop(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_pre_compact hooks
|
||||
pub async fn run_on_pre_compact(&self, ctx: &PreCompactContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_pre_compact(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_post_compact hooks
|
||||
pub async fn run_on_post_compact(&self, ctx: &PostCompactContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_post_compact(ctx).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Built-in Hooks ──
|
||||
|
||||
/// 取消检查 Hook — 在每次工具执行前检查用户是否中止了会话。
|
||||
pub struct CancellationHook {
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
}
|
||||
|
||||
impl CancellationHook {
|
||||
pub fn new(cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>) -> Self {
|
||||
CancellationHook { cancelled_runs }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for CancellationHook {
|
||||
fn name(&self) -> &str {
|
||||
"CancellationHook"
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
if let Ok(runs) = self.cancelled_runs.lock() {
|
||||
if runs.contains(&ctx.session_id) {
|
||||
warn!(
|
||||
"[CancellationHook] 会话 {} 已被用户取消,阻止工具 {} 执行",
|
||||
ctx.session_id, ctx.tool_name
|
||||
);
|
||||
return PreToolUseAction::Block {
|
||||
reason: "用户已手动中止执行".to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
// 清理取消状态
|
||||
if let Ok(mut runs) = self.cancelled_runs.lock() {
|
||||
runs.remove(&ctx.session_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 指标采集 Hook — 自动收集工具调用统计,支持快照查询
|
||||
pub struct MetricsHook {
|
||||
data: Arc<std::sync::Mutex<MetricsData>>,
|
||||
}
|
||||
|
||||
impl Default for MetricsHook {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl MetricsHook {
|
||||
pub fn new() -> Self {
|
||||
MetricsHook {
|
||||
data: Arc::new(std::sync::Mutex::new(MetricsData::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从已有的 Arc<Mutex<MetricsData>> 创建(共享数据引用)
|
||||
pub fn from_arc(data: Arc<std::sync::Mutex<MetricsData>>) -> Self {
|
||||
MetricsHook { data }
|
||||
}
|
||||
|
||||
/// 返回当前指标快照(锁异常时返回 None)
|
||||
pub fn snapshot(&self) -> Option<MetricsData> {
|
||||
self.data.lock().ok().map(|d| d.clone())
|
||||
}
|
||||
|
||||
/// 获取 Arc 引用,供外部持有
|
||||
pub fn data_arc(&self) -> Arc<std::sync::Mutex<MetricsData>> {
|
||||
self.data.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for MetricsHook {
|
||||
fn name(&self) -> &str {
|
||||
"MetricsHook"
|
||||
}
|
||||
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
data.session_id = Some(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
info!(
|
||||
"[Metrics] step={} tool={} is_error={} elapsed={}ms",
|
||||
ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms
|
||||
);
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
*data
|
||||
.tool_call_counts
|
||||
.entry(ctx.tool_name.clone())
|
||||
.or_insert(0) += 1;
|
||||
data.total_steps = ctx.step;
|
||||
if ctx.is_error {
|
||||
data.total_errors += 1;
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
if ctx.step.is_multiple_of(3) {
|
||||
info!(
|
||||
"[Metrics] step {}/{} | messages={} | tokens≈{}/{}",
|
||||
ctx.step, ctx.max_steps, ctx.messages_count, ctx.estimated_tokens, ctx.token_limit
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Metrics] 会话 {} 结束: {} (total_steps={})",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description(),
|
||||
ctx.total_steps
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 审计日志 Hook — 记录所有工具调用到 SQLite agent_audit_log 表
|
||||
pub struct AuditLogHook {
|
||||
db: SqlitePool,
|
||||
}
|
||||
|
||||
impl AuditLogHook {
|
||||
pub fn new(db: SqlitePool) -> Self {
|
||||
AuditLogHook { db }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for AuditLogHook {
|
||||
fn name(&self) -> &str {
|
||||
"AuditLogHook"
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
let status = if ctx.is_error { "FAIL" } else { "OK" };
|
||||
let preview: String = ctx.output_content.chars().take(200).collect();
|
||||
|
||||
info!(
|
||||
"[Audit] session={} step={} tool={} status={} elapsed={}ms",
|
||||
ctx.session_id, ctx.step, ctx.tool_name, status, ctx.elapsed_ms
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let agent_name = ctx.agent_name.clone();
|
||||
let tool_name = ctx.tool_name.clone();
|
||||
let step = ctx.step;
|
||||
let elapsed = ctx.elapsed_ms;
|
||||
let preview_clone = preview.clone();
|
||||
|
||||
// Fire-and-forget 写入,不阻塞主循环
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview, agent_name) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(step as i32)
|
||||
.bind(&tool_name)
|
||||
.bind(status)
|
||||
.bind(elapsed as i32)
|
||||
.bind(&preview_clone)
|
||||
.bind(&agent_name)
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Audit] 会话 {} 终止原因: {}",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description()
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let total_steps = ctx.total_steps;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview) \
|
||||
VALUES (?, ?, 'session', 'SESSION_STOP', 0, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(total_steps as i32)
|
||||
.bind(format!("会话终止,共 {} 步", total_steps))
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
struct TestHook {
|
||||
name: String,
|
||||
pre_called: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
impl TestHook {
|
||||
fn new(name: &str) -> Self {
|
||||
TestHook {
|
||||
name: name.to_string(),
|
||||
pre_called: std::sync::Mutex::new(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for TestHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
*self.pre_called.lock().unwrap() = true;
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_registry_runs_all_hooks() {
|
||||
let mut registry = HookRegistry::new();
|
||||
let hook1 = TestHook::new("test1");
|
||||
let hook2 = TestHook::new("test2");
|
||||
registry.add(Box::new(hook1));
|
||||
registry.add(Box::new(hook2));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(!result.action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_blocking_hook_stops_chain() {
|
||||
struct BlockingHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for BlockingHook {
|
||||
fn name(&self) -> &str {
|
||||
"blocker"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Block {
|
||||
reason: "test block".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(BlockingHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(result.action.is_blocked());
|
||||
assert_eq!(result.action.block_reason(), Some("test block"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mutate_input_accumulates_context() {
|
||||
struct MutateHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateHook {
|
||||
fn name(&self) -> &str {
|
||||
"mutator"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args: serde_json::json!({"key": "modified"}),
|
||||
additional_context: Some("injected context".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({"key": "original"}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_args, serde_json::json!({"key": "modified"}));
|
||||
assert_eq!(
|
||||
result.additional_context,
|
||||
Some("injected context".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_tool_use_mutate_output() {
|
||||
struct MutateOutputHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateOutputHook {
|
||||
fn name(&self) -> &str {
|
||||
"output_mutator"
|
||||
}
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content: "modified output".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateOutputHook));
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "original output".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
|
||||
let result = registry.run_post_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_content, "modified output");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_blocks_when_cancelled() {
|
||||
use std::collections::HashSet;
|
||||
let mut cancelled = HashSet::new();
|
||||
cancelled.insert("test_session".to_string());
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_allows_when_not_cancelled() {
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(!action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_metrics_hook_accumulates_counts() {
|
||||
let hook = MetricsHook::new();
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
hook.post_tool_use(&ctx).await;
|
||||
|
||||
let ctx2 = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result2".into(),
|
||||
is_error: false,
|
||||
step: 2,
|
||||
elapsed_ms: 200,
|
||||
};
|
||||
hook.post_tool_use(&ctx2).await;
|
||||
|
||||
let snapshot = hook.snapshot().expect("snapshot should succeed in test");
|
||||
assert_eq!(snapshot.tool_call_counts.get("search_papers"), Some(&2));
|
||||
assert_eq!(snapshot.total_steps, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_session_start_hook_called() {
|
||||
struct StartTrackingHook {
|
||||
started: std::sync::Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for StartTrackingHook {
|
||||
fn name(&self) -> &str {
|
||||
"start_tracker"
|
||||
}
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
self.started.lock().unwrap().push(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let hook = StartTrackingHook {
|
||||
started: std::sync::Mutex::new(Vec::new()),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(hook));
|
||||
|
||||
let ctx = SessionStartContext {
|
||||
session_id: "test_sid".into(),
|
||||
turn_index: 1,
|
||||
is_resume: false,
|
||||
};
|
||||
registry.run_on_session_start(&ctx).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_new_lifecycle_events_called() {
|
||||
struct LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex<bool>,
|
||||
subagent_stop: std::sync::Mutex<bool>,
|
||||
pre_compact: std::sync::Mutex<bool>,
|
||||
post_compact: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for LifecycleTracker {
|
||||
fn name(&self) -> &str {
|
||||
"lifecycle_tracker"
|
||||
}
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {
|
||||
*self.subagent_start.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {
|
||||
*self.subagent_stop.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {
|
||||
*self.pre_compact.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {
|
||||
*self.post_compact.lock().unwrap() = true;
|
||||
}
|
||||
}
|
||||
|
||||
let tracker = LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex::new(false),
|
||||
subagent_stop: std::sync::Mutex::new(false),
|
||||
pre_compact: std::sync::Mutex::new(false),
|
||||
post_compact: std::sync::Mutex::new(false),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(tracker));
|
||||
|
||||
registry
|
||||
.run_on_subagent_start(&SubagentStartContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
prompt: "test".into(),
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_subagent_stop(&SubagentStopContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
result_summary: "done".into(),
|
||||
steps: 3,
|
||||
is_error: false,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_pre_compact(&PreCompactContext {
|
||||
session_id: "s1".into(),
|
||||
message_count: 50,
|
||||
estimated_tokens: 10000,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_post_compact(&PostCompactContext {
|
||||
session_id: "s1".into(),
|
||||
new_message_count: 10,
|
||||
compression_method: "micro".into(),
|
||||
})
|
||||
.await;
|
||||
|
||||
// If no panic, all hooks were called successfully
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
// src/agent/memory/age.rs
|
||||
//
|
||||
// 记忆时效性追踪 — 参考 Claude Code memdir/memoryAge.ts。
|
||||
//
|
||||
// LLM 不擅长日期计算,"2026-01-15" 不会触发过时判断,
|
||||
// 但 "47 天前" 会。本模块提供人类可读的时效标签,
|
||||
// 注入到 system prompt 中以引导模型在使用记忆中之前核实。
|
||||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// 返回当前 Unix 时间戳(秒)
|
||||
pub fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
/// 距离给定 Unix 时间戳的天数(向下取整)。
|
||||
/// 0 = 今天, 1 = 昨天, 2+ = 更早。
|
||||
/// 负数输入(未来时间/时钟偏差)截断为 0。
|
||||
pub fn memory_age_days(mtime_secs: u64) -> u64 {
|
||||
let now = now_secs();
|
||||
if mtime_secs >= now {
|
||||
return 0;
|
||||
}
|
||||
(now - mtime_secs) / 86_400
|
||||
}
|
||||
|
||||
/// 人类可读的时效标签。
|
||||
pub fn memory_age_label(mtime_secs: u64) -> String {
|
||||
let days = memory_age_days(mtime_secs);
|
||||
match days {
|
||||
0 => "今天".to_string(),
|
||||
1 => "昨天".to_string(),
|
||||
n => format!("{} 天前", n),
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回时效警告文本,如果记忆超过 1 天则返回 Some。
|
||||
/// 新鲜记忆(今天/昨天)返回 None — 此时警告只是噪音。
|
||||
pub fn memory_freshness_text(mtime_secs: u64) -> Option<String> {
|
||||
let days = memory_age_days(mtime_secs);
|
||||
if days <= 1 {
|
||||
return None;
|
||||
}
|
||||
Some(format!(
|
||||
"此记忆已有 {} 天。记忆是时间点快照,不是实时状态 — \
|
||||
关于代码行为或文件:行号的声明可能已过时。\
|
||||
请在断言为事实前与当前代码进行核对。",
|
||||
days
|
||||
))
|
||||
}
|
||||
|
||||
/// 包裹在 <system-reminder> 标签中的时效注释。
|
||||
/// 对于 ≤ 1 天的记忆返回空字符串。
|
||||
pub fn memory_freshness_note(mtime_secs: u64) -> String {
|
||||
match memory_freshness_text(mtime_secs) {
|
||||
Some(text) => format!("<system-reminder>{}</system-reminder>", text),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_age_days_today() {
|
||||
let now = now_secs();
|
||||
assert_eq!(memory_age_days(now), 0);
|
||||
assert_eq!(memory_age_days(now - 100), 0); // 100 秒前仍是今天
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_age_days_yesterday() {
|
||||
let now = now_secs();
|
||||
assert_eq!(memory_age_days(now - 86_400), 1);
|
||||
assert_eq!(memory_age_days(now - 86_400 - 100), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_age_days_older() {
|
||||
let now = now_secs();
|
||||
assert_eq!(memory_age_days(now - 86_400 * 3), 3);
|
||||
assert_eq!(memory_age_days(now - 86_400 * 47), 47);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_age_days_clamps_future_to_zero() {
|
||||
let future = now_secs() + 86_400 * 10;
|
||||
assert_eq!(memory_age_days(future), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_age_label() {
|
||||
let now = now_secs();
|
||||
assert_eq!(memory_age_label(now), "今天");
|
||||
assert_eq!(memory_age_label(now - 86_400), "昨天");
|
||||
assert_eq!(memory_age_label(now - 86_400 * 5), "5 天前");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_freshness_text_none_for_fresh() {
|
||||
let now = now_secs();
|
||||
assert!(memory_freshness_text(now).is_none()); // 今天
|
||||
assert!(memory_freshness_text(now - 86_400).is_none()); // 昨天
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_freshness_text_some_for_old() {
|
||||
let now = now_secs();
|
||||
let text = memory_freshness_text(now - 86_400 * 2);
|
||||
assert!(text.is_some());
|
||||
assert!(text.unwrap().contains("2 天"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_freshness_note_contains_tags() {
|
||||
let now = now_secs();
|
||||
let note = memory_freshness_note(now - 86_400 * 3);
|
||||
assert!(note.contains("<system-reminder>"));
|
||||
assert!(note.contains("</system-reminder>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_freshness_note_empty_for_fresh() {
|
||||
let now = now_secs();
|
||||
assert_eq!(memory_freshness_note(now), "");
|
||||
assert_eq!(memory_freshness_note(now - 86_400), "");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
// src/agent/memory/decay.rs
|
||||
//
|
||||
// 指数时间衰减评分 — 参考 Martian-Engineering/agent-memory。
|
||||
//
|
||||
// 核心理念:LLM 判断语义相关性,数学衰减提供时序排序。
|
||||
// "semantic decay adds LLM judgment, recency scoring adds temporal ordering"
|
||||
// — 两者互补,4天前的关键偏好可胜过1天前的临时备注。
|
||||
//
|
||||
// 公式: score = e^(-λ × days_old)
|
||||
// λ = ln(2) / half_life_days
|
||||
//
|
||||
// 30天半衰期下: 0天=1.0, 15天≈0.707, 30天=0.5, 60天=0.25, 90天≈0.125
|
||||
|
||||
/// 默认半衰期(天)
|
||||
pub const DEFAULT_HALF_LIFE_DAYS: f64 = 30.0;
|
||||
|
||||
/// 计算指数时间衰减评分。
|
||||
///
|
||||
/// score = e^(-λ × days_old),λ = ln(2) / half_life
|
||||
pub fn decay_score(mtime_secs: u64, half_life_days: f64) -> f64 {
|
||||
let _now = super::age::now_secs();
|
||||
let days = super::age::memory_age_days(mtime_secs) as f64;
|
||||
let lambda = std::f64::consts::LN_2 / half_life_days;
|
||||
(-lambda * days).exp()
|
||||
}
|
||||
|
||||
/// 便捷函数:使用默认 30 天半衰期
|
||||
pub fn decay_score_default(mtime_secs: u64) -> f64 {
|
||||
decay_score(mtime_secs, DEFAULT_HALF_LIFE_DAYS)
|
||||
}
|
||||
|
||||
/// Hebbian 激活分级 — 参考 OpenClaw Hot/Warm/Cool 模型。
|
||||
///
|
||||
/// 结合访问频率使用时更强大;
|
||||
/// 当前基于纯 recency 实现(无需额外追踪基础设施)。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ActivationTier {
|
||||
/// ≤7 天 — 高频访问,完整权重
|
||||
Hot,
|
||||
/// 8-30 天 — 中频,衰减中
|
||||
Warm,
|
||||
/// >30 天 — 低频,可能归档
|
||||
Cool,
|
||||
}
|
||||
|
||||
/// 按年龄分级
|
||||
pub fn activation_tier(mtime_secs: u64) -> ActivationTier {
|
||||
let days = super::age::memory_age_days(mtime_secs);
|
||||
if days <= 7 {
|
||||
ActivationTier::Hot
|
||||
} else if days <= 30 {
|
||||
ActivationTier::Warm
|
||||
} else {
|
||||
ActivationTier::Cool
|
||||
}
|
||||
}
|
||||
|
||||
/// 激活等级的排序权重乘数
|
||||
pub fn activation_multiplier(tier: ActivationTier) -> f64 {
|
||||
match tier {
|
||||
ActivationTier::Hot => 1.0,
|
||||
ActivationTier::Warm => 0.7,
|
||||
ActivationTier::Cool => 0.3,
|
||||
}
|
||||
}
|
||||
|
||||
/// 用于日志/显示的可读标签
|
||||
pub fn activation_label(tier: ActivationTier) -> &'static str {
|
||||
match tier {
|
||||
ActivationTier::Hot => "活跃",
|
||||
ActivationTier::Warm => "温",
|
||||
ActivationTier::Cool => "冷",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::memory::age;
|
||||
|
||||
#[test]
|
||||
fn test_decay_score_fresh() {
|
||||
let now = age::now_secs();
|
||||
let score = decay_score(now, 30.0);
|
||||
assert!((score - 1.0).abs() < 0.01, "今天应为 1.0,实际 {}", score);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decay_score_half_life() {
|
||||
let now = age::now_secs();
|
||||
let thirty_days_ago = now - 86_400 * 30;
|
||||
let score = decay_score(thirty_days_ago, 30.0);
|
||||
assert!((score - 0.5).abs() < 0.01, "30天应为 0.5,实际 {}", score);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decay_score_60_days() {
|
||||
let now = age::now_secs();
|
||||
let sixty_days_ago = now - 86_400 * 60;
|
||||
let score = decay_score(sixty_days_ago, 30.0);
|
||||
assert!(score < 0.26 && score > 0.24, "60天应≈0.25,实际 {}", score);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_activation_tier_hot() {
|
||||
let now = age::now_secs();
|
||||
assert_eq!(activation_tier(now), ActivationTier::Hot);
|
||||
assert_eq!(activation_tier(now - 86_400 * 7), ActivationTier::Hot);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_activation_tier_warm() {
|
||||
let now = age::now_secs();
|
||||
assert_eq!(activation_tier(now - 86_400 * 8), ActivationTier::Warm);
|
||||
assert_eq!(activation_tier(now - 86_400 * 30), ActivationTier::Warm);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_activation_tier_cool() {
|
||||
let now = age::now_secs();
|
||||
assert_eq!(activation_tier(now - 86_400 * 31), ActivationTier::Cool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_activation_multiplier_ranges() {
|
||||
assert_eq!(activation_multiplier(ActivationTier::Hot), 1.0);
|
||||
assert_eq!(activation_multiplier(ActivationTier::Warm), 0.7);
|
||||
assert_eq!(activation_multiplier(ActivationTier::Cool), 0.3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decay_score_different_half_life() {
|
||||
let now = age::now_secs();
|
||||
let ago = now - 86_400 * 15;
|
||||
// 15天半衰期下,15天后应为0.5
|
||||
let score = decay_score(ago, 15.0);
|
||||
assert!(
|
||||
(score - 0.5).abs() < 0.01,
|
||||
"15天半衰期15天后应为0.5,实际{}",
|
||||
score
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
// src/agent/memory/dedup.rs
|
||||
//
|
||||
// 记忆去重支持 — 参考 Claude Code memdir 提示词中的去重规则。
|
||||
//
|
||||
// 在保存新记忆前检查是否有可更新的现有条目,
|
||||
// 构建现有记忆的 manifest 供 LLM 参考以减少重复写入。
|
||||
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use super::types::MemoryEntry;
|
||||
|
||||
/// 构建现有记忆的清单预览(供 LLM 了解已存在的内容)。
|
||||
/// 在 save_memory 成功后注入到工具输出中。
|
||||
pub fn build_manifest_preview(entries: &[MemoryEntry]) -> String {
|
||||
if entries.is_empty() {
|
||||
return "当前无其他记忆条目。".to_string();
|
||||
}
|
||||
|
||||
let mut lines = vec!["当前记忆清单:".to_string()];
|
||||
for entry in entries {
|
||||
let type_label = match entry.memory_type {
|
||||
super::types::MemoryType::User => "[偏好]",
|
||||
super::types::MemoryType::Feedback => "[反馈]",
|
||||
super::types::MemoryType::Project => "[项目]",
|
||||
super::types::MemoryType::Reference => "[参考]",
|
||||
};
|
||||
lines.push(format!(
|
||||
"- {} `{}` {}: {}",
|
||||
type_label, entry.slug, entry.name, entry.description
|
||||
));
|
||||
}
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
/// 检查 slug 是否在磁盘上已存在。
|
||||
pub fn slug_exists(memory_dir: &Path, slug: &str) -> bool {
|
||||
let file_path = memory_dir.join(format!("{}.md", slug));
|
||||
file_path.exists()
|
||||
}
|
||||
|
||||
/// 列出所有现有 slug(从磁盘直接读取,避免依赖 MemoryManager 状态)。
|
||||
pub fn list_existing_slugs(memory_dir: &Path) -> Vec<String> {
|
||||
let mut slugs = Vec::new();
|
||||
if let Ok(entries) = fs::read_dir(memory_dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) {
|
||||
if file_name == "MEMORY.md" || !file_name.ends_with(".md") {
|
||||
continue;
|
||||
}
|
||||
if let Some(slug) = file_name.strip_suffix(".md") {
|
||||
slugs.push(slug.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
slugs.sort();
|
||||
slugs
|
||||
}
|
||||
|
||||
// ── Jaccard 相似度去重 ──
|
||||
|
||||
/// 计算两个字符串的 Jaccard 相似度(基于字符级 bigram)。
|
||||
///
|
||||
/// 使用 bigram 而非词级分词以正确处理中文(不依赖分词器)。
|
||||
/// 值域 [0.0, 1.0],阈值 ≥0.70 通常视为重复。
|
||||
///
|
||||
/// 参考 Martian-Engineering/agent-memory 的 70% Jaccard 门控。
|
||||
pub fn jaccard_similarity(a: &str, b: &str) -> f64 {
|
||||
let bigrams_a: std::collections::HashSet<String> = bigrams(a);
|
||||
let bigrams_b: std::collections::HashSet<String> = bigrams(b);
|
||||
|
||||
if bigrams_a.is_empty() && bigrams_b.is_empty() {
|
||||
return 1.0; // 两个空字符串完全相同
|
||||
}
|
||||
|
||||
let intersection = bigrams_a.intersection(&bigrams_b).count();
|
||||
let union = bigrams_a.union(&bigrams_b).count();
|
||||
|
||||
if union == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
intersection as f64 / union as f64
|
||||
}
|
||||
|
||||
/// 提取字符串的字符级 bigram 集合。
|
||||
fn bigrams(s: &str) -> std::collections::HashSet<String> {
|
||||
let chars: Vec<char> = s.chars().collect();
|
||||
let mut set = std::collections::HashSet::new();
|
||||
if chars.len() < 2 {
|
||||
// 单字符内容:将单字符本身作为 bigram
|
||||
if !chars.is_empty() {
|
||||
set.insert(chars[0].to_string());
|
||||
}
|
||||
return set;
|
||||
}
|
||||
for window in chars.windows(2) {
|
||||
set.insert(format!("{}{}", window[0], window[1]));
|
||||
}
|
||||
set
|
||||
}
|
||||
|
||||
/// 检查新内容与现有记忆是否高度重复。
|
||||
/// 返回重复的 slug,或在无重复时返回 None。
|
||||
pub fn find_duplicate_by_content(
|
||||
new_content: &str,
|
||||
existing_entries: &[MemoryEntry],
|
||||
threshold: f64,
|
||||
) -> Option<String> {
|
||||
for entry in existing_entries {
|
||||
if !entry.status.is_active() {
|
||||
continue;
|
||||
}
|
||||
let sim = jaccard_similarity(new_content, &entry.content);
|
||||
if sim >= threshold {
|
||||
return Some(entry.slug.clone());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// ── 写入时内容质量门控 ──
|
||||
|
||||
/// 内容质量检查结果
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum QualityCheck {
|
||||
/// 通过质量检查
|
||||
Accept,
|
||||
/// 太短:有效字符不足
|
||||
TooShort(usize),
|
||||
/// 瞬时状态描述
|
||||
TransientState,
|
||||
/// 模糊语言
|
||||
VagueLanguage(String),
|
||||
/// 纯代码片段
|
||||
CodePattern,
|
||||
}
|
||||
|
||||
/// 瞬时状态关键词(中文 + 英文)
|
||||
const TRANSIENT_PATTERNS: &[&str] = &[
|
||||
"正在做",
|
||||
"正在写",
|
||||
"正在调试",
|
||||
"正在看",
|
||||
"准备做",
|
||||
"is working on",
|
||||
"currently",
|
||||
"right now",
|
||||
"at the moment",
|
||||
];
|
||||
|
||||
/// 模糊语言关键词
|
||||
const VAGUE_PATTERNS: &[(&str, &str)] = &[
|
||||
("maybe", "可能"),
|
||||
("probably", "大概"),
|
||||
("perhaps", "也许"),
|
||||
("might be", "或许"),
|
||||
("似乎", "似乎"),
|
||||
("好像", "好像"),
|
||||
];
|
||||
|
||||
/// 代码模式检测(纯代码片段不应作为记忆)
|
||||
const CODE_PATTERNS: &[&str] = &[
|
||||
"fn ",
|
||||
"impl ",
|
||||
"struct ",
|
||||
"pub fn",
|
||||
"use crate",
|
||||
"function ",
|
||||
"const ",
|
||||
"let mut",
|
||||
"&mut",
|
||||
"import {",
|
||||
"from \"",
|
||||
"export ",
|
||||
];
|
||||
|
||||
/// 最小内容长度(有效字符)。
|
||||
/// 中文信息密度高,10 字即可表达完整语义。
|
||||
const MIN_CONTENT_CHARS: usize = 10;
|
||||
|
||||
/// 检查内容质量(写入时门控)。
|
||||
///
|
||||
/// 仅返回警告 — 不强制拒绝,由 LLM 最终决定。
|
||||
/// 参考 OpenClaw claw-mem 写入时门控 + agent-memory 写入规则。
|
||||
pub fn check_content_quality(content: &str) -> QualityCheck {
|
||||
let trimmed = content.trim();
|
||||
|
||||
// 1. 长度检查
|
||||
let char_count = trimmed.chars().count();
|
||||
if char_count < MIN_CONTENT_CHARS {
|
||||
return QualityCheck::TooShort(char_count);
|
||||
}
|
||||
|
||||
// 2. 瞬时状态检查
|
||||
let lower = trimmed.to_lowercase();
|
||||
for pattern in TRANSIENT_PATTERNS {
|
||||
if lower.contains(pattern) {
|
||||
return QualityCheck::TransientState;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 模糊语言检查
|
||||
for (en, zh) in VAGUE_PATTERNS {
|
||||
if lower.contains(en) || lower.contains(zh) {
|
||||
return QualityCheck::VagueLanguage(if lower.contains(en) {
|
||||
en.to_string()
|
||||
} else {
|
||||
zh.to_string()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 代码模式检查
|
||||
for pattern in CODE_PATTERNS {
|
||||
if trimmed.contains(pattern) {
|
||||
return QualityCheck::CodePattern;
|
||||
}
|
||||
}
|
||||
|
||||
QualityCheck::Accept
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn make_entry(slug: &str, name: &str, desc: &str) -> MemoryEntry {
|
||||
MemoryEntry {
|
||||
slug: slug.to_string(),
|
||||
name: name.to_string(),
|
||||
description: desc.to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: 1000,
|
||||
content: desc.to_string(),
|
||||
path: PathBuf::from(slug),
|
||||
status: super::super::types::MemoryStatus::Active,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_manifest_preview_empty() {
|
||||
let preview = build_manifest_preview(&[]);
|
||||
assert!(preview.contains("无其他记忆条目"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_manifest_preview_with_entries() {
|
||||
let entries = vec![
|
||||
make_entry("user-role", "用户角色", "数据科学家"),
|
||||
make_entry("feedback-tests", "测试反馈", "不要 mock 数据库"),
|
||||
];
|
||||
let preview = build_manifest_preview(&entries);
|
||||
assert!(preview.contains("user-role"));
|
||||
assert!(preview.contains("feedback-tests"));
|
||||
assert!(preview.contains("数据科学家"));
|
||||
assert!(preview.contains("不要 mock 数据库"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_slug_exists_true() {
|
||||
let dir = std::env::temp_dir().join("astro_memory_test_dedup");
|
||||
fs::create_dir_all(&dir).unwrap();
|
||||
fs::write(dir.join("existing.md"), "test").unwrap();
|
||||
assert!(slug_exists(&dir, "existing"));
|
||||
fs::remove_dir_all(&dir).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_slug_exists_false() {
|
||||
let dir = std::env::temp_dir().join("astro_memory_test_dedup_nonexist");
|
||||
assert!(!slug_exists(&dir, "nonexistent"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_existing_slugs() {
|
||||
let dir = std::env::temp_dir().join("astro_memory_test_list_slugs");
|
||||
fs::create_dir_all(&dir).unwrap();
|
||||
fs::write(dir.join("alpha.md"), "a").unwrap();
|
||||
fs::write(dir.join("beta.md"), "b").unwrap();
|
||||
fs::write(dir.join("MEMORY.md"), "index").unwrap();
|
||||
|
||||
let slugs = list_existing_slugs(&dir);
|
||||
assert!(slugs.contains(&"alpha".to_string()));
|
||||
assert!(slugs.contains(&"beta".to_string()));
|
||||
assert!(!slugs.contains(&"MEMORY".to_string()));
|
||||
|
||||
fs::remove_dir_all(&dir).unwrap();
|
||||
}
|
||||
|
||||
// ── Jaccard 相似度测试 ──
|
||||
|
||||
#[test]
|
||||
fn test_jaccard_identical() {
|
||||
let sim = jaccard_similarity("hello world", "hello world");
|
||||
assert!((sim - 1.0).abs() < 0.01, "完全相同应为 1.0,实际 {}", sim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jaccard_completely_different() {
|
||||
let sim = jaccard_similarity("hello world", "abc xyz");
|
||||
assert!(sim < 0.3, "完全不同应较低,实际 {}", sim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jaccard_high_overlap() {
|
||||
let sim = jaccard_similarity(
|
||||
"用户偏好使用 Rust 开发后端服务",
|
||||
"用户偏好使用 Rust 开发后端",
|
||||
);
|
||||
assert!(sim > 0.5, "高重叠应 >0.5,实际 {}", sim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jaccard_chinese_bigram() {
|
||||
let sim = jaccard_similarity("天体物理学研究", "天体物理研究");
|
||||
assert!(sim > 0.5, "中文 bigram 应能正确匹配,实际 {}", sim);
|
||||
}
|
||||
|
||||
// ── 内容质量检查测试 ──
|
||||
|
||||
#[test]
|
||||
fn test_quality_too_short() {
|
||||
assert_eq!(check_content_quality("太短"), QualityCheck::TooShort(2));
|
||||
// 刚好 10 个中文字符(可通过最低长度)
|
||||
let ten = "一二三四五六七八九十";
|
||||
assert_eq!(char_count(ten), 10);
|
||||
assert_eq!(check_content_quality(ten), QualityCheck::Accept);
|
||||
}
|
||||
|
||||
fn char_count(s: &str) -> usize {
|
||||
s.chars().count()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_transient_state() {
|
||||
assert_eq!(
|
||||
check_content_quality("用户正在调试登录模块的问题"),
|
||||
QualityCheck::TransientState
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_vague_language() {
|
||||
assert_eq!(
|
||||
check_content_quality("可能需要在后续版本中优化"),
|
||||
QualityCheck::VagueLanguage("可能".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_code_pattern() {
|
||||
assert_eq!(
|
||||
check_content_quality("fn main() { println!(\"hello\"); }"),
|
||||
QualityCheck::CodePattern
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_accept_good_content() {
|
||||
assert_eq!(
|
||||
check_content_quality(
|
||||
"用户是天体物理学家,主要研究星系演化。偏好使用 Kim 的径向速度拟合方法。"
|
||||
),
|
||||
QualityCheck::Accept
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_duplicate_by_content_detects_high_overlap() {
|
||||
let base = "用户偏好使用 Rust 开发后端服务";
|
||||
let entries = vec![make_entry("memory-a", "A", base)];
|
||||
let dup = find_duplicate_by_content("用户偏好使用 Rust 开发后端系统", &entries, 0.40);
|
||||
assert!(dup.is_some(), "高重叠内容应检测为重复");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_duplicate_rejects_low_overlap() {
|
||||
let entries = vec![make_entry("a", "A", "用户偏好使用 Rust 开发后端")];
|
||||
let dup = find_duplicate_by_content("天体物理学中星系演化研究的最新进展", &entries, 0.40);
|
||||
assert!(dup.is_none(), "低重叠内容不应检测为重复");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_duplicate_skips_historical() {
|
||||
let mut entries = vec![MemoryEntry {
|
||||
slug: "historical-one".to_string(),
|
||||
name: "历史记忆".to_string(),
|
||||
description: "已过时".to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: 1000,
|
||||
content: "用户偏好使用 Rust 开发后端".to_string(),
|
||||
path: PathBuf::from("historical-one.md"),
|
||||
status: super::super::types::MemoryStatus::Historical {
|
||||
superseded_by: Some("new-one".to_string()),
|
||||
},
|
||||
}];
|
||||
// historical 应被跳过,不匹配
|
||||
assert!(find_duplicate_by_content("用户偏好使用 Rust 开发后端", &entries, 0.6,).is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
// src/agent/memory/extraction.rs
|
||||
//
|
||||
// 自动记忆提取 — 参考 Claude Code services/extractMemories/。
|
||||
//
|
||||
// 在每次会话结束时,使用受限子代理分析对话内容并自动提取
|
||||
// 值得保留的记忆条目。提取是 fire-and-forget 的,不影响主会话关闭。
|
||||
//
|
||||
// 设计决策:
|
||||
// - 默认关闭(EXTRACT_MEMORY_ENABLED=false),避免意外的 LLM 费用
|
||||
// - 如果主代理已通过 save_memory 工具写入,则跳过提取
|
||||
// - 使用节流避免每轮都提取
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::agent::memory::dedup;
|
||||
use crate::agent::memory::MemoryManager;
|
||||
use crate::agent::subagent::SubAgentRunner;
|
||||
use crate::agent::tools::memory::SaveMemoryTool;
|
||||
use crate::agent::tools::{GlobFilesTool, GrepFilesTool, ReadFileTool, ToolRegistry};
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 自动提取配置(从环境变量加载)
|
||||
pub struct ExtractionConfig {
|
||||
/// 是否启用自动提取
|
||||
pub enabled: bool,
|
||||
/// 最小提取间隔(轮次)
|
||||
pub throttle_turns: usize,
|
||||
/// 子代理最大 ReAct 步数
|
||||
pub max_steps: usize,
|
||||
}
|
||||
|
||||
impl Default for ExtractionConfig {
|
||||
fn default() -> Self {
|
||||
ExtractionConfig {
|
||||
enabled: false,
|
||||
throttle_turns: 3,
|
||||
max_steps: 3,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ExtractionConfig {
|
||||
/// 从环境变量加载配置。
|
||||
/// - EXTRACT_MEMORY_ENABLED=true/false(默认 false)
|
||||
/// - EXTRACT_MEMORY_THROTTLE_TURNS(默认 3)
|
||||
/// - EXTRACT_MEMORY_MAX_STEPS(默认 3)
|
||||
pub fn from_env() -> Self {
|
||||
let enabled = std::env::var("EXTRACT_MEMORY_ENABLED")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<bool>().ok())
|
||||
.unwrap_or(false);
|
||||
|
||||
let throttle_turns = std::env::var("EXTRACT_MEMORY_THROTTLE_TURNS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
.unwrap_or(3);
|
||||
|
||||
let max_steps = std::env::var("EXTRACT_MEMORY_MAX_STEPS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
.unwrap_or(3);
|
||||
|
||||
ExtractionConfig {
|
||||
enabled,
|
||||
throttle_turns,
|
||||
max_steps,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 提取追踪器 — 存储在 MemoryManager 中以跨轮次追踪状态。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ExtractionTracker {
|
||||
/// 自上次提取以来的轮次数
|
||||
pub turns_since_last_extraction: usize,
|
||||
/// 主代理在本会话中是否已写入记忆
|
||||
pub main_agent_saved_this_session: bool,
|
||||
}
|
||||
|
||||
/// 提取系统提示词
|
||||
const EXTRACTION_SYSTEM_PROMPT: &str = "\
|
||||
你是一个记忆提取助手。分析最近的对话,提取值得持久化保存的信息。
|
||||
|
||||
## 记忆类型
|
||||
- **user**: 用户角色、偏好、知识背景
|
||||
- **feedback**: 用户给出的修正或确认的方法论(包含 Why 和 How to apply)
|
||||
- **project**: 项目上下文、目标、约束(不可从代码推导的部分)
|
||||
- **reference**: 外部资源指针(URL、仪表盘、工单系统)
|
||||
|
||||
## 不应保存
|
||||
- 代码模式、架构详情(可从项目状态推导)
|
||||
- Git 历史、调试方案
|
||||
- 已在 CLAUDE.md 中的内容
|
||||
- 临时任务状态";
|
||||
|
||||
/// 构建提取子代理的用户提示词。
|
||||
fn build_extraction_prompt(new_message_count: usize, existing_manifest: &str) -> String {
|
||||
let manifest_section = if existing_manifest.is_empty() {
|
||||
"当前无记忆条目。".to_string()
|
||||
} else {
|
||||
format!(
|
||||
"## 现有记忆清单\n\n{}\n\n检查此清单 — 更新现有文件而非创建重复项。",
|
||||
existing_manifest
|
||||
)
|
||||
};
|
||||
|
||||
format!(
|
||||
"分析最近约 {} 条消息,提取值得持久化保存的信息。\n\n{}\n\n\
|
||||
## 操作指南\n\
|
||||
1. 先读取需要更新的现有记忆文件(如果有)\n\
|
||||
2. 然后使用 save_memory 工具保存新记忆或更新现有记忆\n\
|
||||
3. 只保存非显而易见的、在后续对话中仍有用的信息\n\
|
||||
4. 不要浪费时间验证或搜索其他内容 — 仅基于对话内容",
|
||||
new_message_count, manifest_section
|
||||
)
|
||||
}
|
||||
|
||||
/// 构建受限工具注册表(只读 + save_memory)。
|
||||
fn build_extraction_tool_registry(memory_manager: Arc<Mutex<MemoryManager>>) -> ToolRegistry {
|
||||
let mut registry = ToolRegistry::empty();
|
||||
|
||||
// 只读工具
|
||||
registry.add_tool(Box::new(ReadFileTool));
|
||||
registry.add_tool(Box::new(GrepFilesTool));
|
||||
registry.add_tool(Box::new(GlobFilesTool));
|
||||
|
||||
// 写入仅限记忆目录
|
||||
registry.add_tool(Box::new(SaveMemoryTool::new(memory_manager)));
|
||||
|
||||
registry
|
||||
}
|
||||
|
||||
/// 运行自动记忆提取(fire-and-forget,调用者应通过 tokio::spawn 运行)。
|
||||
///
|
||||
/// 永不 panic — 所有错误都只记录日志。
|
||||
pub async fn run_extraction(
|
||||
app_state: Arc<AppState>,
|
||||
session_id: String,
|
||||
memory_manager: Arc<Mutex<MemoryManager>>,
|
||||
config: ExtractionConfig,
|
||||
) {
|
||||
if !config.enabled {
|
||||
return;
|
||||
}
|
||||
|
||||
// 检查节流和主代理写入
|
||||
{
|
||||
let mut mgr = memory_manager.lock().await;
|
||||
mgr.extraction_tracker.turns_since_last_extraction += 1;
|
||||
|
||||
if mgr.extraction_tracker.turns_since_last_extraction < config.throttle_turns {
|
||||
return;
|
||||
}
|
||||
if mgr.extraction_tracker.main_agent_saved_this_session {
|
||||
info!("[Extraction] 跳过 — 主代理已通过 save_memory 写入");
|
||||
mgr.extraction_tracker.main_agent_saved_this_session = false;
|
||||
mgr.extraction_tracker.turns_since_last_extraction = 0;
|
||||
return;
|
||||
}
|
||||
mgr.extraction_tracker.turns_since_last_extraction = 0;
|
||||
}
|
||||
|
||||
info!("[Extraction] 开始会话 {} 的自动记忆提取", session_id);
|
||||
|
||||
// 获取现有记忆清单
|
||||
let existing_manifest = {
|
||||
let mgr = memory_manager.lock().await;
|
||||
dedup::build_manifest_preview(mgr.entries())
|
||||
};
|
||||
|
||||
// 构建受限工具集
|
||||
let tool_registry = build_extraction_tool_registry(memory_manager.clone());
|
||||
|
||||
// 构建子代理
|
||||
let runner = SubAgentRunner::new_with_registry(app_state, tool_registry);
|
||||
|
||||
// 构建提示词
|
||||
let prompt = build_extraction_prompt(20, &existing_manifest);
|
||||
|
||||
// 运行子代理(同步等待,但调用者通过 tokio::spawn 异步化)
|
||||
let result = runner
|
||||
.run(EXTRACTION_SYSTEM_PROMPT, &prompt, config.max_steps)
|
||||
.await;
|
||||
|
||||
if result.is_error {
|
||||
warn!(
|
||||
"[Extraction] 子代理返回错误: {}",
|
||||
result.content.chars().take(200).collect::<String>()
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"[Extraction] 提取完成: {}",
|
||||
result.content.chars().take(150).collect::<String>()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extraction_config_default() {
|
||||
let config = ExtractionConfig::default();
|
||||
assert!(!config.enabled);
|
||||
assert_eq!(config.throttle_turns, 3);
|
||||
assert_eq!(config.max_steps, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extraction_config_from_env_disabled() {
|
||||
// 未设置环境变量时应为默认(禁用)
|
||||
std::env::remove_var("EXTRACT_MEMORY_ENABLED");
|
||||
let config = ExtractionConfig::from_env();
|
||||
assert!(!config.enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extraction_tracker_default() {
|
||||
let tracker = ExtractionTracker::default();
|
||||
assert_eq!(tracker.turns_since_last_extraction, 0);
|
||||
assert!(!tracker.main_agent_saved_this_session);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_extraction_prompt() {
|
||||
let prompt = build_extraction_prompt(10, "");
|
||||
assert!(prompt.contains("10 条消息"));
|
||||
assert!(prompt.contains("当前无记忆条目"));
|
||||
|
||||
let prompt_with_manifest = build_extraction_prompt(5, "- user-role: 用户角色");
|
||||
assert!(prompt_with_manifest.contains("现有记忆清单"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
// src/agent/memory/guardrails.rs
|
||||
//
|
||||
// 记忆保存/使用护栏 — 参考 Claude Code memoryTypes.ts。
|
||||
//
|
||||
// 提供两个维度的防护:
|
||||
// 1. WHAT_NOT_TO_SAVE — 不应保存为记忆的内容(即使被要求)
|
||||
// 2. TRUST_BUT_VERIFY — 从记忆中推荐前先核实的提示词
|
||||
//
|
||||
// 这些提示词经过 Claude Code 评估验证(memory-prompt-iteration.eval.ts):
|
||||
// - 排除规则明确告知模型 "即使用户要求保存" 也不应保存噪音内容
|
||||
// - 验证提示词需要放在决策点(系统提示词中记忆段落之后),
|
||||
// 不能埋在通用指南中,否则模型会忽略
|
||||
|
||||
/// 不应保存为记忆的内容。
|
||||
/// 即使用户明确要求保存,这些规则也适用。
|
||||
pub const WHAT_NOT_TO_SAVE: &str = "\
|
||||
不应保存为记忆的内容:
|
||||
- 代码模式、惯例、架构详情、文件路径 — 可从当前项目状态推导
|
||||
- Git 历史、最近修改、谁改了什么 — `git log` / `git blame` 是权威来源
|
||||
- 调试方案或错误临时解决方案 — 修复在代码中,commit message 有上下文
|
||||
- 已在 CLAUDE.md 或项目文档中的内容
|
||||
- 临时任务细节:进行中工作、当前对话上下文
|
||||
|
||||
即使用户明确要求保存以上内容,请询问其中哪些部分是*意外的*或*非常规的* — 那些才是值得保存的。";
|
||||
|
||||
/// "从记忆中推荐前先核实" 提示词。
|
||||
/// 必须放在决策点(记忆段落后),不能在通用指南中。
|
||||
/// Claude Code 评估:放在 "When to access memories" 下时 0/3,
|
||||
/// 放在独立段落标题下时 3/3 — 标题权重影响模型行为。
|
||||
pub const VERIFY_BEFORE_RECOMMENDING: &str = "\
|
||||
## 从记忆中推荐前先核实
|
||||
|
||||
记忆中提到特定函数、文件或标志是一种声明,声称它们*在记忆写入时*存在。
|
||||
但函数可能已被重命名、移除或从未合并。在据此推荐前:
|
||||
|
||||
- 如果记忆提到了文件路径:确认该文件存在
|
||||
- 如果记忆提到了函数或标志:用 grep 搜索
|
||||
- 如果用户将基于你的推荐采取行动(不仅是询问历史),先核实
|
||||
|
||||
\"记忆说 X 存在\" 不等于 \"X 现在存在\"。";
|
||||
|
||||
/// 构建注入到 system prompt 的验证提醒。
|
||||
/// 放在记忆条目之后、`</project-memory-context>` 之前。
|
||||
pub fn build_verification_reminder() -> String {
|
||||
VERIFY_BEFORE_RECOMMENDING.to_string()
|
||||
}
|
||||
|
||||
/// 构建保存记忆的排除规则提示(供工具描述使用)。
|
||||
pub fn build_exclusion_reminder() -> String {
|
||||
WHAT_NOT_TO_SAVE.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_what_not_to_save_is_non_empty() {
|
||||
assert!(!WHAT_NOT_TO_SAVE.is_empty());
|
||||
// 关键短语验证
|
||||
assert!(WHAT_NOT_TO_SAVE.contains("代码模式"));
|
||||
assert!(WHAT_NOT_TO_SAVE.contains("即使用户明确要求"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_before_recommending_is_non_empty() {
|
||||
assert!(!VERIFY_BEFORE_RECOMMENDING.is_empty());
|
||||
assert!(VERIFY_BEFORE_RECOMMENDING.contains("从记忆中推荐前先核实"));
|
||||
assert!(VERIFY_BEFORE_RECOMMENDING.contains("记忆说 X 存在"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_verification_reminder() {
|
||||
let reminder = build_verification_reminder();
|
||||
assert_eq!(reminder, VERIFY_BEFORE_RECOMMENDING);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_exclusion_reminder() {
|
||||
let reminder = build_exclusion_reminder();
|
||||
assert!(reminder.contains("不应保存"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,420 @@
|
||||
// src/agent/memory/mod.rs
|
||||
//
|
||||
// 项目记忆管理器。
|
||||
// 参考 Claude Code memdir 设计。
|
||||
//
|
||||
// 在 {library_dir}/memory/ 目录下维护:
|
||||
// - MEMORY.md — 索引文件(最多 200 行,25KB)
|
||||
// - {slug}.md — 每个记忆一个文件,YAML frontmatter + Markdown 内容
|
||||
//
|
||||
// 自动在 Agent 的 system prompt 中注入最近的记忆条目。
|
||||
// 提供 save_memory 工具供 Agent 写入记忆。
|
||||
|
||||
pub mod age;
|
||||
pub mod decay;
|
||||
pub mod dedup;
|
||||
pub mod extraction;
|
||||
pub mod guardrails;
|
||||
pub mod selection;
|
||||
pub mod types;
|
||||
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use self::types::{entry_from_frontmatter, parse_frontmatter, MemoryEntry, MemoryType};
|
||||
|
||||
/// 索引文件最大行数
|
||||
const MAX_ENTRYPOINT_LINES: usize = 200;
|
||||
/// 索引文件最大字节数
|
||||
const MAX_ENTRYPOINT_BYTES: usize = 25_000;
|
||||
|
||||
/// 记忆管理器
|
||||
pub struct MemoryManager {
|
||||
/// 记忆目录
|
||||
memory_dir: PathBuf,
|
||||
/// 已加载的记忆条目
|
||||
entries: Vec<MemoryEntry>,
|
||||
/// 自动提取追踪状态
|
||||
pub extraction_tracker: extraction::ExtractionTracker,
|
||||
}
|
||||
|
||||
impl MemoryManager {
|
||||
/// 创建并加载记忆。
|
||||
/// `library_dir` 是项目配置中的 library 目录。
|
||||
pub fn new(library_dir: PathBuf) -> Self {
|
||||
let memory_dir = library_dir.join("memory");
|
||||
let mut manager = MemoryManager {
|
||||
memory_dir,
|
||||
entries: Vec::new(),
|
||||
extraction_tracker: extraction::ExtractionTracker::default(),
|
||||
};
|
||||
manager.reload();
|
||||
manager
|
||||
}
|
||||
|
||||
/// 重新从磁盘加载所有记忆。
|
||||
pub fn reload(&mut self) {
|
||||
// 确保目录存在
|
||||
if let Err(e) = fs::create_dir_all(&self.memory_dir) {
|
||||
warn!("[Memory] 无法创建记忆目录 {:?}: {}", self.memory_dir, e);
|
||||
return;
|
||||
}
|
||||
|
||||
self.entries.clear();
|
||||
|
||||
// 扫描 .md 文件(排除 MEMORY.md 和目录)
|
||||
match fs::read_dir(&self.memory_dir) {
|
||||
Ok(dir_entries) => {
|
||||
for entry in dir_entries {
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if file_name == "MEMORY.md" || !file_name.ends_with(".md") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let slug = file_name.strip_suffix(".md").unwrap_or(file_name);
|
||||
let mtime = entry
|
||||
.metadata()
|
||||
.ok()
|
||||
.and_then(|m| m.modified().ok())
|
||||
.map(|t| {
|
||||
t.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
|
||||
match fs::read_to_string(&path) {
|
||||
Ok(raw) => {
|
||||
let (fields, content) = parse_frontmatter(&raw);
|
||||
if let Some(mem_entry) =
|
||||
entry_from_frontmatter(slug, &fields, &content, path.clone(), mtime)
|
||||
{
|
||||
self.entries.push(mem_entry);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[Memory] 无法读取 {:?}: {}", path, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[Memory] 无法扫描记忆目录: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// 按修改时间排序(最新在前)
|
||||
self.entries.sort_by(|a, b| b.mtime.cmp(&a.mtime));
|
||||
|
||||
info!(
|
||||
"[Memory] 加载了 {} 条记忆从 {:?}",
|
||||
self.entries.len(),
|
||||
self.memory_dir
|
||||
);
|
||||
}
|
||||
|
||||
/// 获取所有记忆条目
|
||||
pub fn entries(&self) -> &[MemoryEntry] {
|
||||
&self.entries
|
||||
}
|
||||
|
||||
/// 获取记忆目录路径
|
||||
pub fn memory_dir(&self) -> &std::path::Path {
|
||||
&self.memory_dir
|
||||
}
|
||||
|
||||
/// 标记主代理在本会话中已写入记忆(抑制自动提取)。
|
||||
pub fn mark_main_agent_wrote(&mut self) {
|
||||
self.extraction_tracker.main_agent_saved_this_session = true;
|
||||
}
|
||||
|
||||
/// 语义匹配:委托给 selection 模块使用 LLM 结构化选择。
|
||||
///
|
||||
/// 选择后应用指数时间衰减排序(更近的 active 记忆获得更高权重)。
|
||||
/// 失败时回退到 recency-based 选择(跳过已展示的条目)。
|
||||
pub async fn select_relevant_memories(
|
||||
&self,
|
||||
llm: &crate::clients::llm::LlmClient,
|
||||
context: &str,
|
||||
max_entries: usize,
|
||||
) -> Vec<&MemoryEntry> {
|
||||
if self.entries.len() <= max_entries {
|
||||
return self.entries.iter().collect();
|
||||
}
|
||||
|
||||
let sel_ctx = selection::SelectionContext {
|
||||
max_entries,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let indices = selection::select_structured(&self.entries, llm, context, &sel_ctx).await;
|
||||
|
||||
// 应用指数时间衰减排序(更近的 active 记忆在前,historical 在后)
|
||||
let sorted =
|
||||
selection::apply_decay_scoring(&indices, &self.entries, decay::DEFAULT_HALF_LIFE_DAYS);
|
||||
|
||||
sorted.iter().filter_map(|&i| self.entries.get(i)).collect()
|
||||
}
|
||||
|
||||
/// 从指定条目列表构建 system reminder(而非全部条目)
|
||||
pub fn build_system_reminder_from(&self, selected: &[&MemoryEntry]) -> Option<String> {
|
||||
if selected.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut lines = vec![
|
||||
"<project-memory-context>".to_string(),
|
||||
String::new(),
|
||||
"[PROJECT MEMORY]".to_string(),
|
||||
String::new(),
|
||||
];
|
||||
|
||||
for entry in selected {
|
||||
let type_tag = match entry.memory_type {
|
||||
MemoryType::User => "[偏好]",
|
||||
MemoryType::Feedback => "[反馈]",
|
||||
MemoryType::Project => "[项目]",
|
||||
MemoryType::Reference => "[参考]",
|
||||
};
|
||||
// historical 记忆标记
|
||||
let status_tag = if !entry.status.is_active() {
|
||||
match entry.status.superseded_by() {
|
||||
Some(new_slug) => format!(" [已更新→{}]", new_slug),
|
||||
None => " [已更新]".to_string(),
|
||||
}
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
let preview: String = entry
|
||||
.content
|
||||
.lines()
|
||||
.take(3)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n ");
|
||||
lines.push(format!(
|
||||
"{} {}{}: {}\n {}",
|
||||
type_tag, entry.name, status_tag, entry.description, preview
|
||||
));
|
||||
// 注入时效警告(超过1天的记忆)
|
||||
let freshness = age::memory_freshness_note(entry.mtime);
|
||||
if !freshness.is_empty() {
|
||||
lines.push(freshness);
|
||||
}
|
||||
}
|
||||
|
||||
lines.push(String::new());
|
||||
lines.push(
|
||||
"使用 save_memory 工具保存重要信息。记忆内容可能过时,请在使用前验证。".to_string(),
|
||||
);
|
||||
// 注入验证提醒(从记忆推荐前先核实)
|
||||
lines.push(String::new());
|
||||
lines.push(guardrails::build_verification_reminder());
|
||||
lines.push("</project-memory-context>".to_string());
|
||||
|
||||
Some(lines.join("\n"))
|
||||
}
|
||||
|
||||
/// 保存一条新的记忆。
|
||||
///
|
||||
/// 如果 slug 已存在且内容为 Active,旧文件归档为 `{slug}_v1.md`
|
||||
/// 并将状态标记为 historical(永不删除旧记忆)。
|
||||
pub fn save_memory(
|
||||
&mut self,
|
||||
slug: &str,
|
||||
name: &str,
|
||||
description: &str,
|
||||
memory_type: MemoryType,
|
||||
content: &str,
|
||||
) -> std::io::Result<()> {
|
||||
let file_path = self.memory_dir.join(format!("{}.md", slug));
|
||||
|
||||
// 如果已有活跃版本,归档旧版本
|
||||
let old_path = self.memory_dir.join(format!("{}_v1.md", slug));
|
||||
if file_path.exists() {
|
||||
if let Ok(old_content) = fs::read_to_string(&file_path) {
|
||||
fs::write(&old_path, &old_content)?;
|
||||
info!("[Memory] 归档旧版本: {} → {}", slug, old_path.display());
|
||||
}
|
||||
}
|
||||
|
||||
let frontmatter = format!(
|
||||
"---\nname: {}\ndescription: {}\ntype: {}\nstatus: active\n---\n",
|
||||
name,
|
||||
description,
|
||||
memory_type.as_str()
|
||||
);
|
||||
let full_content = format!("{}{}", frontmatter, content);
|
||||
|
||||
fs::write(&file_path, &full_content)?;
|
||||
|
||||
// 更新 MEMORY.md 索引
|
||||
self.update_index(slug, name, description, &memory_type)?;
|
||||
|
||||
// 重新加载
|
||||
self.reload();
|
||||
|
||||
info!("[Memory] 已保存记忆: {} ({})", name, slug);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 更新 MEMORY.md 索引文件。
|
||||
fn update_index(
|
||||
&self,
|
||||
slug: &str,
|
||||
name: &str,
|
||||
description: &str,
|
||||
memory_type: &MemoryType,
|
||||
) -> std::io::Result<()> {
|
||||
let index_path = self.memory_dir.join("MEMORY.md");
|
||||
let line = format!(
|
||||
"- [{}]({}.md) — {} (type: {})",
|
||||
name,
|
||||
slug,
|
||||
description,
|
||||
memory_type.as_str()
|
||||
);
|
||||
|
||||
let mut content = if index_path.exists() {
|
||||
let existing = fs::read_to_string(&index_path).unwrap_or_default();
|
||||
// 检查是否已有此 slug 的条目
|
||||
let slug_marker = format!("]({}.md)", slug);
|
||||
let lines: Vec<&str> = existing.lines().collect();
|
||||
|
||||
// 行数检查
|
||||
if lines.len() >= MAX_ENTRYPOINT_LINES {
|
||||
// 移除最旧的行(索引头部保持不变)
|
||||
warn!(
|
||||
"[Memory] MEMORY.md 行数已满 ({}), 移除最旧条目",
|
||||
lines.len()
|
||||
);
|
||||
let keep = MAX_ENTRYPOINT_LINES - 1;
|
||||
format!("{}\n{}", lines[..keep.min(lines.len())].join("\n"), line)
|
||||
} else {
|
||||
// 检查是否需要替换已存在的条目
|
||||
let has_entry = lines.iter().any(|l| l.contains(&slug_marker));
|
||||
if has_entry {
|
||||
// 替换已存在的行
|
||||
lines
|
||||
.iter()
|
||||
.map(|l| {
|
||||
if l.contains(&slug_marker) {
|
||||
line.as_str()
|
||||
} else {
|
||||
*l
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
} else {
|
||||
format!("{}\n{}", existing, line)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
format!("# Project Memory\n\n{}", line)
|
||||
};
|
||||
|
||||
// 字节数检查(在大约 25KB 处截断)
|
||||
if content.len() > MAX_ENTRYPOINT_BYTES {
|
||||
let truncated: String = content
|
||||
.char_indices()
|
||||
.take_while(|(i, _)| *i < MAX_ENTRYPOINT_BYTES - 100)
|
||||
.map(|(_, c)| c)
|
||||
.collect();
|
||||
content = format!(
|
||||
"{}\n\n[MEMORY.md 已达到 {}KB 上限,旧条目已截断]",
|
||||
truncated,
|
||||
MAX_ENTRYPOINT_BYTES / 1024
|
||||
);
|
||||
}
|
||||
|
||||
fs::write(&index_path, &content)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 生成 system prompt 中注入的记忆段落。
|
||||
///
|
||||
/// 包含最近的记忆条目(最多 10 条),并在前面注明可信度提醒。
|
||||
pub fn build_system_reminder(&self, max_entries: usize) -> Option<String> {
|
||||
if self.entries.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut lines = vec![
|
||||
"<project-memory-context>".to_string(),
|
||||
"".to_string(),
|
||||
"[PROJECT MEMORY]".to_string(),
|
||||
"".to_string(),
|
||||
];
|
||||
|
||||
let count = max_entries.min(self.entries.len());
|
||||
for entry in self.entries.iter().take(count) {
|
||||
let type_tag = match entry.memory_type {
|
||||
MemoryType::User => "[偏好]",
|
||||
MemoryType::Feedback => "[反馈]",
|
||||
MemoryType::Project => "[项目]",
|
||||
MemoryType::Reference => "[参考]",
|
||||
};
|
||||
// historical 记忆标记
|
||||
let status_tag = if !entry.status.is_active() {
|
||||
match entry.status.superseded_by() {
|
||||
Some(new_slug) => format!(" [已更新→{}]", new_slug),
|
||||
None => " [已更新]".to_string(),
|
||||
}
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
let preview: String = entry
|
||||
.content
|
||||
.lines()
|
||||
.take(3)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n ");
|
||||
lines.push(format!(
|
||||
"{} {}{}: {}\n {}",
|
||||
type_tag, entry.name, status_tag, entry.description, preview
|
||||
));
|
||||
// 注入时效警告(超过1天的记忆)
|
||||
let freshness = age::memory_freshness_note(entry.mtime);
|
||||
if !freshness.is_empty() {
|
||||
lines.push(freshness);
|
||||
}
|
||||
}
|
||||
|
||||
lines.push("".to_string());
|
||||
lines.push(
|
||||
"使用 save_memory 工具保存重要信息。记忆内容可能过时,请在使用前验证。".to_string(),
|
||||
);
|
||||
// 注入验证提醒(从记忆推荐前先核实)
|
||||
lines.push(String::new());
|
||||
lines.push(guardrails::build_verification_reminder());
|
||||
lines.push("</project-memory-context>".to_string());
|
||||
|
||||
Some(lines.join("\n"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_build_empty_reminder() {
|
||||
let manager = MemoryManager {
|
||||
memory_dir: PathBuf::from("/tmp/nonexistent"),
|
||||
entries: Vec::new(),
|
||||
extraction_tracker: extraction::ExtractionTracker::default(),
|
||||
};
|
||||
assert!(manager.build_system_reminder(10).is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,351 @@
|
||||
// src/agent/memory/selection.rs
|
||||
//
|
||||
// 改进的记忆相关性选择 — 参考 Claude Code findRelevantMemories.ts。
|
||||
//
|
||||
// 相比旧实现 (mod.rs 中直接调用 chat_completion + regex 解析):
|
||||
// 1. 使用结构化 JSON 提示词 + 更强健的解析
|
||||
// 2. 支持 SelectionContext 跟踪已展示的记忆(避免重复选择)
|
||||
// 3. 支持工具感知过滤(排除最近使用工具相关的记忆)
|
||||
// 4. 失败时优雅降级到 recency 回退
|
||||
|
||||
use crate::clients::llm::LlmClient;
|
||||
use tracing::info;
|
||||
|
||||
use super::types::MemoryEntry;
|
||||
|
||||
/// 相关性选择结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SelectionContext {
|
||||
/// 已在前几轮展示过的记忆索引(避免重复选择)
|
||||
pub already_surfaced: Vec<usize>,
|
||||
/// 最近使用的工具名(关于这些工具的记忆降权)
|
||||
pub recent_tools: Vec<String>,
|
||||
/// 最大返回条目数
|
||||
pub max_entries: usize,
|
||||
}
|
||||
|
||||
impl Default for SelectionContext {
|
||||
fn default() -> Self {
|
||||
SelectionContext {
|
||||
already_surfaced: Vec::new(),
|
||||
recent_tools: Vec::new(),
|
||||
max_entries: 5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用 LLM 从候选记忆中选出最相关的。
|
||||
///
|
||||
/// 先尝试 LLM 结构化选择,失败时回退到 recency-based 选择。
|
||||
/// `sel_ctx.already_surfaced` 中的条目会被排除在候选之外。
|
||||
pub async fn select_structured(
|
||||
entries: &[MemoryEntry],
|
||||
llm: &LlmClient,
|
||||
context: &str,
|
||||
sel_ctx: &SelectionContext,
|
||||
) -> Vec<usize> {
|
||||
// 过滤已展示的条目
|
||||
let candidates: Vec<(usize, &MemoryEntry)> = entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| !sel_ctx.already_surfaced.contains(i))
|
||||
.collect();
|
||||
|
||||
if candidates.is_empty() {
|
||||
return fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced);
|
||||
}
|
||||
|
||||
if candidates.len() <= sel_ctx.max_entries {
|
||||
return candidates.iter().map(|(i, _)| *i).collect();
|
||||
}
|
||||
|
||||
// 构建候选目录
|
||||
let catalog: Vec<String> = candidates
|
||||
.iter()
|
||||
.map(|(_, e)| {
|
||||
format!(
|
||||
"[{}] [{}] {}: {}",
|
||||
e.slug,
|
||||
e.memory_type.as_str(),
|
||||
e.name,
|
||||
e.description
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 工具提示(可选)
|
||||
let tools_hint = if sel_ctx.recent_tools.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(
|
||||
"\n\n近期使用的工具: {}。不要选择这些工具的使用参考或 API 文档类记忆。",
|
||||
sel_ctx.recent_tools.join(", ")
|
||||
)
|
||||
};
|
||||
|
||||
let prompt = format!(
|
||||
"用户当前话题:\n{}\n\n从以下记忆目录中选择最多 {} 条最相关的。\
|
||||
仅选择明确有帮助的,不确定则不选。返回 JSON 数组如 [\"slug1\", \"slug2\"]。\n\n{}\n{}",
|
||||
context,
|
||||
sel_ctx.max_entries,
|
||||
catalog.join("\n"),
|
||||
tools_hint,
|
||||
);
|
||||
|
||||
match llm
|
||||
.chat_completion(
|
||||
"你是一个记忆检索助手。根据用户话题选择最相关的记忆。仅返回 JSON 字符串数组。",
|
||||
&prompt,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => match extract_slugs(&response, &candidates) {
|
||||
Ok(selected) => {
|
||||
info!(
|
||||
"[Memory] LLM 选择: {}/{} 条相关记忆",
|
||||
selected.len(),
|
||||
candidates.len()
|
||||
);
|
||||
selected
|
||||
}
|
||||
Err(_) => {
|
||||
info!("[Memory] JSON 解析失败,回退到 recency");
|
||||
fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced)
|
||||
}
|
||||
},
|
||||
Err(_) => fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 LLM 响应中提取 slug 列表。
|
||||
/// 尝试两种格式:["slug1","slug2"] 或 [0, 1, 3](数字索引回退)
|
||||
fn extract_slugs(response: &str, candidates: &[(usize, &MemoryEntry)]) -> Result<Vec<usize>, ()> {
|
||||
// 方法1: 查找 JSON 字符串数组
|
||||
if let Some(start) = response.find('[') {
|
||||
if let Some(end) = response.rfind(']') {
|
||||
let json_str = &response[start..=end];
|
||||
|
||||
// 尝试解析为字符串数组
|
||||
if let Ok(slugs) = serde_json::from_str::<Vec<String>>(json_str) {
|
||||
let indices: Vec<usize> = slugs
|
||||
.iter()
|
||||
.filter_map(|s| {
|
||||
candidates
|
||||
.iter()
|
||||
.find(|(_, e)| e.slug == *s)
|
||||
.map(|(i, _)| *i)
|
||||
})
|
||||
.collect();
|
||||
if !indices.is_empty() {
|
||||
return Ok(indices);
|
||||
}
|
||||
}
|
||||
|
||||
// 方法2: 回退到数字索引
|
||||
if let Ok(indices) = serde_json::from_str::<Vec<usize>>(json_str) {
|
||||
let valid: Vec<usize> = indices
|
||||
.into_iter()
|
||||
.filter(|i| candidates.iter().any(|(idx, _)| *idx == *i))
|
||||
.collect();
|
||||
if !valid.is_empty() {
|
||||
return Ok(valid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(())
|
||||
}
|
||||
|
||||
/// 将指数时间衰减评分应用于已选中的记忆,按最终得分降序排序。
|
||||
///
|
||||
/// 每项最终得分 = decay_score(mtime) × 1.0(LLM 选中即默认置信度)。
|
||||
/// 仅影响排序顺序,不减少选中的数量。
|
||||
/// 仅活跃记忆参与衰减排序;historical 记忆保持原顺序。
|
||||
pub fn apply_decay_scoring(
|
||||
indices: &[usize],
|
||||
entries: &[MemoryEntry],
|
||||
half_life_days: f64,
|
||||
) -> Vec<usize> {
|
||||
use super::decay::decay_score;
|
||||
|
||||
let mut scored: Vec<(usize, f64)> = indices
|
||||
.iter()
|
||||
.filter_map(|&i| {
|
||||
entries.get(i).map(|entry| {
|
||||
if entry.status.is_active() {
|
||||
(i, decay_score(entry.mtime, half_life_days))
|
||||
} else {
|
||||
// historical 记忆赋予最低分,排在最后
|
||||
(i, 0.01)
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 按得分降序排列(高分在前)
|
||||
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
scored.into_iter().map(|(i, _)| i).collect()
|
||||
}
|
||||
|
||||
/// Recency 回退:取最近的 max_entries 条(排除已展示的)
|
||||
fn fallback_recency(
|
||||
entries: &[MemoryEntry],
|
||||
max_entries: usize,
|
||||
already_surfaced: &[usize],
|
||||
) -> Vec<usize> {
|
||||
entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| !already_surfaced.contains(i))
|
||||
.take(max_entries)
|
||||
.map(|(i, _)| i)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_entry(slug: &str, name: &str, desc: &str) -> MemoryEntry {
|
||||
MemoryEntry {
|
||||
slug: slug.to_string(),
|
||||
name: name.to_string(),
|
||||
description: desc.to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: 1000,
|
||||
content: String::new(),
|
||||
path: std::path::PathBuf::from(slug),
|
||||
status: super::super::types::MemoryStatus::Active,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fallback_recency_basic() {
|
||||
let entries = vec![
|
||||
make_entry("a", "A", "desc a"),
|
||||
make_entry("b", "B", "desc b"),
|
||||
make_entry("c", "C", "desc c"),
|
||||
];
|
||||
let result = fallback_recency(&entries, 2, &[]);
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0], 0); // 最近的在前
|
||||
assert_eq!(result[1], 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fallback_skips_already_surfaced() {
|
||||
let entries = vec![
|
||||
make_entry("a", "A", "desc a"),
|
||||
make_entry("b", "B", "desc b"),
|
||||
make_entry("c", "C", "desc c"),
|
||||
];
|
||||
let result = fallback_recency(&entries, 3, &[0]); // skip index 0
|
||||
assert_eq!(result, vec![1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_slugs_string_array() {
|
||||
let entries = vec![
|
||||
make_entry("alpha", "Alpha", "first"),
|
||||
make_entry("beta", "Beta", "second"),
|
||||
make_entry("gamma", "Gamma", "third"),
|
||||
];
|
||||
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
|
||||
|
||||
let result = extract_slugs(r#"["alpha", "gamma"]"#, &candidates);
|
||||
assert_eq!(result, Ok(vec![0, 2]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_slugs_numeric_fallback() {
|
||||
let entries = vec![make_entry("x", "X", "x"), make_entry("y", "Y", "y")];
|
||||
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
|
||||
|
||||
let result = extract_slugs("[0]", &candidates);
|
||||
assert_eq!(result, Ok(vec![0]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_slugs_invalid_json() {
|
||||
let entries = vec![make_entry("x", "X", "x")];
|
||||
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
|
||||
|
||||
assert!(extract_slugs("not json at all", &candidates).is_err());
|
||||
assert!(extract_slugs("no brackets here", &candidates).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_selection_context_default() {
|
||||
let ctx = SelectionContext::default();
|
||||
assert!(ctx.already_surfaced.is_empty());
|
||||
assert!(ctx.recent_tools.is_empty());
|
||||
assert_eq!(ctx.max_entries, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apply_decay_scoring_sorts_by_freshness() {
|
||||
use crate::agent::memory::age;
|
||||
let now = age::now_secs();
|
||||
let entries = vec![
|
||||
MemoryEntry {
|
||||
slug: "old".to_string(),
|
||||
name: "旧记忆".to_string(),
|
||||
description: "old".to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: now - 86_400 * 60, // 60天前
|
||||
content: String::new(),
|
||||
path: std::path::PathBuf::from("old.md"),
|
||||
status: super::super::types::MemoryStatus::Active,
|
||||
},
|
||||
MemoryEntry {
|
||||
slug: "fresh".to_string(),
|
||||
name: "新鲜记忆".to_string(),
|
||||
description: "fresh".to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: now, // 今天
|
||||
content: String::new(),
|
||||
path: std::path::PathBuf::from("fresh.md"),
|
||||
status: super::super::types::MemoryStatus::Active,
|
||||
},
|
||||
];
|
||||
// 旧记忆在前(index 0),新记忆在后(index 1)
|
||||
let sorted = apply_decay_scoring(&[0, 1], &entries, 30.0);
|
||||
// 新记忆应排在旧记忆前面
|
||||
assert_eq!(sorted[0], 1);
|
||||
assert_eq!(sorted[1], 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apply_decay_puts_historical_last() {
|
||||
use crate::agent::memory::age;
|
||||
let now = age::now_secs();
|
||||
let entries = vec![
|
||||
MemoryEntry {
|
||||
slug: "active".to_string(),
|
||||
name: "活跃".to_string(),
|
||||
description: "active".to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: now,
|
||||
content: String::new(),
|
||||
path: std::path::PathBuf::from("active.md"),
|
||||
status: super::super::types::MemoryStatus::Active,
|
||||
},
|
||||
MemoryEntry {
|
||||
slug: "historical".to_string(),
|
||||
name: "历史".to_string(),
|
||||
description: "historical".to_string(),
|
||||
memory_type: super::super::types::MemoryType::User,
|
||||
mtime: now, // 也很新,但是 historical
|
||||
content: String::new(),
|
||||
path: std::path::PathBuf::from("historical.md"),
|
||||
status: super::super::types::MemoryStatus::Historical {
|
||||
superseded_by: Some("active".to_string()),
|
||||
},
|
||||
},
|
||||
];
|
||||
let sorted = apply_decay_scoring(&[0, 1], &entries, 30.0);
|
||||
// historical 应排在最后
|
||||
assert_eq!(sorted[0], 0, "活跃记忆应在前");
|
||||
assert_eq!(sorted[1], 1, "historical 应在后");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
// src/agent/memory/types.rs
|
||||
//
|
||||
// 项目记忆系统 — 类型定义。
|
||||
// 参考 Claude Code memdir/memoryTypes.ts 设计。
|
||||
//
|
||||
// 四种记忆类型:
|
||||
// - User — 用户角色、偏好、目标
|
||||
// - Feedback — 用户提供的反馈(修正 + 确认)
|
||||
// - Project — 项目状态、进行中的工作、目标
|
||||
// - Reference — 外部资源的指针
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// 记忆类型
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum MemoryType {
|
||||
User,
|
||||
Feedback,
|
||||
Project,
|
||||
Reference,
|
||||
}
|
||||
|
||||
impl MemoryType {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
MemoryType::User => "user",
|
||||
MemoryType::Feedback => "feedback",
|
||||
MemoryType::Project => "project",
|
||||
MemoryType::Reference => "reference",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for MemoryType {
|
||||
type Err = ();
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"user" => Ok(MemoryType::User),
|
||||
"feedback" => Ok(MemoryType::Feedback),
|
||||
"project" => Ok(MemoryType::Project),
|
||||
"reference" => Ok(MemoryType::Reference),
|
||||
_ => Err(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MemoryType {
|
||||
#[allow(clippy::should_implement_trait)]
|
||||
pub fn from_str(s: &str) -> Option<Self> {
|
||||
s.parse().ok()
|
||||
}
|
||||
}
|
||||
|
||||
/// 记忆生命周期状态 — 参考 Martian-Engineering agent-memory。
|
||||
///
|
||||
/// 事实永远不会被删除,只会从 Active 转换为 Historical。
|
||||
/// supersedes 链保留了"理解如何演变"的完整历史。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
#[derive(Default)]
|
||||
pub enum MemoryStatus {
|
||||
/// 当前有效
|
||||
#[default]
|
||||
Active,
|
||||
/// 已被更新事实取代,superseded_by 指向新 slug
|
||||
Historical { superseded_by: Option<String> },
|
||||
}
|
||||
|
||||
impl MemoryStatus {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
MemoryStatus::Active => "active",
|
||||
MemoryStatus::Historical { .. } => "historical",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str(s: &str, superseded_by: Option<String>) -> Self {
|
||||
match s {
|
||||
"historical" => MemoryStatus::Historical { superseded_by },
|
||||
_ => MemoryStatus::Active,
|
||||
}
|
||||
}
|
||||
|
||||
/// 是否为活跃状态
|
||||
pub fn is_active(&self) -> bool {
|
||||
matches!(self, MemoryStatus::Active)
|
||||
}
|
||||
|
||||
/// 获取取代此记忆的新 slug(如果有)
|
||||
pub fn superseded_by(&self) -> Option<&str> {
|
||||
match self {
|
||||
MemoryStatus::Historical {
|
||||
superseded_by: Some(s),
|
||||
} => Some(s.as_str()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// 记忆条目(从 .md 文件解析)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MemoryEntry {
|
||||
/// 文件名(不含扩展名,作为 slug)
|
||||
pub slug: String,
|
||||
/// 记忆标题
|
||||
pub name: String,
|
||||
/// 简短描述(用于相关性匹配)
|
||||
pub description: String,
|
||||
/// 记忆类型
|
||||
pub memory_type: MemoryType,
|
||||
/// 文件修改时间
|
||||
pub mtime: u64,
|
||||
/// 完整内容(不含 frontmatter)
|
||||
pub content: String,
|
||||
/// 文件路径
|
||||
pub path: PathBuf,
|
||||
/// 生命周期状态(默认 Active)
|
||||
pub status: MemoryStatus,
|
||||
}
|
||||
|
||||
/// MEMORY.md 索引中的一行
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MemoryIndexLine {
|
||||
pub slug: String,
|
||||
pub title: String,
|
||||
pub description: String,
|
||||
pub memory_type: MemoryType,
|
||||
}
|
||||
|
||||
/// 解析 Markdown 文件的 YAML frontmatter。
|
||||
///
|
||||
/// Frontmatter 格式(向后兼容,status/superseded_by 可选):
|
||||
/// ```markdown
|
||||
/// ---
|
||||
/// name: my-memory
|
||||
/// description: Short description
|
||||
/// metadata:
|
||||
/// type: user
|
||||
/// status: active # 可选,缺失默认 active
|
||||
/// superseded_by: "" # 可选,historical 时指向新 slug
|
||||
/// ---
|
||||
/// Content here...
|
||||
/// ```
|
||||
///
|
||||
/// 返回 (frontmatter_fields, content)。
|
||||
pub fn parse_frontmatter(raw: &str) -> (Vec<(String, String)>, String) {
|
||||
let mut fields = Vec::new();
|
||||
let content;
|
||||
|
||||
if let Some(rest) = raw.strip_prefix("---\n") {
|
||||
if let Some(end_pos) = rest.find("\n---\n") {
|
||||
let fm = &rest[..end_pos];
|
||||
content = rest[end_pos + 5..].to_string();
|
||||
|
||||
for line in fm.lines() {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Some(colon_pos) = line.find(':') {
|
||||
let key = line[..colon_pos].trim().to_string();
|
||||
let value = line[colon_pos + 1..].trim().to_string();
|
||||
fields.push((key, value));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
content = raw.to_string();
|
||||
}
|
||||
} else {
|
||||
content = raw.to_string();
|
||||
}
|
||||
|
||||
(fields, content)
|
||||
}
|
||||
|
||||
/// 从 frontmatter 字段构建 MemoryEntry。
|
||||
pub fn entry_from_frontmatter(
|
||||
slug: &str,
|
||||
fields: &[(String, String)],
|
||||
content: &str,
|
||||
path: PathBuf,
|
||||
mtime: u64,
|
||||
) -> Option<MemoryEntry> {
|
||||
let mut name = String::new();
|
||||
let mut description = String::new();
|
||||
let mut memory_type = MemoryType::User; // default
|
||||
let mut status_str = String::new();
|
||||
let mut superseded_by = String::new();
|
||||
|
||||
for (key, value) in fields {
|
||||
match key.as_str() {
|
||||
"name" => name = value.clone(),
|
||||
"description" => description = value.clone(),
|
||||
"type" | "memory_type" => {
|
||||
if let Some(t) = MemoryType::from_str(value) {
|
||||
memory_type = t;
|
||||
}
|
||||
}
|
||||
"status" => status_str = value.clone(),
|
||||
"superseded_by" => superseded_by = value.clone(),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if name.is_empty() {
|
||||
name = slug.to_string();
|
||||
}
|
||||
|
||||
let status = MemoryStatus::from_str(
|
||||
&status_str,
|
||||
if superseded_by.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(superseded_by)
|
||||
},
|
||||
);
|
||||
|
||||
Some(MemoryEntry {
|
||||
slug: slug.to_string(),
|
||||
name,
|
||||
description,
|
||||
memory_type,
|
||||
mtime,
|
||||
content: content.to_string(),
|
||||
path,
|
||||
status,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_basic() {
|
||||
let raw = "---\nname: test-memory\ndescription: A test\nmetadata:\n type: user\n---\nThis is the content.";
|
||||
let (fields, content) = parse_frontmatter(raw);
|
||||
assert_eq!(content.trim(), "This is the content.");
|
||||
assert!(fields
|
||||
.iter()
|
||||
.any(|(k, v)| k == "name" && v == "test-memory"));
|
||||
assert!(fields
|
||||
.iter()
|
||||
.any(|(k, v)| k == "description" && v == "A test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_no_frontmatter() {
|
||||
let raw = "Just content, no frontmatter.";
|
||||
let (fields, content) = parse_frontmatter(raw);
|
||||
assert_eq!(content, raw);
|
||||
assert!(fields.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_type_from_str() {
|
||||
assert_eq!(MemoryType::from_str("user"), Some(MemoryType::User));
|
||||
assert_eq!(MemoryType::from_str("feedback"), Some(MemoryType::Feedback));
|
||||
assert_eq!(MemoryType::from_str("project"), Some(MemoryType::Project));
|
||||
assert_eq!(
|
||||
MemoryType::from_str("reference"),
|
||||
Some(MemoryType::Reference)
|
||||
);
|
||||
assert_eq!(MemoryType::from_str("invalid"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_status_default_active() {
|
||||
let status = MemoryStatus::default();
|
||||
assert_eq!(status, MemoryStatus::Active);
|
||||
assert!(status.is_active());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_status_parse_active() {
|
||||
let status = MemoryStatus::from_str("active", None);
|
||||
assert_eq!(status, MemoryStatus::Active);
|
||||
assert!(status.is_active());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_status_parse_historical_with_superseded_by() {
|
||||
let status = MemoryStatus::from_str("historical", Some("new-version".to_string()));
|
||||
assert!(!status.is_active());
|
||||
assert_eq!(status.superseded_by(), Some("new-version"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_status_parse_unknown_defaults_to_active() {
|
||||
let status = MemoryStatus::from_str("invalid", None);
|
||||
assert_eq!(status, MemoryStatus::Active);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_entry_from_frontmatter_parses_status_and_superseded_by() {
|
||||
let fields = vec![
|
||||
("name".to_string(), "test-mem".to_string()),
|
||||
("description".to_string(), "A description".to_string()),
|
||||
("status".to_string(), "historical".to_string()),
|
||||
("superseded_by".to_string(), "better-slug".to_string()),
|
||||
];
|
||||
let entry = entry_from_frontmatter(
|
||||
"test-mem",
|
||||
&fields,
|
||||
"Content here",
|
||||
PathBuf::from("test-mem.md"),
|
||||
1000,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!entry.status.is_active());
|
||||
assert_eq!(entry.status.superseded_by(), Some("better-slug"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_entry_from_frontmatter_missing_status_defaults_to_active() {
|
||||
let fields = vec![
|
||||
("name".to_string(), "test-mem".to_string()),
|
||||
("description".to_string(), "A description".to_string()),
|
||||
];
|
||||
let entry = entry_from_frontmatter(
|
||||
"test-mem",
|
||||
&fields,
|
||||
"Content here",
|
||||
PathBuf::from("test-mem.md"),
|
||||
1000,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(entry.status.is_active());
|
||||
}
|
||||
}
|
||||
+21
-3
@@ -1,6 +1,24 @@
|
||||
// src/agent/mod.rs
|
||||
// 科研智能体模块
|
||||
// 基于 ReAct 框架实现 Thought -> Action -> Observation 循环
|
||||
//
|
||||
// 科研智能体模块 — 基于 ReAct 框架实现 Thought → Action → Observation 循环。
|
||||
//
|
||||
// 模块结构(参考 Claude Code 分层设计):
|
||||
// tools/ — 工具定义与注册(按功能域拆分)
|
||||
// runtime — ReAct 循环引擎 + Streaming + 会话管理
|
||||
// compact — 三层上下文压缩(micro/auto/manual)
|
||||
// terminal — 循环终止信号(结构化退出原因)
|
||||
// hooks — 生命周期事件系统(PreToolUse/PostToolUse/Stop)
|
||||
|
||||
pub mod tools;
|
||||
pub mod autonomous;
|
||||
pub mod background;
|
||||
pub mod compact;
|
||||
pub mod hooks;
|
||||
pub mod memory;
|
||||
pub mod runtime;
|
||||
pub mod skills;
|
||||
pub mod subagent;
|
||||
pub mod task_board;
|
||||
pub mod team;
|
||||
pub mod terminal;
|
||||
pub mod tools;
|
||||
pub mod trajectory;
|
||||
|
||||
@@ -1,786 +0,0 @@
|
||||
// src/agent/runtime.rs
|
||||
//
|
||||
// 科研智能体运行时核心模块。
|
||||
// 实现 ReAct 循环:Thought -> Action (工具调用) -> Observation -> Thought...
|
||||
// 支持会话持久化、上下文压缩、死循环检测和 SSE 流式输出。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn, error};
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{
|
||||
ChatMessage, LlmClient, MessageRole, StreamEvent,
|
||||
};
|
||||
use super::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// Agent 配置参数
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentConfig {
|
||||
/// 最大 ReAct 迭代次数
|
||||
pub max_steps: usize,
|
||||
/// 同质调用检测阈值(连续相同调用次数)
|
||||
pub duplicate_call_threshold: usize,
|
||||
/// 工具执行超时时间(秒)
|
||||
pub tool_timeout_secs: u64,
|
||||
/// 工具输出最大字符数
|
||||
pub max_tool_output_chars: usize,
|
||||
/// 上下文 Token 估算上限(触发自动摘要压缩)
|
||||
pub context_char_limit: usize,
|
||||
}
|
||||
|
||||
impl Default for AgentConfig {
|
||||
fn default() -> Self {
|
||||
AgentConfig {
|
||||
max_steps: std::env::var("AGENT_MAX_STEPS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(8),
|
||||
duplicate_call_threshold: 3,
|
||||
tool_timeout_secs: std::env::var("AGENT_TOOL_TIMEOUT_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(120),
|
||||
max_tool_output_chars: std::env::var("AGENT_MAX_TOOL_OUTPUT_CHARS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(4000),
|
||||
context_char_limit: std::env::var("AGENT_CONTEXT_CHAR_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(16000),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SSE 流式事件(发送给前端)
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AgentStreamEvent {
|
||||
/// 会话创建/恢复
|
||||
#[serde(rename = "session")]
|
||||
Session {
|
||||
session_id: String,
|
||||
title: String,
|
||||
},
|
||||
/// 智能体思考过程
|
||||
#[serde(rename = "thought")]
|
||||
Thought {
|
||||
content: String,
|
||||
step: usize,
|
||||
},
|
||||
/// 工具调用开始
|
||||
#[serde(rename = "tool_call")]
|
||||
ToolCall {
|
||||
name: String,
|
||||
arguments: serde_json::Value,
|
||||
step: usize,
|
||||
},
|
||||
/// 工具执行结果(Observation)
|
||||
#[serde(rename = "tool_result")]
|
||||
ToolResult {
|
||||
name: String,
|
||||
output: String,
|
||||
is_error: bool,
|
||||
metadata: serde_json::Value,
|
||||
step: usize,
|
||||
},
|
||||
/// 文本增量流式输出(最终回答)
|
||||
#[serde(rename = "text_delta")]
|
||||
TextDelta {
|
||||
content: String,
|
||||
},
|
||||
/// Token 使用统计
|
||||
#[serde(rename = "usage")]
|
||||
Usage {
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
total_tokens: u32,
|
||||
},
|
||||
/// 错误通知
|
||||
#[serde(rename = "error")]
|
||||
Error {
|
||||
message: String,
|
||||
},
|
||||
/// 完成标记
|
||||
#[serde(rename = "done")]
|
||||
Done,
|
||||
}
|
||||
|
||||
/// 同质调用检测器
|
||||
#[derive(Debug, Default)]
|
||||
struct DuplicateDetector {
|
||||
last_call: Option<(String, String)>, // (tool_name, arguments)
|
||||
consecutive_count: usize,
|
||||
}
|
||||
|
||||
impl DuplicateDetector {
|
||||
/// 记录一次调用,返回是否检测到死循环
|
||||
fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool {
|
||||
let key = (tool_name.to_string(), arguments.to_string());
|
||||
if self.last_call.as_ref() == Some(&key) {
|
||||
self.consecutive_count += 1;
|
||||
if self.consecutive_count >= threshold {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.last_call = Some(key);
|
||||
self.consecutive_count = 1;
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 智能体运行时
|
||||
pub struct AgentRuntime {
|
||||
app_state: Arc<AppState>,
|
||||
config: AgentConfig,
|
||||
tool_registry: ToolRegistry,
|
||||
}
|
||||
|
||||
impl AgentRuntime {
|
||||
/// 创建新的运行时实例
|
||||
pub fn new(app_state: Arc<AppState>) -> Self {
|
||||
AgentRuntime {
|
||||
app_state,
|
||||
config: AgentConfig::default(),
|
||||
tool_registry: ToolRegistry::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带自定义配置的运行时实例
|
||||
pub fn with_config(app_state: Arc<AppState>, config: AgentConfig) -> Self {
|
||||
AgentRuntime {
|
||||
app_state,
|
||||
config,
|
||||
tool_registry: ToolRegistry::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行完整的智能体对话回合(流式 SSE 输出)
|
||||
///
|
||||
/// 流程:
|
||||
/// 1. 加载或创建会话
|
||||
/// 2. 构建消息上下文
|
||||
/// 3. ReAct 循环:LLM 调用 -> 工具执行 -> 结果注入 -> 再次调用 ...
|
||||
/// 4. 最终回答流式输出
|
||||
/// 5. 持久化所有消息
|
||||
pub async fn run_turn(
|
||||
&self,
|
||||
session_id: Option<String>,
|
||||
question: &str,
|
||||
tx: mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
let db = &self.app_state.db;
|
||||
let llm = &self.app_state.llm;
|
||||
|
||||
// 1. 创建或恢复会话
|
||||
let sid = match session_id {
|
||||
Some(id) => {
|
||||
// 验证会话存在
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)"
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
|
||||
}
|
||||
id
|
||||
}
|
||||
None => {
|
||||
let new_id = uuid::Uuid::new_v4().to_string();
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)"
|
||||
)
|
||||
.bind(&new_id)
|
||||
.bind("")
|
||||
.bind(llm.model())
|
||||
.execute(db)
|
||||
.await?;
|
||||
new_id
|
||||
}
|
||||
};
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Session {
|
||||
session_id: sid.clone(),
|
||||
title: String::new(),
|
||||
});
|
||||
|
||||
// 2. 加载历史消息(过滤掉 thought 字段,仅保留纯对话上下文)
|
||||
let mut messages = self.load_history_for_llm(db, &sid).await?;
|
||||
|
||||
// 获取当前轮次号
|
||||
let turn_index: i32 = sqlx::query_scalar(
|
||||
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?"
|
||||
)
|
||||
.bind(&sid)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 3. 构建系统提示词
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(self.system_prompt()));
|
||||
}
|
||||
|
||||
// 4. 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
self.save_message(db, &sid, turn_index, 0, &ChatMessage::user(question), None).await?;
|
||||
|
||||
// 5. ReAct 循环
|
||||
let tool_defs = self.tool_registry.definitions();
|
||||
let tool_ctx = ToolContext {
|
||||
app_state: Arc::clone(&self.app_state),
|
||||
};
|
||||
let mut duplicate_detector = DuplicateDetector::default();
|
||||
let mut step = 0;
|
||||
|
||||
loop {
|
||||
step += 1;
|
||||
// 检查是否被用户手动中止
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.remove(&sid) {
|
||||
warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 上下文安全检查
|
||||
let context_chars: usize = messages.iter()
|
||||
.filter_map(|m| m.content.as_ref())
|
||||
.map(|c| c.len())
|
||||
.sum();
|
||||
|
||||
if context_chars > self.config.context_char_limit {
|
||||
info!("[AgentRuntime] 上下文超限 ({} > {}),触发压缩", context_chars, self.config.context_char_limit);
|
||||
self.compress_context(&mut messages, llm).await;
|
||||
}
|
||||
|
||||
// 调用 LLM(使用 `chat_stream` 实时流式读取,支持思维过程/最终回答的流式发送和中止检测)
|
||||
let mut stream_rx = match llm.chat_stream(&messages, &tool_defs).await {
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[AgentRuntime] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
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 usage: Option<crate::clients::llm::TokenUsage> = None;
|
||||
let mut is_tool_call_step = false;
|
||||
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
enum StreamLoopResult {
|
||||
Success,
|
||||
Error(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
let stream_loop_res = {
|
||||
let mut cancel_pinned = Box::pin(cancel_fut);
|
||||
let mut error_msg = None;
|
||||
let mut cancelled = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
event_opt = stream_rx.recv() => {
|
||||
match event_opt {
|
||||
Some(event) => {
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
// 实时流式发送思考过程给前端
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
// 如果目前还没发现是工具调用步骤,就实时流式发送文本给前端作为最终回答
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
is_tool_call_step = true;
|
||||
accumulated_tool_calls = Some(tool_calls);
|
||||
}
|
||||
StreamEvent::ToolCallDelta { .. } => {}
|
||||
StreamEvent::Usage(u) => {
|
||||
usage = Some(u);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
break;
|
||||
}
|
||||
StreamEvent::Error(e) => {
|
||||
error_msg = Some(e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = &mut cancel_pinned => {
|
||||
cancelled = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if cancelled {
|
||||
StreamLoopResult::Cancelled
|
||||
} else if let Some(e) = error_msg {
|
||||
StreamLoopResult::Error(e)
|
||||
} else {
|
||||
StreamLoopResult::Success
|
||||
}
|
||||
};
|
||||
|
||||
match stream_loop_res {
|
||||
StreamLoopResult::Success => {}
|
||||
StreamLoopResult::Error(e_str) => {
|
||||
error!("[AgentRuntime] 流式读取错误: {}", e_str);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式读取失败: {}", e_str),
|
||||
});
|
||||
break;
|
||||
}
|
||||
StreamLoopResult::Cancelled => {
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
cancelled.remove(&sid);
|
||||
}
|
||||
warn!("[AgentRuntime] 在流式调用期间被用户手动中止,会话 ID: {}", sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 处理 Thought(优先使用原生推理内容,否则如果属于工具调用步骤,使用 accumulated_content 存储)
|
||||
let mut thought_content = None;
|
||||
if !accumulated_reasoning.is_empty() {
|
||||
thought_content = Some(accumulated_reasoning.clone());
|
||||
}
|
||||
|
||||
if thought_content.is_none() && is_tool_call_step {
|
||||
if !accumulated_content.is_empty() {
|
||||
// 有工具调用时,content 被视为 Thought
|
||||
thought_content = Some(accumulated_content.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 如果是在工具调用步骤中产生的前言描述,而我们之前没实时以 Thought 发送过,此时统一作为 Thought 发送给前端展示
|
||||
if is_tool_call_step {
|
||||
if let Some(ref thought_text) = thought_content {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: thought_text.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let reasoning_option = if accumulated_reasoning.is_empty() { None } else { Some(accumulated_reasoning.clone()) };
|
||||
|
||||
// 无工具调用 = 最终回答
|
||||
if accumulated_tool_calls.is_none() || accumulated_tool_calls.as_ref().unwrap().is_empty() {
|
||||
// 保存助手最终回答消息
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
Some(accumulated_content.clone()),
|
||||
reasoning_option.clone(),
|
||||
None,
|
||||
);
|
||||
|
||||
self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?;
|
||||
messages.push(assistant_msg);
|
||||
|
||||
// 如果有原生推理内容且之前没发送过,发送给前端展示最终思维链
|
||||
if let Some(ref thought_text) = reasoning_option {
|
||||
if thought_content.is_none() {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: thought_text.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 流式发送 Done 或 Token 消耗
|
||||
if let Some(u) = usage {
|
||||
let _ = tx.send(AgentStreamEvent::Usage {
|
||||
prompt_tokens: u.prompt_tokens,
|
||||
completion_tokens: u.completion_tokens,
|
||||
total_tokens: u.total_tokens,
|
||||
});
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
let tool_calls = accumulated_tool_calls.unwrap();
|
||||
|
||||
// 有工具调用 —— 构建 assistant 消息(含 tool_calls 和 reasoning_content)
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
if accumulated_content.is_empty() { None } else { Some(accumulated_content.clone()) },
|
||||
reasoning_option.clone(),
|
||||
Some(tool_calls.clone()),
|
||||
);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?;
|
||||
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;
|
||||
|
||||
// 死循环检测
|
||||
if duplicate_detector.record(tool_name, tool_args_str, self.config.duplicate_call_threshold) {
|
||||
warn!("[AgentRuntime] 检测到死循环:{} 连续调用 {} 次", tool_name, self.config.duplicate_call_threshold);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用,已自动终止循环。", tool_name),
|
||||
});
|
||||
|
||||
// 注入错误 tool result 让 LLM 知道要停止
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!("错误:工具 {} 被连续重复调用 {} 次,参数完全相同。请停止重复调用并直接给出目前收集到的答案。", tool_name, self.config.duplicate_call_threshold),
|
||||
);
|
||||
messages.push(error_msg);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 解析参数
|
||||
let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
let error_output = format!("工具参数 JSON 解析失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?;
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 发送工具调用事件
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
name: tool_name.clone(),
|
||||
arguments: args.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
enum ToolResultEnum {
|
||||
Success(ToolOutput),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
// 执行工具(带超时保护和手动中止检测)
|
||||
let tool_res = match self.tool_registry.get(tool_name) {
|
||||
Some(tool) => {
|
||||
let timeout = std::time::Duration::from_secs(self.config.tool_timeout_secs);
|
||||
let tool_fut = tool.execute(args, &tool_ctx);
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
res = tokio::time::timeout(timeout, tool_fut) => {
|
||||
match res {
|
||||
Ok(output) => ToolResultEnum::Success(output),
|
||||
Err(_) => ToolResultEnum::Success(ToolOutput::error(format!("工具 {} 执行超时({}秒)", tool_name, self.config.tool_timeout_secs))),
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
ToolResultEnum::Cancelled
|
||||
}
|
||||
}
|
||||
}
|
||||
None => ToolResultEnum::Success(ToolOutput::error(format!("未知工具: {}", tool_name))),
|
||||
};
|
||||
|
||||
let output = match tool_res {
|
||||
ToolResultEnum::Success(out) => out,
|
||||
ToolResultEnum::Cancelled => {
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
cancelled.remove(&sid);
|
||||
}
|
||||
warn!("[AgentRuntime] 在工具 {} 执行期间被用户手动中止,会话 ID: {}", tool_name, sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// 发送工具结果事件
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
// 截断工具输出
|
||||
let truncated_content = if output.content.len() > self.config.max_tool_output_chars {
|
||||
let truncated: String = output.content.chars().take(self.config.max_tool_output_chars).collect();
|
||||
format!("{}...\n[输出已截断,原始长度: {} 字符]", truncated, output.content.len())
|
||||
} else {
|
||||
output.content.clone()
|
||||
};
|
||||
|
||||
// 构建 tool result 消息
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated_content);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?;
|
||||
messages.push(tool_msg);
|
||||
}
|
||||
}
|
||||
|
||||
// 6. 更新会话元信息
|
||||
let new_turn_count: i32 = sqlx::query_scalar(
|
||||
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?"
|
||||
)
|
||||
.bind(&sid)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 首轮自动生成标题
|
||||
if new_turn_count <= 1 {
|
||||
let title = self.generate_title(question);
|
||||
sqlx::query("UPDATE agent_sessions SET title = ?, turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?")
|
||||
.bind(&title)
|
||||
.bind(new_turn_count)
|
||||
.bind(&sid)
|
||||
.execute(db)
|
||||
.await?;
|
||||
} else {
|
||||
sqlx::query("UPDATE agent_sessions SET turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?")
|
||||
.bind(new_turn_count)
|
||||
.bind(&sid)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
|
||||
Ok(sid)
|
||||
}
|
||||
|
||||
/// 系统提示词
|
||||
fn system_prompt(&self) -> String {
|
||||
"你是一位专业的天体物理学研究助手,具备丰富的天文学知识。你可以使用以下工具帮助用户进行科研工作:\n\
|
||||
\n\
|
||||
- search_papers: 检索天文学文献(ADS/arXiv)\n\
|
||||
- get_paper_content: 获取文献全文内容(自动下载、解析)\n\
|
||||
- read_local_file: 快速读取已解析的本地文献\n\
|
||||
- rag_search: 在已向量化的文献库中进行语义检索\n\
|
||||
- query_target: 查询天体物理属性(坐标、光谱型等)\n\
|
||||
\n\
|
||||
请遵循以下原则:\n\
|
||||
1. 先思考用户的问题需要什么信息,再决定调用哪些工具。\n\
|
||||
2. 优先使用已有的本地文献资源(read_local_file / rag_search),必要时再检索新文献。\n\
|
||||
3. 回答时引用具体文献来源,使用 ADS bibcode 标注。\n\
|
||||
4. 对于数学公式,使用标准 LaTeX 格式。\n\
|
||||
5. 用中文回答用户的问题,但保持科学术语的准确性(可附带英文原文)。\n\
|
||||
6. 如果一个工具调用失败,不要重复使用完全相同的参数重试,尝试换一种方式。".to_string()
|
||||
}
|
||||
|
||||
/// 从数据库加载历史消息(包含 thought 作为 reasoning_content,以备原生思考模型使用)
|
||||
async fn load_history_for_llm(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let rows: Vec<(String, String, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? ORDER BY id ASC"
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (role_str, content, tool_calls_json, tool_call_id, _thought) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"system" => MessageRole::System,
|
||||
"user" => MessageRole::User,
|
||||
"assistant" => MessageRole::Assistant,
|
||||
"tool" => MessageRole::Tool,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = tool_calls_json
|
||||
.and_then(|json_str| serde_json::from_str(&json_str).ok());
|
||||
|
||||
messages.push(ChatMessage {
|
||||
role,
|
||||
content: if content.is_empty() { None } else { Some(content) },
|
||||
tool_call_id,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 保存消息到数据库
|
||||
async fn save_message(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step_index: i32,
|
||||
msg: &ChatMessage,
|
||||
thought: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let role = match msg.role {
|
||||
MessageRole::System => "system",
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
MessageRole::Tool => "tool",
|
||||
};
|
||||
|
||||
let content = msg.content.as_deref().unwrap_or("");
|
||||
let tool_calls_json = msg.tool_calls.as_ref()
|
||||
.map(|tc| serde_json::to_string(tc).unwrap_or_default());
|
||||
let tool_call_id = msg.tool_call_id.as_deref();
|
||||
let token_count = content.len() as i32 / 4; // 粗略估算
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step_index)
|
||||
.bind(role)
|
||||
.bind(content)
|
||||
.bind(thought)
|
||||
.bind(&tool_calls_json)
|
||||
.bind(tool_call_id)
|
||||
.bind(token_count)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 上下文压缩:保留系统提示、最近的 user 消息、以及最近的 tool_calls/tool 对
|
||||
async fn compress_context(
|
||||
&self,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
) {
|
||||
if messages.len() <= 4 {
|
||||
return;
|
||||
}
|
||||
|
||||
// 保留系统消息
|
||||
let system_msg = messages.first().cloned();
|
||||
|
||||
// 找到安全切割点:必须保证 assistant(tool_calls) 和后续 tool(result) 不被切断
|
||||
// 策略:保留最近 6 条消息 + 系统消息
|
||||
let keep_count = 6.min(messages.len() - 1);
|
||||
let to_summarize = &messages[1..messages.len() - keep_count];
|
||||
|
||||
if to_summarize.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// 生成摘要
|
||||
let summary_content: String = to_summarize.iter()
|
||||
.filter_map(|m| {
|
||||
let role = match m.role {
|
||||
MessageRole::User => "用户",
|
||||
MessageRole::Assistant => "助手",
|
||||
MessageRole::Tool => "工具",
|
||||
_ => return None,
|
||||
};
|
||||
m.content.as_ref().map(|c| {
|
||||
let preview: String = c.chars().take(200).collect();
|
||||
format!("[{}] {}", role, preview)
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let summary_prompt = format!(
|
||||
"请用简洁的中文总结以下对话历史的要点(不超过500字):\n\n{}",
|
||||
summary_content
|
||||
);
|
||||
|
||||
let summary = match llm.chat_completion(
|
||||
"你是一个对话摘要助手。请提取对话的关键信息和结论。",
|
||||
&summary_prompt,
|
||||
).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
warn!("[AgentRuntime] 上下文摘要生成失败: {},回退为简单截断", e);
|
||||
format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len())
|
||||
}
|
||||
};
|
||||
|
||||
// 重建消息列表
|
||||
let recent = messages[messages.len() - keep_count..].to_vec();
|
||||
messages.clear();
|
||||
if let Some(sys) = system_msg {
|
||||
messages.push(sys);
|
||||
}
|
||||
messages.push(ChatMessage::user(format!("[历史对话摘要]\n{}", summary)));
|
||||
messages.extend(recent);
|
||||
|
||||
info!("[AgentRuntime] 上下文压缩完成,消息数: {}", messages.len());
|
||||
}
|
||||
|
||||
/// 根据用户首条问题生成会话标题
|
||||
fn generate_title(&self, question: &str) -> String {
|
||||
let chars: String = question.chars().take(50).collect();
|
||||
if question.len() > 50 {
|
||||
format!("{}...", chars)
|
||||
} else {
|
||||
chars
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
// src/agent/runtime/circuit_breaker.rs
|
||||
//
|
||||
// 熔断器 — 防止无限自动压缩循环。
|
||||
// 参考 Claude Code MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES 设计。
|
||||
//
|
||||
// 当自动压缩连续失败 MAX_CONSECUTIVE_FAILURES 次后,熔断器打开,
|
||||
// 停止后续压缩尝试,避免无限循环。
|
||||
|
||||
use std::time::Instant;
|
||||
use tracing::warn;
|
||||
|
||||
/// 最大连续失败次数,超出后熔断器打开
|
||||
const MAX_CONSECUTIVE_FAILURES: usize = 3;
|
||||
|
||||
/// 熔断器打开后,经过此时间自动进入 HalfOpen 状态尝试恢复
|
||||
const AUTO_RECOVERY_TIMEOUT_SECS: u64 = 300; // 5 分钟
|
||||
|
||||
/// 熔断器状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum CircuitState {
|
||||
/// 正常工作,允许压缩
|
||||
Closed,
|
||||
/// 熔断,拒绝后续压缩
|
||||
Open,
|
||||
/// 半开:允许一次试探性压缩以决定是否恢复
|
||||
HalfOpen,
|
||||
}
|
||||
|
||||
/// 压缩熔断器
|
||||
#[derive(Debug)]
|
||||
pub struct CompactionCircuitBreaker {
|
||||
/// 连续失败计数
|
||||
consecutive_failures: usize,
|
||||
/// 压缩总次数
|
||||
total_compactions: usize,
|
||||
/// 当前状态
|
||||
state: CircuitState,
|
||||
/// 熔断器打开的时间(用于自动恢复)
|
||||
opened_at: Option<Instant>,
|
||||
}
|
||||
|
||||
impl CompactionCircuitBreaker {
|
||||
/// 创建新的熔断器(初始状态 Closed)
|
||||
pub fn new() -> Self {
|
||||
CompactionCircuitBreaker {
|
||||
consecutive_failures: 0,
|
||||
total_compactions: 0,
|
||||
state: CircuitState::Closed,
|
||||
opened_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录一次成功的压缩(重置失败计数,关闭熔断器)
|
||||
pub fn record_success(&mut self) {
|
||||
self.consecutive_failures = 0;
|
||||
self.total_compactions += 1;
|
||||
self.state = CircuitState::Closed;
|
||||
self.opened_at = None;
|
||||
}
|
||||
|
||||
/// 记录一次失败的压缩(递增失败计数,可能触发熔断)
|
||||
pub fn record_failure(&mut self) {
|
||||
self.consecutive_failures += 1;
|
||||
self.total_compactions += 1;
|
||||
|
||||
if self.consecutive_failures >= MAX_CONSECUTIVE_FAILURES {
|
||||
let was_already_open = self.state == CircuitState::Open;
|
||||
self.state = CircuitState::Open;
|
||||
self.opened_at = Some(Instant::now());
|
||||
if !was_already_open {
|
||||
warn!(
|
||||
"[CircuitBreaker] 熔断器打开!连续 {} 次压缩失败,停止自动压缩。\
|
||||
{} 秒后将自动尝试恢复。",
|
||||
self.consecutive_failures, AUTO_RECOVERY_TIMEOUT_SECS
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 熔断器是否打开(应停止自动压缩)
|
||||
pub fn is_open(&self) -> bool {
|
||||
self.state == CircuitState::Open
|
||||
}
|
||||
|
||||
/// 是否可以尝试压缩。
|
||||
///
|
||||
/// 当熔断器打开超过 AUTO_RECOVERY_TIMEOUT_SECS 时,自动转为 HalfOpen 状态,
|
||||
/// 允许下一次压缩尝试以判断是否恢复。
|
||||
pub fn can_attempt(&mut self) -> bool {
|
||||
match self.state {
|
||||
CircuitState::Closed | CircuitState::HalfOpen => true,
|
||||
CircuitState::Open => {
|
||||
// 检查是否已超时,可自动进入 HalfOpen
|
||||
if let Some(opened) = self.opened_at {
|
||||
if opened.elapsed().as_secs() >= AUTO_RECOVERY_TIMEOUT_SECS {
|
||||
self.state = CircuitState::HalfOpen;
|
||||
warn!(
|
||||
"[CircuitBreaker] 熔断器超时,进入 HalfOpen 状态,\
|
||||
允许下一次压缩尝试"
|
||||
);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 重置熔断器到 Closed 状态
|
||||
pub fn reset(&mut self) {
|
||||
self.consecutive_failures = 0;
|
||||
self.state = CircuitState::Closed;
|
||||
self.opened_at = None;
|
||||
}
|
||||
|
||||
/// 获取连续失败次数
|
||||
pub fn consecutive_failures(&self) -> usize {
|
||||
self.consecutive_failures
|
||||
}
|
||||
|
||||
/// 获取压缩总次数
|
||||
pub fn total_compactions(&self) -> usize {
|
||||
self.total_compactions
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CompactionCircuitBreaker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_opens_after_max_failures() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
assert!(!breaker.is_open());
|
||||
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert!(!breaker.is_open()); // 2 failures, not yet open
|
||||
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open()); // 3 failures, now open
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_on_success() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert_eq!(breaker.consecutive_failures(), 2);
|
||||
|
||||
breaker.record_success();
|
||||
assert_eq!(breaker.consecutive_failures(), 0);
|
||||
assert!(!breaker.is_open());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_method() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open());
|
||||
|
||||
breaker.reset();
|
||||
assert!(!breaker.is_open());
|
||||
assert_eq!(breaker.consecutive_failures(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_can_attempt() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
assert!(breaker.can_attempt());
|
||||
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
// 刚打开,不应允许尝试
|
||||
assert!(!breaker.can_attempt());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_half_open_after_reset_or_success() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
// 触发熔断
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
assert!(breaker.is_open());
|
||||
|
||||
// success 直接重置到 Closed
|
||||
breaker.record_success();
|
||||
assert!(!breaker.is_open());
|
||||
assert!(breaker.can_attempt());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_record_failure_while_open_stays_open() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
assert!(breaker.is_open());
|
||||
// 熔断器打开后再次失败,保持 Open
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open());
|
||||
assert_eq!(breaker.consecutive_failures(), 4);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
// src/agent/runtime/context.rs
|
||||
//
|
||||
// 上下文构建:加载历史消息、注入系统提示词、添加用户消息、
|
||||
// 从数据库恢复持久化的任务状态。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::info;
|
||||
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
use super::session;
|
||||
|
||||
/// 构建初始 LLM 上下文:加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。
|
||||
pub async fn build_initial_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
system_prompt: &str,
|
||||
question: &str,
|
||||
_turn_index: i32,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let mut messages = session::load_history_for_llm(db, session_id).await?;
|
||||
|
||||
// 注入系统提示词(如果历史中没有)
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(system_prompt));
|
||||
}
|
||||
|
||||
// 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
|
||||
// 从数据库恢复持久化的任务状态
|
||||
if let Some(task_reminder) = restore_tasks_from_db(db, session_id).await {
|
||||
messages.push(ChatMessage::user(task_reminder));
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 从 agent_tasks 表恢复任务状态,返回格式化的提醒文本。
|
||||
///
|
||||
/// 如果表不存在或没有任务记录,返回 None。
|
||||
async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option<String> {
|
||||
let rows: Vec<(String, String, String, String, Option<String>)> = sqlx::query_as(
|
||||
"SELECT task_id, content, status, blocked_by, owner \
|
||||
FROM agent_tasks WHERE session_id = ? AND (owner = '' OR owner = 'lead') \
|
||||
ORDER BY created_at ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
if rows.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut lines: Vec<String> = Vec::new();
|
||||
for (task_id, content, status, blocked_by, owner) in &rows {
|
||||
let icon = match status.as_str() {
|
||||
"in_progress" => "🔄",
|
||||
"completed" => "✅",
|
||||
_ => "⏳",
|
||||
};
|
||||
|
||||
let blocked: Vec<String> = serde_json::from_str(blocked_by).unwrap_or_default();
|
||||
let mut line = format!("{} [{}] {}", icon, task_id, content);
|
||||
if !blocked.is_empty() {
|
||||
line.push_str(&format!(" (依赖: {})", blocked.join(", ")));
|
||||
}
|
||||
if let Some(o) = owner {
|
||||
if !o.is_empty() && o != "lead" {
|
||||
line.push_str(&format!(" (指派: {})", o));
|
||||
}
|
||||
}
|
||||
lines.push(line);
|
||||
}
|
||||
|
||||
info!("[Context] 从数据库恢复了 {} 个任务状态", rows.len());
|
||||
|
||||
Some(format!(
|
||||
"[当前任务状态]\n以下是上次会话中持久化的任务计划,请基于最新状态继续工作:\n\n{}\n\n\
|
||||
使用 todo_write 工具更新任务进度。",
|
||||
lines.join("\n")
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
// src/agent/runtime/error_recovery.rs
|
||||
//
|
||||
// 错误恢复阶梯。
|
||||
// 参考 Claude Code error recovery ladder 设计。
|
||||
//
|
||||
// 当 LLM 流返回可恢复的错误(如 prompt_too_long)时,
|
||||
// 按阶梯顺序尝试恢复:
|
||||
// 1. Aggressive Compact — 激进微压缩(保留更少的工具结果)
|
||||
// 2. Reactive Compact — 使用 LLM 摘要压缩对话历史
|
||||
// 3. Escalate Tokens — 临时提升 token 上限到 64k
|
||||
// 4. Multi-Turn — 注入 metacognitive 消息分步处理
|
||||
// 5. Surface — 放弃恢复,暴露错误给用户
|
||||
//
|
||||
// 每一步都有 `has_attempted` 守卫,防止无限循环。
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::token_budget::TokenBudget;
|
||||
|
||||
/// 错误类型分类
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ErrorKind {
|
||||
/// 上下文过长(prompt too long / 413)
|
||||
PromptTooLong,
|
||||
/// Token 耗尽
|
||||
TokenExhausted,
|
||||
/// 模型错误
|
||||
ModelError(String),
|
||||
/// 超时
|
||||
Timeout,
|
||||
/// 限流(HTTP 429)
|
||||
RateLimited,
|
||||
/// 服务过载(HTTP 529)
|
||||
Overloaded,
|
||||
}
|
||||
|
||||
/// 恢复步骤
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum RecoveryStep {
|
||||
/// 尝试更激进的 micro_compact
|
||||
AggressiveCompact,
|
||||
/// 使用 LLM 摘要压缩
|
||||
ReactiveCompact,
|
||||
/// 提升 token 上限
|
||||
EscalateTokens { new_hard_limit: usize },
|
||||
/// 分轮恢复(注入 meta 消息)
|
||||
MultiTurn,
|
||||
/// 放弃,暴露错误
|
||||
Surface,
|
||||
/// 指数退避重试(用于 429/529 瞬态错误)
|
||||
RetryWithBackoff { attempt: u32, delay_ms: u64 },
|
||||
}
|
||||
|
||||
/// 恢复尝试追踪
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RecoveryAttempts {
|
||||
pub aggressive_compact: bool,
|
||||
pub reactive_compact: bool,
|
||||
pub escalate_tokens: bool,
|
||||
pub multi_turn: bool,
|
||||
}
|
||||
|
||||
impl RecoveryAttempts {
|
||||
pub fn new() -> Self {
|
||||
RecoveryAttempts {
|
||||
aggressive_compact: false,
|
||||
reactive_compact: false,
|
||||
escalate_tokens: false,
|
||||
multi_turn: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 是否有未尝试的恢复步骤
|
||||
pub fn has_remaining(&self) -> bool {
|
||||
!self.aggressive_compact
|
||||
|| !self.reactive_compact
|
||||
|| !self.escalate_tokens
|
||||
|| !self.multi_turn
|
||||
}
|
||||
|
||||
/// 获取下一个应尝试的恢复步骤
|
||||
pub fn next_step(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
|
||||
// 429/529 使用退避重试,不消耗上下文恢复步骤
|
||||
if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) {
|
||||
return Some(RecoveryStep::RetryWithBackoff {
|
||||
attempt: 0,
|
||||
delay_ms: 500,
|
||||
});
|
||||
}
|
||||
|
||||
match error_kind {
|
||||
ErrorKind::PromptTooLong | ErrorKind::TokenExhausted => {
|
||||
if !self.aggressive_compact {
|
||||
self.aggressive_compact = true;
|
||||
return Some(RecoveryStep::AggressiveCompact);
|
||||
}
|
||||
if !self.reactive_compact {
|
||||
self.reactive_compact = true;
|
||||
return Some(RecoveryStep::ReactiveCompact);
|
||||
}
|
||||
if !self.escalate_tokens {
|
||||
self.escalate_tokens = true;
|
||||
return Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000,
|
||||
});
|
||||
}
|
||||
if !self.multi_turn {
|
||||
self.multi_turn = true;
|
||||
return Some(RecoveryStep::MultiTurn);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// 非上下文相关错误,直接暴露
|
||||
if !self.multi_turn {
|
||||
self.multi_turn = true;
|
||||
return Some(RecoveryStep::Surface);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RecoveryAttempts {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 错误恢复器
|
||||
pub struct ErrorRecovery {
|
||||
/// 恢复步骤追踪
|
||||
pub attempts: RecoveryAttempts,
|
||||
/// Token 预算(用于 escalate 步骤)
|
||||
pub token_budget: TokenBudget,
|
||||
}
|
||||
|
||||
impl ErrorRecovery {
|
||||
/// 创建新的错误恢复器
|
||||
pub fn new(token_budget: TokenBudget) -> Self {
|
||||
ErrorRecovery {
|
||||
attempts: RecoveryAttempts::new(),
|
||||
token_budget,
|
||||
}
|
||||
}
|
||||
|
||||
/// 尝试从错误中恢复。
|
||||
///
|
||||
/// 返回 `Some(RecoveryStep)` 表示找到了恢复步骤(调用方应执行该步骤后重试)。
|
||||
/// 返回 `None` 表示所有步骤已尝试完毕,应暴露错误给用户。
|
||||
pub fn try_recover(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
|
||||
let step = self.attempts.next_step(error_kind);
|
||||
match &step {
|
||||
Some(RecoveryStep::AggressiveCompact) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 1/4: AggressiveCompact");
|
||||
}
|
||||
Some(RecoveryStep::ReactiveCompact) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 2/4: ReactiveCompact");
|
||||
}
|
||||
Some(RecoveryStep::EscalateTokens { new_hard_limit }) => {
|
||||
info!(
|
||||
"[ErrorRecovery] 尝试步骤 3/4: EscalateTokens → {}",
|
||||
new_hard_limit
|
||||
);
|
||||
self.token_budget.escalate_hard_limit(*new_hard_limit);
|
||||
}
|
||||
Some(RecoveryStep::RetryWithBackoff { attempt, delay_ms }) => {
|
||||
info!(
|
||||
"[ErrorRecovery] 退避重试: attempt={}, delay={}ms",
|
||||
attempt, delay_ms
|
||||
);
|
||||
}
|
||||
Some(RecoveryStep::MultiTurn) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 4/4: MultiTurn");
|
||||
}
|
||||
Some(RecoveryStep::Surface) => {
|
||||
warn!("[ErrorRecovery] 无法恢复,暴露错误");
|
||||
}
|
||||
None => {
|
||||
warn!("[ErrorRecovery] 所有恢复步骤已尝试完毕");
|
||||
}
|
||||
}
|
||||
step
|
||||
}
|
||||
|
||||
/// 生成 multi-turn 恢复消息(注入到对话中以继续处理)
|
||||
pub fn multi_turn_message() -> String {
|
||||
"由于 token 限制,当前回答被截断。请基于已收集的信息继续分析,\
|
||||
重点关注尚未完成的部分。你可以:\n\
|
||||
1. 总结已有发现\n\
|
||||
2. 使用 compress_context 手动压缩上下文\n\
|
||||
3. 分步完成剩余工作"
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// 检查是否需要恢复(错误是否可恢复)
|
||||
pub fn is_recoverable(error_kind: &ErrorKind) -> bool {
|
||||
matches!(
|
||||
error_kind,
|
||||
ErrorKind::PromptTooLong
|
||||
| ErrorKind::TokenExhausted
|
||||
| ErrorKind::RateLimited
|
||||
| ErrorKind::Overloaded
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算指数退避延迟(毫秒)。
|
||||
///
|
||||
/// 公式:min(500 * 2^attempt, 32000) + 25% 随机抖动
|
||||
/// 如果有 Retry-After header,优先使用。
|
||||
pub fn backoff_delay(attempt: u32, retry_after_secs: Option<u64>) -> u64 {
|
||||
if let Some(ra) = retry_after_secs {
|
||||
return ra * 1000;
|
||||
}
|
||||
let base = 500u64 * 2u64.pow(attempt.min(6)); // cap at 2^6 = 64 → 32000ms
|
||||
let base = base.min(32_000);
|
||||
// Simple deterministic jitter using attempt (avoid rand dependency)
|
||||
let jitter = (base / 4) * (attempt as u64 % 5) / 5;
|
||||
base + jitter
|
||||
}
|
||||
|
||||
/// 解析错误字符串中的 Retry-After header 值。
|
||||
///
|
||||
/// 期望格式: `retry_after=Some(N)` 出现在错误消息中。
|
||||
pub fn parse_retry_after(error_str: &str) -> Option<u64> {
|
||||
if let Some(pos) = error_str.find("retry_after=Some(") {
|
||||
let prefix_len = "retry_after=Some(".len(); // 19
|
||||
let rest = &error_str[pos + prefix_len..];
|
||||
if let Some(end) = rest.find(')') {
|
||||
return rest[..end].parse().ok();
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 从错误字符串分类错误类型。
|
||||
/// 解析 LLM API 返回的错误消息,映射到 ErrorKind。
|
||||
pub fn classify_error(error_str: &str) -> ErrorKind {
|
||||
let lower = error_str.to_lowercase();
|
||||
|
||||
// 先检测限流/过载(HTTP 状态码检查)
|
||||
if lower.contains("429")
|
||||
|| lower.contains("rate limit")
|
||||
|| lower.contains("rate_limit")
|
||||
|| lower.contains("too many requests")
|
||||
{
|
||||
return ErrorKind::RateLimited;
|
||||
}
|
||||
if lower.contains("529")
|
||||
|| lower.contains("overloaded")
|
||||
|| lower.contains("overload")
|
||||
|| lower.contains("service overloaded")
|
||||
{
|
||||
return ErrorKind::Overloaded;
|
||||
}
|
||||
|
||||
if lower.contains("prompt_too_long")
|
||||
|| lower.contains("prompt too long")
|
||||
|| lower.contains("context length")
|
||||
|| lower.contains("413")
|
||||
|| lower.contains("context_window_exceeded")
|
||||
|| lower.contains("input length")
|
||||
{
|
||||
return ErrorKind::PromptTooLong;
|
||||
}
|
||||
|
||||
if lower.contains("max_tokens")
|
||||
|| lower.contains("token limit")
|
||||
|| lower.contains("token_exhausted")
|
||||
|| lower.contains("maximum context length")
|
||||
|| lower.contains("reduce the length")
|
||||
{
|
||||
return ErrorKind::TokenExhausted;
|
||||
}
|
||||
|
||||
if lower.contains("timeout")
|
||||
|| lower.contains("timed out")
|
||||
|| lower.contains("deadline exceeded")
|
||||
|| lower.contains("408")
|
||||
|| lower.contains("504")
|
||||
{
|
||||
return ErrorKind::Timeout;
|
||||
}
|
||||
|
||||
// 默认归类为模型错误
|
||||
ErrorKind::ModelError(error_str.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_all_steps_sequence() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::AggressiveCompact)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::ReactiveCompact)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::MultiTurn)
|
||||
);
|
||||
// 所有步骤已尝试
|
||||
assert_eq!(attempts.next_step(&ErrorKind::PromptTooLong), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_error_goes_straight_to_surface() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::ModelError("test".into())),
|
||||
Some(RecoveryStep::Surface)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::ModelError("test".into())),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_remaining() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert!(attempts.has_remaining());
|
||||
|
||||
// 消耗所有步骤
|
||||
for _ in 0..4 {
|
||||
attempts.next_step(&ErrorKind::PromptTooLong);
|
||||
}
|
||||
assert!(!attempts.has_remaining());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_recoverable() {
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::PromptTooLong));
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::TokenExhausted));
|
||||
assert!(!ErrorRecovery::is_recoverable(&ErrorKind::ModelError(
|
||||
"test".into()
|
||||
)));
|
||||
assert!(!ErrorRecovery::is_recoverable(&ErrorKind::Timeout));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_rate_limited_429() {
|
||||
let kind = classify_error("HTTP 429: Too Many Requests");
|
||||
assert_eq!(kind, ErrorKind::RateLimited);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_overloaded_529() {
|
||||
let kind = classify_error("HTTP 529: Service Overloaded");
|
||||
assert_eq!(kind, ErrorKind::Overloaded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limited_is_recoverable() {
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited));
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_retry_after() {
|
||||
let err = "HTTP 429: retry_after=Some(30)";
|
||||
assert_eq!(parse_retry_after(err), Some(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_retry_after_none() {
|
||||
let err = "HTTP 500: Internal Server Error";
|
||||
assert_eq!(parse_retry_after(err), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_delay() {
|
||||
// Attempt 0: 500 + jitter
|
||||
let d0 = backoff_delay(0, None);
|
||||
assert!(d0 >= 500 && d0 <= 700);
|
||||
|
||||
// Attempt 3: 500*8=4000 + jitter
|
||||
let d3 = backoff_delay(3, None);
|
||||
assert!(d3 >= 4000 && d3 <= 5000);
|
||||
|
||||
// Capped at 32s
|
||||
let d10 = backoff_delay(10, None);
|
||||
assert!(d10 <= 40_000);
|
||||
|
||||
// Retry-After takes priority
|
||||
let d_ra = backoff_delay(0, Some(15));
|
||||
assert_eq!(d_ra, 15000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limited_goes_to_retry() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
let step = attempts.next_step(&ErrorKind::RateLimited);
|
||||
assert!(matches!(step, Some(RecoveryStep::RetryWithBackoff { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_escalation() {
|
||||
let budget = TokenBudget::new(32_000, 40_000);
|
||||
let mut recovery = ErrorRecovery::new(budget);
|
||||
assert_eq!(recovery.token_budget.hard_limit, 40_000);
|
||||
|
||||
recovery.attempts.next_step(&ErrorKind::PromptTooLong); // aggressive
|
||||
recovery.attempts.next_step(&ErrorKind::PromptTooLong); // reactive
|
||||
let step = recovery.try_recover(&ErrorKind::PromptTooLong); // escalate
|
||||
|
||||
assert_eq!(
|
||||
step,
|
||||
Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000
|
||||
})
|
||||
);
|
||||
assert_eq!(recovery.token_budget.hard_limit, 64_000);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
// src/agent/runtime/executor.rs
|
||||
//
|
||||
// 工具调用执行器:验证 → PreToolUse hooks → 并行执行 → 结果收集 → PostToolUse hooks。
|
||||
|
||||
use futures_util::stream::FuturesUnordered;
|
||||
use futures_util::StreamExt;
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{ChatMessage, ToolCall};
|
||||
|
||||
use super::file_cache::FileStateCache;
|
||||
use super::permission::PermissionChecker;
|
||||
use super::{AgentStreamEvent, DuplicateDetector};
|
||||
use crate::agent::hooks::{HookRegistry, PostToolUseContext, PreToolUseContext};
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result;
|
||||
use crate::agent::tools::{InterruptBehavior, ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 准备好的工具调用
|
||||
#[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,
|
||||
}
|
||||
|
||||
/// 验证工具调用:死循环检测 + 参数解析。
|
||||
///
|
||||
/// 返回 (prepared_calls, has_duplicate)。
|
||||
/// 死循环或参数无效时,错误消息直接注入到 messages。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn validate_and_prepare(
|
||||
tool_calls: &[ToolCall],
|
||||
duplicate_detector: &mut DuplicateDetector,
|
||||
duplicate_threshold: usize,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
) -> (Vec<PreparedCall>, bool) {
|
||||
let mut prepared_calls: Vec<PreparedCall> = Vec::new();
|
||||
let mut has_duplicate = false;
|
||||
|
||||
for tool_call in tool_calls {
|
||||
let tool_name = &tool_call.function.name;
|
||||
let tool_args_str = &tool_call.function.arguments;
|
||||
|
||||
// 死循环检测
|
||||
if duplicate_detector.record(tool_name, tool_args_str, duplicate_threshold) {
|
||||
warn!(
|
||||
"[Executor] 检测到死循环:{} 连续调用 {} 次",
|
||||
tool_name, duplicate_threshold
|
||||
);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用,已自动终止循环。", tool_name),
|
||||
});
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!(
|
||||
"错误:工具 {} 被连续重复调用 {} 次,参数完全相同。\
|
||||
请停止重复调用并直接给出目前收集到的答案。",
|
||||
tool_name, duplicate_threshold
|
||||
),
|
||||
);
|
||||
messages.push(error_msg);
|
||||
has_duplicate = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
// 解析参数
|
||||
let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
let error_output = format!("工具参数 JSON 解析失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output);
|
||||
save_tool_message_sync(db, session_id, turn_index, step, &tool_msg);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
prepared_calls.push(PreparedCall {
|
||||
tool_call_id: tool_call.id.clone(),
|
||||
tool_name: tool_name.clone(),
|
||||
args,
|
||||
});
|
||||
}
|
||||
|
||||
(prepared_calls, has_duplicate)
|
||||
}
|
||||
|
||||
/// 并行执行所有准备好的工具调用。
|
||||
///
|
||||
/// 流程:
|
||||
/// 1. 权限检查(deny 规则阻止不可执行工具)
|
||||
/// 2. 发送 ToolCall SSE 事件
|
||||
/// 3. 运行 PreToolUse hooks
|
||||
/// 4. 工具分区 + 并行执行(并发安全工具一批并行,不安全工具单独串行)
|
||||
/// 5. 收集结果、发送 ToolResult SSE、运行 PostToolUse hooks
|
||||
/// 6. 返回 ToolResultMessage 列表供调用方推入 messages
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn execute_parallel(
|
||||
prepared_calls: &[PreparedCall],
|
||||
tool_registry: &ToolRegistry,
|
||||
app_state: Arc<AppState>,
|
||||
hook_registry: &HookRegistry,
|
||||
_permission_checker: Option<&PermissionChecker>,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
agent_name: &str,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
tool_timeout_secs: u64,
|
||||
max_output_chars: usize,
|
||||
read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
) -> ToolExecutionResult {
|
||||
if prepared_calls.is_empty() {
|
||||
return ToolExecutionResult {
|
||||
tool_messages: Vec::new(),
|
||||
was_cancelled: false,
|
||||
had_duplicate: false,
|
||||
};
|
||||
}
|
||||
|
||||
let sid = session_id.to_string();
|
||||
|
||||
// Phase 1: 发送 ToolCall SSE 事件
|
||||
for prep in prepared_calls {
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
name: prep.tool_name.clone(),
|
||||
arguments: prep.args.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
|
||||
// Phase 2: PreToolUse hooks — 收集修改后的参数和附加上下文
|
||||
let exec_start = std::time::Instant::now();
|
||||
let mut mutated_args: Vec<serde_json::Value> = Vec::new();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
for prep in prepared_calls {
|
||||
let hook_ctx = PreToolUseContext {
|
||||
session_id: sid.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
tool_args: prep.args.clone(),
|
||||
step,
|
||||
};
|
||||
let result = hook_registry.run_pre_tool_use(&hook_ctx).await;
|
||||
if result.action.is_blocked() {
|
||||
let reason = result.action.block_reason().unwrap_or("unknown");
|
||||
warn!(
|
||||
"[Executor] PreToolUse hook 阻止了 {} 的执行: {}",
|
||||
prep.tool_name, reason
|
||||
);
|
||||
}
|
||||
// 使用 hook 可能修改后的参数
|
||||
mutated_args.push(result.final_args);
|
||||
if let Some(ctx) = result.additional_context {
|
||||
additional_contexts.push(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 3: 并行执行
|
||||
let cancelled = Arc::new(AtomicBool::new(false));
|
||||
let cancel_flag = cancelled.clone();
|
||||
let app_state_ref = app_state.clone();
|
||||
let sid_ref = sid.clone();
|
||||
|
||||
let cancel_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(locked) = app_state_ref.cancelled_runs.lock() {
|
||||
if locked.contains(&sid_ref) {
|
||||
cancel_flag.store(true, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs);
|
||||
|
||||
// Phase 3: 使用 FuturesUnordered 进行渐进式并行执行。
|
||||
// 每个工具完成后立即发送 SSE ToolResult 事件到前端(非阻塞),
|
||||
// 而后台继续等待其他工具完成。快工具的结果不会因慢工具而延迟。
|
||||
let mut exec_futs: FuturesUnordered<_> = prepared_calls
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, prep)| {
|
||||
let tool_name = prep.tool_name.clone();
|
||||
let args = mutated_args
|
||||
.get(i)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx = ToolContext::with_file_cache(app_state.clone(), read_file_state.clone());
|
||||
let cancelled = cancelled.clone();
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
Box::pin(async move {
|
||||
let output = match tool_opt {
|
||||
Some(tool) => {
|
||||
let interrupt_behavior = tool.interrupt_behavior();
|
||||
let is_blocking = interrupt_behavior == InterruptBehavior::Block;
|
||||
|
||||
let tool_fut = tool.execute(args, &tool_ctx);
|
||||
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(_) => ToolOutput::error(format!(
|
||||
"工具 {} 执行超时({}秒)",
|
||||
tool_name,
|
||||
timeout_dur.as_secs()
|
||||
)),
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
ToolOutput::error("执行已被用户取消")
|
||||
}
|
||||
}
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
let was_cancelled = cancelled.load(Ordering::SeqCst);
|
||||
(
|
||||
prep.tool_call_id.clone(),
|
||||
prep.tool_name.clone(),
|
||||
prep.args.clone(),
|
||||
output,
|
||||
was_cancelled,
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut tool_messages: Vec<ToolResultMessage> = Vec::new();
|
||||
let mut was_cancelled = false;
|
||||
|
||||
// 渐进式处理结果:每个工具一完成就立即处理(SSE 事件 + PostToolUse hook + 持久化)
|
||||
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
|
||||
exec_futs.next().await
|
||||
{
|
||||
if cancelled_flag {
|
||||
was_cancelled = true;
|
||||
}
|
||||
|
||||
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
|
||||
|
||||
// SSE 事件 — 立即推送到前端
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
// 输出处理:小结果直接传递,大结果持久化到磁盘并返回 stub
|
||||
let tool_results_dir = app_state.config.library_dir.join("tool-results");
|
||||
let (processed_content, _persisted_path) = maybe_persist_tool_result(
|
||||
&output.content,
|
||||
&tool_call_id,
|
||||
max_output_chars,
|
||||
&tool_results_dir,
|
||||
);
|
||||
|
||||
// PostToolUse hook
|
||||
let post_ctx = PostToolUseContext {
|
||||
session_id: sid.clone(),
|
||||
agent_name: agent_name.to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
tool_args,
|
||||
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;
|
||||
|
||||
let chat_message = ChatMessage::tool_result(&tool_call_id, &final_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,
|
||||
});
|
||||
}
|
||||
|
||||
cancel_handle.abort();
|
||||
|
||||
ToolExecutionResult {
|
||||
tool_messages,
|
||||
was_cancelled,
|
||||
had_duplicate: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 同步保存 tool 角色消息到数据库。
|
||||
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.content.as_deref().unwrap_or("").to_string();
|
||||
let tool_call_id = msg.tool_call_id.clone();
|
||||
// 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, 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("lead")
|
||||
.execute(&db_clone)
|
||||
.await
|
||||
{
|
||||
warn!("[Executor] 保存 tool 消息失败(非致命): {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,571 @@
|
||||
// src/agent/runtime/file_cache.rs
|
||||
//
|
||||
// 文件状态缓存 — 参考 Claude Code FileStateCache 设计。
|
||||
//
|
||||
// 在 Read 工具调用前检查缓存:
|
||||
// 1. 路径已缓存 → 读取磁盘 mtime → mtime 相同 + offset/limit 一致 → 返回 stub
|
||||
// 2. mtime 不同或新文件 → 正常读取 → 写入缓存
|
||||
//
|
||||
// 压缩时:
|
||||
// - 压缩前:快照缓存到普通对象
|
||||
// - 压缩后:清空缓存,将最近 N 个文件作为上下文注入
|
||||
//
|
||||
// 缓存上限:100 个条目,25MB 内容总大小(LRU 自动淘汰)。
|
||||
|
||||
use lru::LruCache;
|
||||
use std::num::NonZeroUsize;
|
||||
use std::path::Path;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// 缓存条目最大数量
|
||||
pub const MAX_ENTRIES: usize = 100;
|
||||
/// 缓存内容总大小上限(25MB)
|
||||
pub const MAX_CACHE_SIZE_BYTES: usize = 25 * 1024 * 1024;
|
||||
|
||||
/// 文件不变时的占位消息(参考 Claude Code FILE_UNCHANGED_STUB)
|
||||
pub const FILE_UNCHANGED_STUB: &str =
|
||||
"File unchanged since last read. The content from the earlier read_file tool_result \
|
||||
in this conversation is still current — refer to that instead of re-reading.";
|
||||
|
||||
/// 压缩后恢复的最大文件数
|
||||
pub const POST_COMPACT_MAX_FILES_TO_RESTORE: usize = 5;
|
||||
/// 压缩后恢复的每文件最大 token 数(~字符数)
|
||||
pub const POST_COMPACT_MAX_CHARS_PER_FILE: usize = 4_000;
|
||||
|
||||
/// 单个文件的缓存状态
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FileState {
|
||||
/// 上次读取的文件内容
|
||||
pub content: String,
|
||||
/// 文件修改时间(Unix 时间戳,秒级)
|
||||
pub timestamp: i64,
|
||||
/// 读取起始行(1-based)
|
||||
pub offset: usize,
|
||||
/// 行数限制
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// 文件状态快照(用于压缩前后传递,纯数据,不含 LRU 结构)
|
||||
pub type FileStateSnapshot = Vec<(String, FileState)>;
|
||||
|
||||
/// 文件状态缓存。
|
||||
///
|
||||
/// 包装 `LruCache<String, FileState>` + 内容总大小追踪。
|
||||
/// 通过 `Arc<Mutex<FileStateCache>>` 在工具调用间共享。
|
||||
pub struct FileStateCache {
|
||||
cache: LruCache<String, FileState>,
|
||||
/// 当前缓存中所有内容的字节数总和(近似,使用 content.len())
|
||||
current_size_bytes: usize,
|
||||
/// 最大字节数
|
||||
max_size_bytes: usize,
|
||||
}
|
||||
|
||||
impl FileStateCache {
|
||||
/// 创建新的缓存实例
|
||||
pub fn new() -> Self {
|
||||
let max_entries = NonZeroUsize::new(MAX_ENTRIES).unwrap();
|
||||
FileStateCache {
|
||||
cache: LruCache::new(max_entries),
|
||||
current_size_bytes: 0,
|
||||
max_size_bytes: MAX_CACHE_SIZE_BYTES,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带自定义参数的新缓存
|
||||
pub fn with_limits(max_entries: usize, max_size_bytes: usize) -> Self {
|
||||
let max_entries =
|
||||
NonZeroUsize::new(max_entries).unwrap_or(NonZeroUsize::new(MAX_ENTRIES).unwrap());
|
||||
FileStateCache {
|
||||
cache: LruCache::new(max_entries),
|
||||
current_size_bytes: 0,
|
||||
max_size_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
/// 规范化路径 key(确保一致性)
|
||||
fn normalize_key(path: &str) -> String {
|
||||
// 去除尾随斜杠,规范化重复斜杠
|
||||
let p = Path::new(path);
|
||||
// 尝试 canonicalize(跟随符号链接),失败则用简单的字符串规范化
|
||||
match p.canonicalize() {
|
||||
Ok(canon) => canon.to_string_lossy().to_string(),
|
||||
Err(_) => {
|
||||
// 简单规范化:折叠重复的 /
|
||||
let mut result = String::with_capacity(path.len());
|
||||
let mut prev_slash = false;
|
||||
for ch in path.chars() {
|
||||
if ch == '/' || ch == '\\' {
|
||||
if !prev_slash {
|
||||
result.push('/');
|
||||
prev_slash = true;
|
||||
}
|
||||
} else {
|
||||
result.push(ch);
|
||||
prev_slash = false;
|
||||
}
|
||||
}
|
||||
// 去除尾随 /
|
||||
if result.ends_with('/') && result.len() > 1 {
|
||||
result.pop();
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取缓存的条目,返回 Some(&FileState) 若存在
|
||||
pub fn get(&mut self, path: &str) -> Option<&FileState> {
|
||||
let key = Self::normalize_key(path);
|
||||
self.cache.get(&key)
|
||||
}
|
||||
|
||||
/// 写入缓存条目,自动处理容量限制
|
||||
pub fn set(&mut self, path: &str, state: FileState) {
|
||||
let key = Self::normalize_key(path);
|
||||
let content_len = state.content.len();
|
||||
|
||||
// 如果 key 已存在,先减去旧内容的 size
|
||||
if let Some(old) = self.cache.get(&key) {
|
||||
self.current_size_bytes = self.current_size_bytes.saturating_sub(old.content.len());
|
||||
}
|
||||
|
||||
// 驱逐旧条目直到有足够空间
|
||||
while self.current_size_bytes + content_len > self.max_size_bytes && !self.cache.is_empty() {
|
||||
if let Some((_, evicted)) = self.cache.pop_lru() {
|
||||
self.current_size_bytes = self
|
||||
.current_size_bytes
|
||||
.saturating_sub(evicted.content.len());
|
||||
}
|
||||
}
|
||||
|
||||
// 如果单个文件超过上限,仍存储但记录警告
|
||||
if content_len > self.max_size_bytes {
|
||||
warn!(
|
||||
"[FileCache] 单个文件内容 ({} bytes) 超过缓存上限 ({} bytes)",
|
||||
content_len, self.max_size_bytes
|
||||
);
|
||||
}
|
||||
|
||||
self.current_size_bytes += content_len;
|
||||
self.cache.push(key, state);
|
||||
|
||||
info!(
|
||||
"[FileCache] 缓存写入: path={}, size={} bytes, cache_entries={}, cache_size={}",
|
||||
path,
|
||||
content_len,
|
||||
self.cache.len(),
|
||||
self.current_size_bytes
|
||||
);
|
||||
}
|
||||
|
||||
/// 检查 key 是否存在
|
||||
pub fn contains(&mut self, path: &str) -> bool {
|
||||
let key = Self::normalize_key(path);
|
||||
self.cache.contains(&key)
|
||||
}
|
||||
|
||||
/// 删除缓存条目
|
||||
pub fn remove(&mut self, path: &str) -> bool {
|
||||
let key = Self::normalize_key(path);
|
||||
if let Some(removed) = self.cache.pop(&key) {
|
||||
self.current_size_bytes = self
|
||||
.current_size_bytes
|
||||
.saturating_sub(removed.content.len());
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 清空缓存
|
||||
pub fn clear(&mut self) {
|
||||
self.cache.clear();
|
||||
self.current_size_bytes = 0;
|
||||
info!("[FileCache] 缓存已清空");
|
||||
}
|
||||
|
||||
/// 缓存条目数
|
||||
pub fn len(&self) -> usize {
|
||||
self.cache.len()
|
||||
}
|
||||
|
||||
/// 缓存是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.cache.len() == 0
|
||||
}
|
||||
|
||||
/// 当前缓存内容总大小(近似字节数)
|
||||
pub fn current_size_bytes(&self) -> usize {
|
||||
self.current_size_bytes
|
||||
}
|
||||
|
||||
// ── 压缩集成 ──
|
||||
|
||||
/// 生成快照(纯数据,不含 LRU 结构)。
|
||||
/// 在压缩前调用,用于压缩后恢复文件上下文。
|
||||
pub fn to_snapshot(&mut self) -> FileStateSnapshot {
|
||||
// 按 timestamp 降序排序(最近读的在前)
|
||||
let mut entries: Vec<(String, FileState)> = self
|
||||
.cache
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.collect();
|
||||
entries.sort_by(|a, b| b.1.timestamp.cmp(&a.1.timestamp));
|
||||
entries
|
||||
}
|
||||
|
||||
/// 从快照恢复指定数量的最近文件到缓存。
|
||||
/// 在压缩后调用。
|
||||
pub fn restore_from_snapshot(&mut self, snapshot: &FileStateSnapshot, max_files: usize) {
|
||||
for (path, state) in snapshot.iter().take(max_files) {
|
||||
// 不恢复过大的文件(已有提示说可能过时)
|
||||
if state.content.len() > POST_COMPACT_MAX_CHARS_PER_FILE {
|
||||
continue;
|
||||
}
|
||||
self.set(path, state.clone());
|
||||
}
|
||||
info!(
|
||||
"[FileCache] 从快照恢复了 {} 个文件 (快照大小: {})",
|
||||
self.len().min(max_files),
|
||||
snapshot.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// 从快照生成上 下文注入文本(用于压缩后注入到对话中)。
|
||||
/// 返回格式化的 markdown 块列表。
|
||||
pub fn build_restore_context(snapshot: &FileStateSnapshot, max_files: usize) -> Vec<String> {
|
||||
if snapshot.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut contexts: Vec<String> = Vec::new();
|
||||
let mut used_chars = 0usize;
|
||||
let total_budget = POST_COMPACT_MAX_CHARS_PER_FILE * max_files;
|
||||
|
||||
for (path, state) in snapshot.iter().take(max_files) {
|
||||
let preview: String = state
|
||||
.content
|
||||
.chars()
|
||||
.take(POST_COMPACT_MAX_CHARS_PER_FILE)
|
||||
.collect();
|
||||
let truncated = if state.content.len() > preview.len() {
|
||||
format!(
|
||||
"{}…\n[内容已截断: {} 字符 → {} 字符]",
|
||||
preview,
|
||||
state.content.len(),
|
||||
preview.len()
|
||||
)
|
||||
} else {
|
||||
preview
|
||||
};
|
||||
|
||||
if used_chars + truncated.len() > total_budget {
|
||||
break;
|
||||
}
|
||||
|
||||
let block = format!(
|
||||
"[压缩后恢复: {}]\n上次读取时间戳: {}\n内容:\n```\n{}\n```",
|
||||
path, state.timestamp, truncated
|
||||
);
|
||||
used_chars += block.len();
|
||||
contexts.push(block);
|
||||
}
|
||||
|
||||
if !contexts.is_empty() {
|
||||
info!(
|
||||
"[FileCache] 生成压缩恢复上下文: {} 文件, {} chars",
|
||||
contexts.len(),
|
||||
used_chars
|
||||
);
|
||||
}
|
||||
|
||||
contexts
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FileStateCache {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取文件的当前 mtime(Unix 时间戳,秒)。
|
||||
/// 失败时返回 None。
|
||||
pub fn get_file_mtime(path: &str) -> Option<i64> {
|
||||
match std::fs::metadata(path) {
|
||||
Ok(meta) => match meta.modified() {
|
||||
Ok(time) => match time.duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(d) => Some(d.as_secs() as i64),
|
||||
Err(_) => None,
|
||||
},
|
||||
Err(_) => None,
|
||||
},
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_new_cache_is_empty() {
|
||||
let cache = FileStateCache::new();
|
||||
assert!(cache.is_empty());
|
||||
assert_eq!(cache.len(), 0);
|
||||
assert_eq!(cache.current_size_bytes(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_and_get() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "hello world".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(!cache.is_empty());
|
||||
|
||||
let entry = cache.get("/tmp/test.txt");
|
||||
assert!(entry.is_some());
|
||||
assert_eq!(entry.unwrap().content, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_path_normalization() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp//test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 规范化后的路径应该能命中
|
||||
assert!(cache.get("/tmp/test.txt").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains() {
|
||||
let mut cache = FileStateCache::new();
|
||||
assert!(!cache.contains("/tmp/test.txt"));
|
||||
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(cache.contains("/tmp/test.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(cache.remove("/tmp/test.txt"));
|
||||
assert!(cache.is_empty());
|
||||
assert!(!cache.remove("/tmp/test.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "a".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/b.txt",
|
||||
FileState {
|
||||
content: "b".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.clear();
|
||||
assert!(cache.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snapshot_sorted_by_timestamp_desc() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/old.txt",
|
||||
FileState {
|
||||
content: "old".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/new.txt",
|
||||
FileState {
|
||||
content: "new".to_string(),
|
||||
timestamp: 3000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/mid.txt",
|
||||
FileState {
|
||||
content: "mid".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
assert_eq!(snapshot.len(), 3);
|
||||
// 按 timestamp 降序
|
||||
assert_eq!(snapshot[0].1.timestamp, 3000);
|
||||
assert_eq!(snapshot[1].1.timestamp, 2000);
|
||||
assert_eq!(snapshot[2].1.timestamp, 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_snapshot() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "a".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/b.txt",
|
||||
FileState {
|
||||
content: "b".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
cache.clear();
|
||||
|
||||
cache.restore_from_snapshot(&snapshot, 1);
|
||||
assert_eq!(cache.len(), 1);
|
||||
// 应该恢复 timestamp 最高的
|
||||
assert!(cache.get("/tmp/b.txt").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_restore_context() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "file a content here".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
let contexts = FileStateCache::build_restore_context(&snapshot, 2);
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert!(contexts[0].contains("/tmp/a.txt"));
|
||||
assert!(contexts[0].contains("file a content here"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_snapshot_builds_no_context() {
|
||||
let contexts = FileStateCache::build_restore_context(&Vec::new(), 5);
|
||||
assert!(contexts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction_on_size() {
|
||||
// 创建一个小容量缓存(仅 200 bytes)
|
||||
let mut cache = FileStateCache::with_limits(100, 200);
|
||||
|
||||
// 写入 3 个 100 字节内容 → 应触发 LRU 淘汰
|
||||
cache.set(
|
||||
"/tmp/1.txt",
|
||||
FileState {
|
||||
content: "x".repeat(100),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/2.txt",
|
||||
FileState {
|
||||
content: "y".repeat(100),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 此时应该有 2 个条目(200 bytes)
|
||||
assert_eq!(cache.len(), 2);
|
||||
|
||||
// 写入第 3 个 → 应淘汰最旧的(1.txt)
|
||||
cache.set(
|
||||
"/tmp/3.txt",
|
||||
FileState {
|
||||
content: "z".repeat(100),
|
||||
timestamp: 3000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 旧条目被淘汰
|
||||
assert!(!cache.contains("/tmp/1.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_file_mtime() {
|
||||
// 创建临时文件
|
||||
let tmp = std::env::temp_dir().join("test_mtime.txt");
|
||||
std::fs::write(&tmp, "test").unwrap();
|
||||
|
||||
let mtime = get_file_mtime(&tmp.to_string_lossy());
|
||||
assert!(mtime.is_some());
|
||||
assert!(mtime.unwrap() > 0);
|
||||
|
||||
// 不存在的文件
|
||||
let mtime = get_file_mtime("/tmp/nonexistent_12345_xxx.txt");
|
||||
assert!(mtime.is_none());
|
||||
|
||||
std::fs::remove_file(&tmp).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_file_unchanged_stub_is_static() {
|
||||
assert!(FILE_UNCHANGED_STUB.contains("File unchanged"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// src/agent/runtime/finalize.rs
|
||||
//
|
||||
// 会话收尾:更新元信息、生成标题、保存指标、发送 Done 事件。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::info;
|
||||
|
||||
use super::{AgentMetrics, AgentStreamEvent};
|
||||
use crate::agent::hooks::{HookRegistry, SessionStopContext};
|
||||
use crate::agent::memory::extraction::{run_extraction, ExtractionConfig};
|
||||
use crate::agent::terminal::TurnTerminal;
|
||||
use crate::agent::trajectory::TrajectoryExporter;
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 完成会话回合:更新 session 元信息、触发 OnSessionStop hook、发送 Done、
|
||||
/// 导出 trajectory 数据。
|
||||
///
|
||||
/// `terminal` 参数允许传入实际的终止原因(取消/错误/超限等),
|
||||
/// 传入 `None` 时默认使用 `Completed`。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn finalize_turn(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
metrics: &AgentMetrics,
|
||||
question: &str,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
hook_registry: &HookRegistry,
|
||||
terminal: Option<TurnTerminal>,
|
||||
// Trajectory 导出所需参数
|
||||
library_dir: Option<&PathBuf>,
|
||||
llm_model: Option<&str>,
|
||||
system_prompt: Option<&str>,
|
||||
// 自动记忆提取所需
|
||||
app_state: Option<Arc<AppState>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let new_turn_count: i32 = sqlx::query_scalar(
|
||||
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
let metrics_json = serde_json::to_value(metrics).unwrap_or_default();
|
||||
info!(
|
||||
"[Finalize] 会话 {} 指标: steps={}, tools={:?}, compressions={}, duplicates={}",
|
||||
session_id,
|
||||
metrics.total_steps,
|
||||
metrics.tool_calls,
|
||||
metrics.compression_count,
|
||||
metrics.duplicate_detections
|
||||
);
|
||||
|
||||
// 首轮自动生成标题
|
||||
if new_turn_count <= 1 {
|
||||
let title = generate_title(question);
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET title = ?, turn_count = ?, metadata = ?, \
|
||||
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
|
||||
)
|
||||
.bind(&title)
|
||||
.bind(new_turn_count)
|
||||
.bind(&metrics_json)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
} else {
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET turn_count = ?, metadata = ?, \
|
||||
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
|
||||
)
|
||||
.bind(new_turn_count)
|
||||
.bind(&metrics_json)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
// OnSessionStop hook — 唯一触发点,使用传入的 terminal 或默认 Completed
|
||||
let default_terminal = TurnTerminal::Completed {
|
||||
session_id: session_id.to_string(),
|
||||
total_steps: metrics.total_steps,
|
||||
};
|
||||
let actual_terminal = terminal.unwrap_or(default_terminal);
|
||||
hook_registry
|
||||
.run_on_session_stop(&SessionStopContext {
|
||||
session_id: session_id.to_string(),
|
||||
terminal: &actual_terminal,
|
||||
total_steps: metrics.total_steps,
|
||||
})
|
||||
.await;
|
||||
|
||||
// ── Trajectory 导出(非阻塞,失败不影响主流程) ──
|
||||
if let (Some(lib_dir), Some(model), Some(sys_prompt)) = (library_dir, llm_model, system_prompt)
|
||||
{
|
||||
if let Err(e) = TrajectoryExporter::export(
|
||||
db,
|
||||
session_id,
|
||||
lib_dir,
|
||||
model,
|
||||
sys_prompt,
|
||||
metrics,
|
||||
Some(&actual_terminal),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("[Finalize] Trajectory 导出失败(非致命): {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// ── 自动记忆提取(fire-and-forget,不阻塞会话关闭) ──
|
||||
if let Some(app_state) = app_state {
|
||||
let extraction_config = ExtractionConfig::from_env();
|
||||
let session_id_owned = session_id.to_string();
|
||||
let mem_mgr = app_state.memory_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
run_extraction(app_state, session_id_owned, mem_mgr, extraction_config).await;
|
||||
});
|
||||
}
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 根据用户首条问题生成会话标题(截取前 50 字符)
|
||||
fn generate_title(question: &str) -> String {
|
||||
let chars: String = question.chars().take(50).collect();
|
||||
if question.len() > 50 {
|
||||
format!("{}...", chars)
|
||||
} else {
|
||||
chars
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,189 @@
|
||||
// src/agent/runtime/partitioner.rs
|
||||
//
|
||||
// 工具调用并发分区器。
|
||||
// 将准备好的工具调用按并发安全性分组成批处理。
|
||||
// 参考 Claude Code partitionToolCalls() 设计。
|
||||
//
|
||||
// 分区规则:
|
||||
// 1. 连续的并发安全工具放在同一个并行批次
|
||||
// 2. 非并发安全的工具单独一个批次(串行执行)
|
||||
// 3. 每个并行批次最多 max_concurrency 个工具
|
||||
|
||||
use tracing::debug;
|
||||
|
||||
use super::executor::PreparedCall;
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
|
||||
/// 一批工具调用
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolBatch {
|
||||
/// 该批次是否可以并行执行
|
||||
pub is_parallel: bool,
|
||||
/// 批次中的工具调用(按原始顺序)
|
||||
pub calls: Vec<PreparedCall>,
|
||||
}
|
||||
|
||||
/// 工具调用分区器
|
||||
pub struct ToolPartitioner {
|
||||
/// 最大并行度(并行批次中最多执行的工具数)
|
||||
max_concurrency: usize,
|
||||
}
|
||||
|
||||
impl ToolPartitioner {
|
||||
/// 创建分区器
|
||||
///
|
||||
/// `max_concurrency` 为 0 时使用默认值 10。
|
||||
pub fn new(max_concurrency: usize) -> Self {
|
||||
let concurrency = if max_concurrency == 0 {
|
||||
10
|
||||
} else {
|
||||
max_concurrency
|
||||
};
|
||||
ToolPartitioner {
|
||||
max_concurrency: concurrency,
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 prepared_calls 分区为顺序批处理。
|
||||
///
|
||||
/// 返回的批次列表按顺序依次执行:
|
||||
/// - `is_parallel = true` 的批次内工具可并发执行
|
||||
/// - `is_parallel = false` 的批次内只有一个工具,需串行执行
|
||||
pub fn partition(
|
||||
&self,
|
||||
prepared_calls: &[PreparedCall],
|
||||
tool_registry: &ToolRegistry,
|
||||
) -> Vec<ToolBatch> {
|
||||
let mut batches: Vec<ToolBatch> = Vec::new();
|
||||
|
||||
for prep in prepared_calls {
|
||||
let is_safe = self.is_concurrent_safe(prep, tool_registry);
|
||||
|
||||
// 尝试追加到上一个并行批次
|
||||
if is_safe {
|
||||
if let Some(last) = batches.last_mut() {
|
||||
if last.is_parallel && last.calls.len() < self.max_concurrency {
|
||||
last.calls.push(prep.clone());
|
||||
continue;
|
||||
}
|
||||
}
|
||||
// 新建并行批次
|
||||
batches.push(ToolBatch {
|
||||
is_parallel: true,
|
||||
calls: vec![prep.clone()],
|
||||
});
|
||||
} else {
|
||||
// 非并发安全,单独一个串行批次
|
||||
// 如果前一个也是串行,可以合并(但为了简单,保持每个非安全工具独立批次)
|
||||
batches.push(ToolBatch {
|
||||
is_parallel: false,
|
||||
calls: vec![prep.clone()],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
debug!(
|
||||
"[Partitioner] 分区完成: {} calls → {} batches",
|
||||
prepared_calls.len(),
|
||||
batches.len()
|
||||
);
|
||||
batches
|
||||
}
|
||||
|
||||
/// 判断单个工具调用是否可并发安全执行
|
||||
fn is_concurrent_safe(&self, prep: &PreparedCall, tool_registry: &ToolRegistry) -> bool {
|
||||
tool_registry
|
||||
.get(&prep.tool_name)
|
||||
.map(|t| t.is_concurrency_safe(&prep.args))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use super::*;
|
||||
use crate::agent::skills::SkillRegistry;
|
||||
|
||||
fn make_registry() -> ToolRegistry {
|
||||
ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
|
||||
"./skills",
|
||||
)))))
|
||||
}
|
||||
|
||||
fn make_prep(name: &str) -> PreparedCall {
|
||||
PreparedCall {
|
||||
tool_call_id: format!("call_{}", name),
|
||||
tool_name: name.to_string(),
|
||||
args: serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_concurrent_in_single_batch() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("rag_search"),
|
||||
make_prep("get_paper_metadata"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// search_papers, rag_search, get_paper_metadata 都是并发安全的
|
||||
// 它们应该在一个并行批次中
|
||||
assert_eq!(batches.len(), 1);
|
||||
assert!(batches[0].is_parallel);
|
||||
assert_eq!(batches[0].calls.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_download_splits_batch() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("download_paper"),
|
||||
make_prep("rag_search"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// search_papers (安全) → download_paper (不安全) → rag_search (安全)
|
||||
// 应分为 3 个批次
|
||||
assert_eq!(batches.len(), 3);
|
||||
assert!(batches[0].is_parallel); // search_papers
|
||||
assert_eq!(batches[0].calls.len(), 1);
|
||||
assert!(!batches[1].is_parallel); // download_paper (串行)
|
||||
assert_eq!(batches[1].calls.len(), 1);
|
||||
assert!(batches[2].is_parallel); // rag_search
|
||||
assert_eq!(batches[2].calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_max_concurrency_limit() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(2);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("rag_search"),
|
||||
make_prep("get_paper_metadata"),
|
||||
make_prep("load_skill"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// 4 个并发安全调用,max_concurrency=2 → 应分为 2 个并行批次
|
||||
assert_eq!(batches.len(), 2);
|
||||
assert!(batches[0].is_parallel);
|
||||
assert_eq!(batches[0].calls.len(), 2);
|
||||
assert!(batches[1].is_parallel);
|
||||
assert_eq!(batches[1].calls.len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// src/agent/runtime/permission.rs
|
||||
//
|
||||
// 权限检查管道 — 优先级排序的规则链。
|
||||
// 参考 Claude Code PermissionChecker 设计。
|
||||
//
|
||||
// 规则优先级(从高到低):
|
||||
// 1. Deny — 不可覆盖的拒绝
|
||||
// 2. Allow — 允许
|
||||
// 3. Ask — 需要用户确认
|
||||
//
|
||||
// 支持通配符 "*" 匹配所有工具。
|
||||
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::PermissionRule;
|
||||
|
||||
/// 权限检查结果
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionResult {
|
||||
/// 被拒绝(不可覆盖)
|
||||
Denied { reason: String },
|
||||
/// 允许
|
||||
Allowed,
|
||||
/// 需要用户确认
|
||||
AskUser { message: String },
|
||||
}
|
||||
|
||||
impl PermissionResult {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, PermissionResult::Allowed)
|
||||
}
|
||||
|
||||
pub fn is_denied(&self) -> bool {
|
||||
matches!(self, PermissionResult::Denied { .. })
|
||||
}
|
||||
}
|
||||
|
||||
/// 权限检查器 — 维护有序规则列表并逐条匹配
|
||||
pub struct PermissionChecker {
|
||||
rules: Vec<PermissionRule>,
|
||||
}
|
||||
|
||||
impl PermissionChecker {
|
||||
/// 创建空的检查器(默认允许所有)
|
||||
pub fn new() -> Self {
|
||||
PermissionChecker { rules: Vec::new() }
|
||||
}
|
||||
|
||||
/// 添加规则。先添加的优先级更高。
|
||||
pub fn add_rule(&mut self, rule: PermissionRule) {
|
||||
self.rules.push(rule);
|
||||
}
|
||||
|
||||
/// 检查指定工具是否可以执行。
|
||||
///
|
||||
/// 遍历规则列表,返回第一个匹配的决策。
|
||||
/// 无匹配规则时默认 Allow。
|
||||
pub fn check(&self, tool_name: &str) -> PermissionResult {
|
||||
for rule in &self.rules {
|
||||
match rule {
|
||||
PermissionRule::Deny {
|
||||
tool_name: name,
|
||||
reason,
|
||||
} if Self::matches(name, tool_name) => {
|
||||
info!("[Permission] 拒绝工具 {}: {}", tool_name, reason);
|
||||
return PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
}
|
||||
PermissionRule::Allow { tool_name: name } if Self::matches(name, tool_name) => {
|
||||
return PermissionResult::Allowed;
|
||||
}
|
||||
PermissionRule::Ask {
|
||||
tool_name: name,
|
||||
message,
|
||||
} if Self::matches(name, tool_name) => {
|
||||
return PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// 默认允许
|
||||
PermissionResult::Allowed
|
||||
}
|
||||
|
||||
/// 检查是否有明确拒绝该工具的规则
|
||||
pub fn is_denied(&self, tool_name: &str) -> bool {
|
||||
self.check(tool_name).is_denied()
|
||||
}
|
||||
|
||||
/// 规则名称匹配:支持精确匹配和通配符 "*"
|
||||
fn matches(pattern: &str, tool_name: &str) -> bool {
|
||||
pattern == "*" || pattern == tool_name
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PermissionChecker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::tools::PermissionRule;
|
||||
|
||||
#[test]
|
||||
fn test_empty_checker_allows_all() {
|
||||
let checker = PermissionChecker::new();
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
assert_eq!(checker.check("download_paper"), PermissionResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_wins_over_allow() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "download_paper".into(),
|
||||
reason: "blocked".into(),
|
||||
});
|
||||
checker.add_rule(PermissionRule::Allow {
|
||||
tool_name: "download_paper".into(),
|
||||
});
|
||||
|
||||
let result = checker.check("download_paper");
|
||||
assert!(result.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_deny_blocks_all() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "*".into(),
|
||||
reason: "all blocked".into(),
|
||||
});
|
||||
|
||||
assert!(checker.check("search_papers").is_denied());
|
||||
assert!(checker.check("download_paper").is_denied());
|
||||
assert!(checker.is_denied("rag_search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_returns_ask_user() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Ask {
|
||||
tool_name: "delete_paper".into(),
|
||||
message: "Are you sure?".into(),
|
||||
});
|
||||
|
||||
let result = checker.check("delete_paper");
|
||||
assert_eq!(
|
||||
result,
|
||||
PermissionResult::AskUser {
|
||||
message: "Are you sure?".into()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_match_defaults_to_allow() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "download_paper".into(),
|
||||
reason: "blocked".into(),
|
||||
});
|
||||
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rule_ordering_first_match_wins() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
// 先添加 Allow,后添加 Deny — Allow 先匹配
|
||||
checker.add_rule(PermissionRule::Allow {
|
||||
tool_name: "search_papers".into(),
|
||||
});
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "search_papers".into(),
|
||||
reason: "should not match".into(),
|
||||
});
|
||||
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
// src/agent/runtime/session.rs
|
||||
//
|
||||
// 会话生命周期管理:创建/恢复/验证 Agent 会话。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::clients::llm::LlmClient;
|
||||
|
||||
/// 会话信息摘要
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionInfo {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
}
|
||||
|
||||
/// 创建新会话或恢复已有会话。
|
||||
///
|
||||
/// 返回会话信息。如果指定的 session_id 不存在则返回错误。
|
||||
pub async fn create_or_resume_session(
|
||||
db: &SqlitePool,
|
||||
session_id: Option<String>,
|
||||
llm: &LlmClient,
|
||||
) -> anyhow::Result<SessionInfo> {
|
||||
match session_id {
|
||||
Some(id) => {
|
||||
// 验证会话存在且未被软删除
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
|
||||
}
|
||||
|
||||
// 计算当前轮次号
|
||||
let turn_index: i32 = sqlx::query_scalar(
|
||||
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?",
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: id,
|
||||
turn_index,
|
||||
})
|
||||
}
|
||||
None => {
|
||||
let new_id = uuid::Uuid::new_v4().to_string();
|
||||
sqlx::query("INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)")
|
||||
.bind(&new_id)
|
||||
.bind("")
|
||||
.bind(llm.model())
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: new_id,
|
||||
turn_index: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 加载会话的历史消息(供 LLM 上下文使用)。
|
||||
///
|
||||
/// `agent_name` 参数用于消息隔离:
|
||||
/// - `"lead"` — 只加载 Lead Agent 自己的消息(默认)
|
||||
/// - `"*"` — 加载所有 agent 的消息(调试/审计用)
|
||||
pub async fn load_history_for_llm(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
load_history_for_agent(db, session_id, "lead").await
|
||||
}
|
||||
|
||||
/// 加载指定 agent 的历史消息。
|
||||
pub async fn load_history_for_agent(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
agent_name: &str,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
let rows: Vec<(
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
)> = if agent_name == "*" {
|
||||
sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? AND agent_name = ? ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(agent_name)
|
||||
.fetch_all(db)
|
||||
.await?
|
||||
};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (role_str, content, tool_calls_json, tool_call_id, thought) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"system" => MessageRole::System,
|
||||
"user" => MessageRole::User,
|
||||
"assistant" => MessageRole::Assistant,
|
||||
"tool" => MessageRole::Tool,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let tool_calls: Option<Vec<crate::clients::llm::ToolCall>> =
|
||||
tool_calls_json.and_then(|json_str| serde_json::from_str(&json_str).ok());
|
||||
|
||||
messages.push(ChatMessage {
|
||||
role,
|
||||
content: if content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(content)
|
||||
},
|
||||
tool_call_id,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content: thought,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
// src/agent/runtime/streaming.rs
|
||||
//
|
||||
// LLM 流式响应处理:消费 chat_stream 返回的 StreamEvent 通道,
|
||||
// 支持并发取消检测,累积推理内容、文本增量和工具调用。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::error;
|
||||
|
||||
use crate::clients::llm::{
|
||||
ChatMessage, LlmClient, StreamEvent, TokenUsage, ToolCall, ToolDefinition,
|
||||
};
|
||||
|
||||
use super::AgentStreamEvent;
|
||||
|
||||
/// 流式处理的结果
|
||||
#[derive(Debug)]
|
||||
pub struct StreamOutput {
|
||||
pub content: String,
|
||||
pub reasoning: Option<String>,
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
pub is_tool_call_step: bool,
|
||||
pub status: StreamStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum StreamStatus {
|
||||
Success,
|
||||
Error(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// 处理 LLM 流式响应。
|
||||
///
|
||||
/// 使用 tokio::select! 在流式读取和取消信号之间竞速。
|
||||
/// 实时发送 Thought/TextDelta SSE 事件给前端。
|
||||
pub async fn process_llm_stream(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
tool_defs: &[ToolDefinition],
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
step: usize,
|
||||
session_id: &str,
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
) -> StreamOutput {
|
||||
// 1. 发起 LLM 流式调用
|
||||
let mut stream_rx = match llm.chat_stream(messages, tool_defs).await {
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[Streaming] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
return StreamOutput {
|
||||
content: String::new(),
|
||||
reasoning: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
is_tool_call_step: false,
|
||||
status: StreamStatus::Error(e.to_string()),
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
let mut accumulated_content = String::new();
|
||||
let mut accumulated_reasoning = String::new();
|
||||
let mut accumulated_tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut usage: Option<TokenUsage> = None;
|
||||
let mut is_tool_call_step = false;
|
||||
|
||||
// 2. 取消监视 future
|
||||
let sid = session_id.to_string();
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 3. tokio::select! 竞速:流式事件 vs 取消信号
|
||||
let mut cancel_pinned = Box::pin(cancel_fut);
|
||||
let mut error_msg = None;
|
||||
let mut was_cancelled = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
event_opt = stream_rx.recv() => {
|
||||
match event_opt {
|
||||
Some(event) => {
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
is_tool_call_step = true;
|
||||
accumulated_tool_calls = Some(tool_calls);
|
||||
}
|
||||
StreamEvent::ToolCallDelta { .. } => {}
|
||||
StreamEvent::Usage(u) => {
|
||||
usage = Some(u);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
break;
|
||||
}
|
||||
StreamEvent::Error(e) => {
|
||||
error_msg = Some(e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = &mut cancel_pinned => {
|
||||
was_cancelled = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 构建结果
|
||||
let status = if was_cancelled {
|
||||
StreamStatus::Cancelled
|
||||
} else if let Some(e) = error_msg {
|
||||
StreamStatus::Error(e)
|
||||
} else {
|
||||
StreamStatus::Success
|
||||
};
|
||||
|
||||
let reasoning = if accumulated_reasoning.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_reasoning)
|
||||
};
|
||||
|
||||
StreamOutput {
|
||||
content: accumulated_content,
|
||||
reasoning,
|
||||
tool_calls: accumulated_tool_calls,
|
||||
usage,
|
||||
is_tool_call_step,
|
||||
status,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
// src/agent/runtime/streaming_executor.rs
|
||||
//
|
||||
// 流式工具执行器。
|
||||
// 参考 Claude Code StreamingToolExecutor 设计。
|
||||
//
|
||||
// 当 LLM 流式输出 tool_use 块时,立即开始执行并发安全的工具。
|
||||
// 非并发安全的工具排队等待。结果按流中顺序 yield。
|
||||
//
|
||||
// 功能:
|
||||
// 1. 流式执行 — tool_use 到达时立即调度
|
||||
// 2. Sibling Abort — 副效应工具报错时中止兄弟并行执行
|
||||
// 3. Progress 流式 — 长时间操作可发送进度更新
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{broadcast, mpsc, oneshot};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::partitioner::ToolPartitioner;
|
||||
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 流式工具执行状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TrackedToolStatus {
|
||||
/// 工具调用已从 LLM 流中接收到
|
||||
Queued,
|
||||
/// 正在执行中
|
||||
Executing,
|
||||
/// 执行完成,等待 yield
|
||||
Completed,
|
||||
/// 结果已 yield 给调用方
|
||||
Yielded,
|
||||
}
|
||||
|
||||
/// 跟踪中的工具执行
|
||||
#[derive(Debug)]
|
||||
struct TrackedTool {
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
status: TrackedToolStatus,
|
||||
/// 执行完成后的输出
|
||||
output: Option<ToolOutput>,
|
||||
/// 取消通道(Sibling Abort 使用)
|
||||
#[allow(dead_code)]
|
||||
cancel_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
/// Sibling Abort 原因
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AbortReason {
|
||||
/// 兄弟工具出错触发的级联取消
|
||||
SiblingError { description: String },
|
||||
/// 用户主动中断
|
||||
UserInterrupted,
|
||||
}
|
||||
|
||||
/// 流式工具执行器
|
||||
pub struct StreamingToolExecutor {
|
||||
/// 所有跟踪中的工具
|
||||
tracked: Vec<TrackedTool>,
|
||||
/// 工具注册表
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
/// 并发分区器(保留用于未来并发策略优化)
|
||||
#[allow(dead_code)]
|
||||
partitioner: ToolPartitioner,
|
||||
/// 工具上下文
|
||||
tool_context: ToolContext,
|
||||
/// Sibling Abort 广播通道 (tx)
|
||||
abort_tx: broadcast::Sender<AbortReason>,
|
||||
/// Sibling Abort 广播通道 (rx)
|
||||
abort_rx: broadcast::Receiver<AbortReason>,
|
||||
/// 当前是否已发生错误(触发 sibling abort)
|
||||
has_errored: bool,
|
||||
/// 出错工具的描述
|
||||
errored_tool_desc: String,
|
||||
/// 下一个 stream_index
|
||||
next_index: usize,
|
||||
/// 最大工具输出字符数
|
||||
max_output_chars: usize,
|
||||
}
|
||||
|
||||
impl StreamingToolExecutor {
|
||||
/// 创建新的流式执行器
|
||||
pub fn new(
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
tool_context: ToolContext,
|
||||
max_concurrency: usize,
|
||||
max_output_chars: usize,
|
||||
) -> Self {
|
||||
let (abort_tx, abort_rx) = broadcast::channel(16);
|
||||
StreamingToolExecutor {
|
||||
tracked: Vec::new(),
|
||||
tool_registry,
|
||||
partitioner: ToolPartitioner::new(max_concurrency),
|
||||
tool_context,
|
||||
abort_tx,
|
||||
abort_rx,
|
||||
has_errored: false,
|
||||
errored_tool_desc: String::new(),
|
||||
next_index: 0,
|
||||
max_output_chars,
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取 abort 广播发送端(供外部注入取消信号)
|
||||
pub fn abort_sender(&self) -> broadcast::Sender<AbortReason> {
|
||||
self.abort_tx.clone()
|
||||
}
|
||||
|
||||
/// 当 LLM 流产生一个新的 tool_use 时调用。
|
||||
///
|
||||
/// 返回 true 表示该工具已立即开始执行(并发安全),false 表示排队。
|
||||
pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool {
|
||||
let _index = self.next_index;
|
||||
self.next_index += 1;
|
||||
|
||||
let is_concurrency_safe = self
|
||||
.tool_registry
|
||||
.get(&name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let (cancel_tx, _cancel_rx) = oneshot::channel();
|
||||
|
||||
let tool = TrackedTool {
|
||||
tool_call_id: call_id.clone(),
|
||||
tool_name: name.clone(),
|
||||
args: args.clone(),
|
||||
status: TrackedToolStatus::Queued,
|
||||
output: None,
|
||||
cancel_tx: Some(cancel_tx),
|
||||
};
|
||||
|
||||
self.tracked.push(tool);
|
||||
|
||||
if is_concurrency_safe {
|
||||
info!("[StreamingExecutor] 立即调度并发安全工具: {}", name);
|
||||
self.try_execute_pending();
|
||||
true
|
||||
} else {
|
||||
info!("[StreamingExecutor] 排队非并发安全工具: {}", name);
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM 流结束后调用,执行所有剩余排队工具。
|
||||
pub async fn flush(&mut self) {
|
||||
info!(
|
||||
"[StreamingExecutor] flush: {} tracked, {} queued",
|
||||
self.tracked.len(),
|
||||
self.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Queued)
|
||||
.count()
|
||||
);
|
||||
|
||||
// 将剩余排队的工具分批执行
|
||||
let queued: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in queued {
|
||||
self.execute_one(idx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 按流顺序获取下一个完成的结果(非阻塞)。
|
||||
pub fn next_result(&mut self) -> Option<(String, ToolOutput)> {
|
||||
for tool in &mut self.tracked {
|
||||
if tool.status == TrackedToolStatus::Completed {
|
||||
tool.status = TrackedToolStatus::Yielded;
|
||||
let output = tool
|
||||
.output
|
||||
.take()
|
||||
.unwrap_or_else(|| ToolOutput::error("工具执行异常:无输出"));
|
||||
return Some((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 是否有未 yield 的结果
|
||||
pub fn has_pending_results(&self) -> bool {
|
||||
self.tracked
|
||||
.iter()
|
||||
.any(|t| t.status == TrackedToolStatus::Completed)
|
||||
}
|
||||
|
||||
/// 是否有未完成的工具
|
||||
pub fn has_unfinished(&self) -> bool {
|
||||
self.tracked.iter().any(|t| {
|
||||
t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing
|
||||
})
|
||||
}
|
||||
|
||||
/// 获取所有已完成的结果(包括已 yield 和未 yield 的)
|
||||
pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> {
|
||||
let mut results = Vec::new();
|
||||
for tool in &mut self.tracked {
|
||||
if let Some(output) = tool.output.take() {
|
||||
results.push((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
// ── 内部方法 ──
|
||||
|
||||
/// 尝试执行可执行的排队工具
|
||||
fn try_execute_pending(&mut self) {
|
||||
// 简单策略:如果有正在执行的且它不是并发的,则不启动新的
|
||||
let has_executing = self
|
||||
.tracked
|
||||
.iter()
|
||||
.any(|t| t.status == TrackedToolStatus::Executing);
|
||||
|
||||
if !has_executing {
|
||||
// 启动所有排队的并发安全工具
|
||||
let indices: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in indices {
|
||||
// 在同步上下文中只能标记状态,实际执行在 async 上下文中
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行单个工具(内部辅助)
|
||||
async fn execute_one(&mut self, idx: usize) {
|
||||
if idx >= self.tracked.len() {
|
||||
return;
|
||||
}
|
||||
|
||||
// 检查 sibling abort
|
||||
if self.has_errored {
|
||||
if let Ok(reason) = self.abort_rx.try_recv() {
|
||||
let msg = match reason {
|
||||
AbortReason::SiblingError { ref description } => {
|
||||
format!("取消:并行工具 {} 出错,已级联取消", description)
|
||||
}
|
||||
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
|
||||
};
|
||||
self.tracked[idx].output = Some(ToolOutput::error(msg));
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
let tool_name = self.tracked[idx].tool_name.clone();
|
||||
|
||||
let output = match self.tool_registry.get(&tool_name) {
|
||||
Some(tool) => {
|
||||
let (progress_tx, _progress_rx) = mpsc::unbounded_channel();
|
||||
|
||||
let tool_args = self.tracked[idx].args.clone();
|
||||
let tool_fut =
|
||||
tool.execute_with_progress(tool_args, &self.tool_context, Some(&progress_tx));
|
||||
|
||||
tool_fut.await
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
// 检查是否需要触发 sibling abort
|
||||
if output.is_error {
|
||||
let causes_abort = self
|
||||
.tool_registry
|
||||
.get(&tool_name)
|
||||
.map(|t| t.causes_sibling_abort())
|
||||
.unwrap_or(false);
|
||||
|
||||
if causes_abort {
|
||||
warn!(
|
||||
"[StreamingExecutor] 工具 {} 出错,触发 sibling abort",
|
||||
tool_name
|
||||
);
|
||||
self.has_errored = true;
|
||||
self.errored_tool_desc = tool_name.clone();
|
||||
let _ = self.abort_tx.send(AbortReason::SiblingError {
|
||||
description: tool_name.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 截断输出
|
||||
let truncated = if output.content.len() > self.max_output_chars {
|
||||
let t: String = output.content.chars().take(self.max_output_chars).collect();
|
||||
ToolOutput {
|
||||
content: format!(
|
||||
"{}...\n[输出已截断,原始长度: {} 字符]",
|
||||
t,
|
||||
output.content.len()
|
||||
),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata,
|
||||
}
|
||||
} else {
|
||||
output
|
||||
};
|
||||
|
||||
self.tracked[idx].output = Some(truncated);
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
// src/agent/runtime/system_prompt.rs
|
||||
//
|
||||
// 模块化系统提示词组装 — 参考 Claude Code s10 System Prompt 设计。
|
||||
//
|
||||
// 将硬编码的提示词拆分为独立 section,运行时按需拼接。
|
||||
// 静态 section 在前以最大化 Anthropic prompt cache 命中率。
|
||||
|
||||
/// 系统提示词组装器
|
||||
pub struct SystemPrompt {
|
||||
sections: Vec<(&'static str, String)>,
|
||||
}
|
||||
|
||||
impl SystemPrompt {
|
||||
pub fn new() -> Self {
|
||||
SystemPrompt {
|
||||
sections: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 添加一个 section(先添加的排在前面)
|
||||
pub fn add_section(&mut self, name: &'static str, content: String) {
|
||||
self.sections.push((name, content));
|
||||
}
|
||||
|
||||
/// 组装最终的系统提示词(section 间用双换行分隔)
|
||||
pub fn assemble(&self) -> String {
|
||||
self.sections
|
||||
.iter()
|
||||
.map(|(_, content)| content.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
}
|
||||
|
||||
/// 是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.sections.is_empty()
|
||||
}
|
||||
|
||||
/// 获取 section 数量
|
||||
pub fn section_count(&self) -> usize {
|
||||
self.sections.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SystemPrompt {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 静态身份 section(始终加载,最大化 prompt cache 命中率)
|
||||
pub const IDENTITY_SECTION: &str = "\
|
||||
你是一位专业的天体物理学研究助手,具备丰富的天文学知识。";
|
||||
|
||||
/// 静态核心原则 section
|
||||
pub const PRINCIPLES_SECTION: &str = "\
|
||||
核心原则:
|
||||
1. 主动使用工具搜索最新文献,不要仅凭训练数据回答。
|
||||
2. 优先使用本地资源(get_paper_content / rag_search),必要时再检索新文献。
|
||||
3. 收集到足够信息后立即给出最终答案,避免无意义的重复工具调用。
|
||||
4. 回答时引用具体文献来源,使用 ADS bibcode 标注。
|
||||
5. 对于数学公式,使用标准 LaTeX 格式。
|
||||
6. 用中文回答,保持科学术语的准确性(可附带英文原文)。
|
||||
7. 对于复杂任务(如文献综述),调用 load_skill 获取方法论指引,再用 todo_write 制定计划。
|
||||
8. 如果某个工具调用失败,不要用相同参数重试,尝试换一种方式或工具。
|
||||
9. 任务状态会在每轮开始时从数据库恢复,请基于最新状态继续工作。";
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_assemble_empty() {
|
||||
let sp = SystemPrompt::new();
|
||||
assert_eq!(sp.assemble(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_assemble_multiple_sections() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
sp.add_section("a", "Section A".to_string());
|
||||
sp.add_section("b", "Section B".to_string());
|
||||
let result = sp.assemble();
|
||||
assert_eq!(result, "Section A\n\nSection B");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_static_sections_first() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
sp.add_section("identity", "I am".to_string());
|
||||
sp.add_section("dynamic", "Tools: ...".to_string());
|
||||
let result = sp.assemble();
|
||||
assert!(result.starts_with("I am"));
|
||||
assert!(result.contains("Tools: ..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_section_count() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
assert_eq!(sp.section_count(), 0);
|
||||
sp.add_section("a", "A".to_string());
|
||||
assert_eq!(sp.section_count(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
// src/agent/runtime/token_budget.rs
|
||||
//
|
||||
// Token 预算管理。
|
||||
// 参考 Claude Code TokenBudget 设计。
|
||||
// 软限制:接近上限时注入 nudging 消息提醒模型。
|
||||
// 硬限制:达到上限时触发强制压缩或终止。
|
||||
|
||||
/// Token 预算管理器。
|
||||
///
|
||||
/// 参考 Claude Code TokenBudget 设计,增加:
|
||||
/// - 多级渐进式 nudge(near_soft / over_soft / over_hard)
|
||||
/// - Diminishing returns 检测(防止模型在死循环中消耗预算)
|
||||
/// - Continuation 计数追踪
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TokenBudget {
|
||||
/// 软限制(触发 nudging 提醒)
|
||||
pub soft_limit: usize,
|
||||
/// 硬限制(触发强制动作)
|
||||
pub hard_limit: usize,
|
||||
/// 已消耗输入 tokens
|
||||
pub input_tokens_spent: usize,
|
||||
/// 已消耗输出 tokens
|
||||
pub output_tokens_spent: usize,
|
||||
/// 延续次数(每个 ReAct 步骤递增)
|
||||
pub continuation_count: usize,
|
||||
/// 上次检查时的总消耗(用于 diminishing returns 检测)
|
||||
last_total_spent: usize,
|
||||
/// 连续无进展次数
|
||||
consecutive_no_progress: usize,
|
||||
/// 是否已触发 diminishing returns
|
||||
pub diminishing_returns: bool,
|
||||
}
|
||||
|
||||
impl TokenBudget {
|
||||
/// 创建预算管理器
|
||||
pub fn new(soft_limit: usize, hard_limit: usize) -> Self {
|
||||
TokenBudget {
|
||||
soft_limit,
|
||||
hard_limit,
|
||||
input_tokens_spent: 0,
|
||||
output_tokens_spent: 0,
|
||||
continuation_count: 0,
|
||||
last_total_spent: 0,
|
||||
consecutive_no_progress: 0,
|
||||
diminishing_returns: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录输入 token 消耗
|
||||
pub fn spend_input(&mut self, tokens: usize) {
|
||||
self.input_tokens_spent += tokens;
|
||||
}
|
||||
|
||||
/// 记录输出 token 消耗
|
||||
pub fn spend_output(&mut self, tokens: usize) {
|
||||
self.output_tokens_spent += tokens;
|
||||
}
|
||||
|
||||
/// 总消耗
|
||||
pub fn total_spent(&self) -> usize {
|
||||
self.input_tokens_spent + self.output_tokens_spent
|
||||
}
|
||||
|
||||
/// 记录一次延续(每个 ReAct 步骤调用一次)。
|
||||
/// 同时检查 diminishing returns。
|
||||
pub fn record_continuation(&mut self) -> bool {
|
||||
self.continuation_count += 1;
|
||||
self.check_diminishing_returns_inner()
|
||||
}
|
||||
|
||||
/// 检测 diminishing returns — 模型在同一问题上打转而不产生实质进展。
|
||||
///
|
||||
/// 触发条件:3 次以上延续,且连续 2 次检查的 token 增量 < 500。
|
||||
/// 返回 true 表示已检测到无进展循环。
|
||||
pub fn check_diminishing_returns(&mut self) -> bool {
|
||||
self.check_diminishing_returns_inner()
|
||||
}
|
||||
|
||||
fn check_diminishing_returns_inner(&mut self) -> bool {
|
||||
if self.diminishing_returns {
|
||||
return true; // 已触发过,保持状态
|
||||
}
|
||||
if self.continuation_count < 3 {
|
||||
return false;
|
||||
}
|
||||
let delta = self.total_spent().saturating_sub(self.last_total_spent);
|
||||
self.last_total_spent = self.total_spent();
|
||||
if delta < 500 {
|
||||
self.consecutive_no_progress += 1;
|
||||
if self.consecutive_no_progress >= 2 {
|
||||
self.diminishing_returns = true;
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.consecutive_no_progress = 0;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 是否接近软限制(超过 80%)
|
||||
pub fn near_soft_limit(&self) -> bool {
|
||||
if self.soft_limit == 0 {
|
||||
return false;
|
||||
}
|
||||
self.total_spent() >= self.soft_limit * 8 / 10
|
||||
}
|
||||
|
||||
/// 是否超过软限制
|
||||
pub fn over_soft_limit(&self) -> bool {
|
||||
self.total_spent() >= self.soft_limit
|
||||
}
|
||||
|
||||
/// 是否超过硬限制
|
||||
pub fn over_hard_limit(&self) -> bool {
|
||||
self.total_spent() >= self.hard_limit
|
||||
}
|
||||
|
||||
/// 已使用预算的百分比(相对于软限制)
|
||||
pub fn usage_pct(&self) -> u32 {
|
||||
if self.soft_limit == 0 {
|
||||
return 0;
|
||||
}
|
||||
(self.total_spent() * 100 / self.soft_limit) as u32
|
||||
}
|
||||
|
||||
/// 剩余可用 tokens(硬限制 - 已消耗)
|
||||
pub fn remaining(&self) -> usize {
|
||||
self.hard_limit.saturating_sub(self.total_spent())
|
||||
}
|
||||
|
||||
/// 生成渐进式 nudging 提醒消息。
|
||||
///
|
||||
/// 三级:
|
||||
/// - near_soft (80-99%): 温和提醒
|
||||
/// - over_soft (100-硬): 明确警告
|
||||
/// - over_hard (>硬限制): 强制完成
|
||||
/// - diminishing_returns: 要求最终答案
|
||||
pub fn nudge_message(&self) -> Option<String> {
|
||||
if self.diminishing_returns {
|
||||
return Some(
|
||||
"⚠️ 已检测到重复操作模式 — 后续步骤未产生新信息。\
|
||||
请基于已收集的全部信息直接给出最终答案,不要再调用工具。"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
if self.over_hard_limit() {
|
||||
Some(format!(
|
||||
"🔴 Token 预算已耗尽({}/{} tokens, {}%)。\
|
||||
请立即总结当前发现并给出最终答案,不要再调用任何工具。",
|
||||
self.total_spent(),
|
||||
self.hard_limit,
|
||||
self.usage_pct()
|
||||
))
|
||||
} else if self.over_soft_limit() {
|
||||
let pct = self.usage_pct();
|
||||
Some(format!(
|
||||
"🟡 Token 预算警告:已使用 {}/{} tokens ({}%)。\
|
||||
请尽快总结关键发现并给出最终答案。如非必要,不要再调用工具。",
|
||||
self.total_spent(),
|
||||
self.soft_limit,
|
||||
pct
|
||||
))
|
||||
} else if self.near_soft_limit() {
|
||||
let pct = self.usage_pct();
|
||||
Some(format!(
|
||||
"💡 Token 预算提示:已使用 {}/{} tokens ({}%)。\
|
||||
请注意控制后续步骤的深度,优先处理最重要的发现。",
|
||||
self.total_spent(),
|
||||
self.soft_limit,
|
||||
pct
|
||||
))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 将硬限制提升到指定值(用于 error recovery 中的 escalate 步骤)
|
||||
pub fn escalate_hard_limit(&mut self, new_limit: usize) {
|
||||
self.hard_limit = new_limit;
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TokenBudget {
|
||||
fn default() -> Self {
|
||||
Self::new(32_000, 40_000)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_spend_tracking() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(500);
|
||||
budget.spend_output(300);
|
||||
assert_eq!(budget.total_spent(), 800);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_near_soft_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
assert!(!budget.near_soft_limit());
|
||||
budget.spend_input(850); // 85% > 80%
|
||||
assert!(budget.near_soft_limit());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_over_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(2100);
|
||||
assert!(budget.over_hard_limit());
|
||||
assert_eq!(budget.remaining(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nudge_message_at_soft_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(1000); // exactly at soft limit
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("Token 预算警告"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nudge_message_at_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(2000); // at hard limit
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("已耗尽"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_nudge_when_under_limit() {
|
||||
let budget = TokenBudget::new(1000, 2000);
|
||||
assert!(budget.nudge_message().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_escalate_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.escalate_hard_limit(64000);
|
||||
assert_eq!(budget.hard_limit, 64000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_budget() {
|
||||
let budget = TokenBudget::default();
|
||||
assert_eq!(budget.soft_limit, 32_000);
|
||||
assert_eq!(budget.hard_limit, 40_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_near_soft_limit_nudge() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(850); // 85% — near soft
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("Token 预算提示"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_not_triggered_early() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
// < 3 continuations — should not trigger
|
||||
budget.record_continuation();
|
||||
assert!(!budget.diminishing_returns);
|
||||
budget.record_continuation();
|
||||
assert!(!budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_triggers_after_stagnation() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
// First 3 continuations establish baseline (all with 0 spending)
|
||||
// After the 3rd, consecutive_no_progress becomes 1 (delta=0 < 500)
|
||||
for _ in 0..3 {
|
||||
budget.record_continuation();
|
||||
}
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
// 4th continuation with tiny spending — 2nd consecutive <500
|
||||
budget.spend_input(100);
|
||||
budget.record_continuation();
|
||||
// consecutive_no_progress is now 2 → triggered
|
||||
assert!(budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_resets_on_progress() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
// Establish baseline with spending BEFORE the 3-continuation threshold
|
||||
budget.spend_input(2000);
|
||||
for _ in 0..3 {
|
||||
budget.record_continuation();
|
||||
}
|
||||
// delta = total - last_total. After the first check, last_total is set.
|
||||
// total_spent=2000, last_total=2000 after first continuation ≥ 3
|
||||
|
||||
budget.spend_input(100); // total=2100
|
||||
budget.record_continuation(); // delta=100 < 500, cons=1
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
budget.spend_input(600); // total=2700
|
||||
budget.record_continuation(); // delta=600 >= 500, cons resets to 0
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
budget.spend_input(100); // total=2800
|
||||
budget.record_continuation(); // delta=100 < 500, cons=1
|
||||
assert!(!budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_nudge_message() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
budget.diminishing_returns = true;
|
||||
let msg = budget.nudge_message().unwrap();
|
||||
assert!(msg.contains("重复操作"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_pct() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(500);
|
||||
assert_eq!(budget.usage_pct(), 50);
|
||||
budget.spend_output(300);
|
||||
assert_eq!(budget.usage_pct(), 80);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,846 @@
|
||||
// src/agent/skills.rs
|
||||
//
|
||||
// 两层技能加载系统(参考 Claude Code src/skills/ + src/tools/SkillTool/ 设计):
|
||||
// Layer 1 — system-reminder 注入:每轮动态列出 skill 名称(~20 tokens/skill)
|
||||
// Layer 2 — LoadSkillTool:LLM 按需调用,注入完整 skill 内容(~2000 tokens/skill)
|
||||
//
|
||||
// Skill 文件格式(对齐 Claude Code 的目录约定):
|
||||
// skills/{skill-name}/SKILL.md ← 必须是目录 + SKILL.md
|
||||
//
|
||||
// SKILL.md 内容(Markdown + YAML frontmatter):
|
||||
// ---
|
||||
// name: skill-name
|
||||
// description: 一句话描述
|
||||
// context: inline | fork
|
||||
// allowed-tools:
|
||||
// - bash
|
||||
// - read
|
||||
// when_to_use: 何时自动触发
|
||||
// model: haiku | sonnet | opus | inherit
|
||||
// disable-model-invocation: false
|
||||
// user-invocable: true
|
||||
// paths:
|
||||
// - "*.rs"
|
||||
// ---
|
||||
//
|
||||
// # Skill 正文
|
||||
// 详细内容...
|
||||
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::SystemTime;
|
||||
use tracing::{info, warn};
|
||||
|
||||
// ── Frontmatter ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Skill 的 YAML frontmatter 结构(serde_yaml 解析)。
|
||||
/// 仅声明与平台相关的字段;未知字段自动忽略。
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
pub struct SkillFrontmatter {
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub description: Option<String>,
|
||||
/// 执行模式:inline | fork
|
||||
#[serde(default)]
|
||||
pub context: Option<String>,
|
||||
/// 工具白名单
|
||||
#[serde(rename = "allowed-tools", default)]
|
||||
pub allowed_tools: Option<Vec<String>>,
|
||||
/// 推荐模型
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
/// 参数提示
|
||||
#[serde(rename = "argument-hint", default)]
|
||||
pub argument_hint: Option<String>,
|
||||
/// 使用场景说明
|
||||
#[serde(rename = "when_to_use", default)]
|
||||
pub when_to_use: Option<String>,
|
||||
/// 禁止模型通过 Skill tool 自动调用
|
||||
#[serde(rename = "disable-model-invocation", default)]
|
||||
pub disable_model_invocation: Option<bool>,
|
||||
/// 用户是否可通过 /skill-name 手动调用
|
||||
#[serde(rename = "user-invocable", default)]
|
||||
pub user_invocable: Option<bool>,
|
||||
/// 版本号
|
||||
#[serde(default)]
|
||||
pub version: Option<String>,
|
||||
/// 条件激活的 glob 模式
|
||||
#[serde(default)]
|
||||
pub paths: Option<Vec<String>>,
|
||||
/// fork 模式下的 agent 类型
|
||||
#[serde(default)]
|
||||
pub agent: Option<String>,
|
||||
/// fork 模式下的 effort 级别
|
||||
#[serde(default)]
|
||||
pub effort: Option<String>,
|
||||
}
|
||||
|
||||
impl SkillFrontmatter {
|
||||
/// 校验必填/推荐字段,返回警告列表
|
||||
pub fn validate(&self, skill_name: &str) -> Vec<String> {
|
||||
let mut warnings = Vec::new();
|
||||
if self.description.is_none() {
|
||||
warnings.push(format!("Skill '{}' 缺少 description 字段", skill_name));
|
||||
}
|
||||
if let Some(ref ctx) = self.context {
|
||||
if ctx != "inline" && ctx != "fork" {
|
||||
warnings.push(format!(
|
||||
"Skill '{}' 的 context 值无效 '{}',应为 inline 或 fork",
|
||||
skill_name, ctx
|
||||
));
|
||||
}
|
||||
}
|
||||
warnings
|
||||
}
|
||||
}
|
||||
|
||||
// ── Skill Data Structures ──────────────────────────────────────────────────
|
||||
|
||||
/// Skill 元信息(Layer 1:出现在 skill 列表中)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SkillMeta {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
/// 执行模式(inline / fork)
|
||||
pub context: Option<String>,
|
||||
/// 工具白名单(空 Vec 表示无限制)
|
||||
pub allowed_tools: Vec<String>,
|
||||
/// 使用场景说明
|
||||
pub when_to_use: Option<String>,
|
||||
/// 是否禁止模型通过 Skill tool 调用
|
||||
pub disable_model_invocation: bool,
|
||||
/// 是否允许用户通过 /skill-name 手动调用
|
||||
pub user_invocable: bool,
|
||||
/// 条件激活的 glob 模式(空 Vec 表示始终激活)
|
||||
pub paths: Vec<String>,
|
||||
}
|
||||
|
||||
/// 完整的 Skill(Layer 2:LLM 调用 load_skill 时注入)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Skill {
|
||||
pub meta: SkillMeta,
|
||||
pub body: String,
|
||||
/// Skill 所在目录,用于 ${SKILL_DIR} 变量替换
|
||||
pub skill_dir: PathBuf,
|
||||
}
|
||||
|
||||
// ── Skill Registry (Caching Layer) ─────────────────────────────────────────
|
||||
|
||||
/// Skill 使用统计
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SkillUsageStat {
|
||||
pub invoke_count: u64,
|
||||
pub last_used_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
/// Skill 注册表 — 缓存已加载的 skills,支持 mtime 增量刷新。
|
||||
///
|
||||
/// 使用方式:
|
||||
/// ```ignore
|
||||
/// let registry = SkillRegistry::new(skills_dir);
|
||||
/// registry.refresh()?;
|
||||
/// let reminder = registry.build_reminder();
|
||||
/// ```
|
||||
#[derive(Debug)]
|
||||
pub struct SkillRegistry {
|
||||
skills_dir: PathBuf,
|
||||
skills: Vec<Skill>,
|
||||
/// 上次扫描时 skills_dir 的 mtime(用于增量刷新)
|
||||
last_scan_mtime: Option<SystemTime>,
|
||||
/// 使用统计(按 skill name 索引)
|
||||
usage_stats: HashMap<String, SkillUsageStat>,
|
||||
}
|
||||
|
||||
impl SkillRegistry {
|
||||
/// 创建新的 skill 注册表(不执行初始扫描,调用 `refresh()` 触发)
|
||||
pub fn new(skills_dir: PathBuf) -> Self {
|
||||
SkillRegistry {
|
||||
skills_dir,
|
||||
skills: Vec::new(),
|
||||
last_scan_mtime: None,
|
||||
usage_stats: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查是否需要重新扫描(目录 mtime 变化或首次加载)
|
||||
pub fn needs_refresh(&self) -> bool {
|
||||
match dir_modified_time(&self.skills_dir) {
|
||||
Some(current_mtime) => match self.last_scan_mtime {
|
||||
Some(last) => current_mtime > last,
|
||||
None => true,
|
||||
},
|
||||
None => !self.skills.is_empty(), // 目录消失但还有缓存 → 保持缓存
|
||||
}
|
||||
}
|
||||
|
||||
/// 扫描 skills 目录并加载/更新所有 skills。
|
||||
/// 始终执行完整重载(简单可靠,skill 数量少时成本可忽略)。
|
||||
pub fn refresh(&mut self) {
|
||||
let current_mtime = dir_modified_time(&self.skills_dir);
|
||||
|
||||
let mut new_skills = Vec::new();
|
||||
for (dir_name, skill_md_path) in discover_skill_dirs(&self.skills_dir) {
|
||||
match load_skill_from_path(&skill_md_path, &dir_name) {
|
||||
Ok(skill) => new_skills.push(skill),
|
||||
Err(e) => {
|
||||
warn!("[SkillRegistry] 加载 skill '{}' 失败: {}", dir_name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 按使用频率排序:常用 skill 排前面
|
||||
new_skills.sort_by(|a, b| {
|
||||
let score_a = self.usage_score(&a.meta.name);
|
||||
let score_b = self.usage_score(&b.meta.name);
|
||||
// 降序排列(高分在前)
|
||||
score_b
|
||||
.partial_cmp(&score_a)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
let count = new_skills.len();
|
||||
self.skills = new_skills;
|
||||
self.last_scan_mtime = current_mtime;
|
||||
info!("[SkillRegistry] 已刷新 {} 个 skill", count);
|
||||
}
|
||||
|
||||
/// 获取所有 skill 的元信息列表
|
||||
pub fn list_skills(&self) -> Vec<SkillMeta> {
|
||||
self.skills.iter().map(|s| s.meta.clone()).collect()
|
||||
}
|
||||
|
||||
/// 按名称获取完整 skill
|
||||
pub fn get_skill(&self, name: &str) -> Option<&Skill> {
|
||||
self.skills.iter().find(|s| s.meta.name == name)
|
||||
}
|
||||
|
||||
/// 检查是否有可用 skills
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.skills.is_empty()
|
||||
}
|
||||
|
||||
/// 技能数量
|
||||
pub fn len(&self) -> usize {
|
||||
self.skills.len()
|
||||
}
|
||||
|
||||
// ── 使用统计 ──
|
||||
|
||||
/// 计算 skill 的使用评分(指数衰减,7 天半衰期)
|
||||
fn usage_score(&self, name: &str) -> f64 {
|
||||
match self.usage_stats.get(name) {
|
||||
Some(stat) => {
|
||||
let count_weight = (stat.invoke_count as f64).ln_1p(); // log(1 + count)
|
||||
let recency_weight = match stat.last_used_at {
|
||||
Some(last) => {
|
||||
let age_hours = chrono::Utc::now()
|
||||
.signed_duration_since(last)
|
||||
.num_hours()
|
||||
.max(0) as f64;
|
||||
// 7 天半衰期
|
||||
0.5_f64.powf(age_hours / (7.0 * 24.0))
|
||||
}
|
||||
None => 0.1,
|
||||
};
|
||||
count_weight * recency_weight
|
||||
}
|
||||
None => 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录一次 skill 调用
|
||||
pub fn record_usage(&mut self, name: &str) {
|
||||
let stat = self.usage_stats.entry(name.to_string()).or_default();
|
||||
stat.invoke_count += 1;
|
||||
stat.last_used_at = Some(chrono::Utc::now());
|
||||
}
|
||||
|
||||
/// 获取所有使用统计的快照
|
||||
pub fn usage_stats(&self) -> &HashMap<String, SkillUsageStat> {
|
||||
&self.usage_stats
|
||||
}
|
||||
|
||||
// ── System Prompt 构建 ──
|
||||
|
||||
/// 构建 skill 列表的 system-reminder 消息(Layer 1,参考 Claude Code)。
|
||||
/// 使用结构化 XML 标签,包含名称和描述。
|
||||
pub fn build_reminder(&self) -> Option<String> {
|
||||
if self.skills.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let invocable: Vec<&Skill> = self
|
||||
.skills
|
||||
.iter()
|
||||
.filter(|s| !s.meta.disable_model_invocation && s.meta.user_invocable)
|
||||
.collect();
|
||||
|
||||
if invocable.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut lines = vec![
|
||||
"<system-reminder>".to_string(),
|
||||
"The following skills are available for use with the Skill tool:".to_string(),
|
||||
];
|
||||
|
||||
for skill in &invocable {
|
||||
let desc = match &skill.meta.when_to_use {
|
||||
Some(wtu) => format!("{} - {}", skill.meta.description, wtu),
|
||||
None => skill.meta.description.clone(),
|
||||
};
|
||||
let context_marker = match skill.meta.context.as_deref() {
|
||||
Some("fork") => " [fork]",
|
||||
_ => "",
|
||||
};
|
||||
lines.push(format!("- {}: {}{}", skill.meta.name, desc, context_marker));
|
||||
}
|
||||
|
||||
lines.push(
|
||||
"When a skill matches the user's request, invoke load_skill BEFORE generating any other response about the task.".to_string(),
|
||||
);
|
||||
lines.push(
|
||||
"If you see a <command-name> tag in the current conversation turn, the skill has ALREADY been loaded - follow the instructions directly instead of calling load_skill again.".to_string(),
|
||||
);
|
||||
lines.push("</system-reminder>".to_string());
|
||||
|
||||
Some(lines.join("\n"))
|
||||
}
|
||||
|
||||
/// 构建 LoadSkillTool 的 description(动态生成,列出可用 skills)
|
||||
pub fn build_tool_description(&self) -> String {
|
||||
if self.skills.is_empty() {
|
||||
return "加载指定的领域技能完整内容。当前没有可用的技能。".to_string();
|
||||
}
|
||||
|
||||
let invocable: Vec<&Skill> = self
|
||||
.skills
|
||||
.iter()
|
||||
.filter(|s| !s.meta.disable_model_invocation)
|
||||
.collect();
|
||||
|
||||
if invocable.is_empty() {
|
||||
return "加载指定的领域技能完整内容。当前没有可用的技能。".to_string();
|
||||
}
|
||||
|
||||
let mut desc = String::from("加载指定的领域技能完整内容。可用技能:\n");
|
||||
for skill in &invocable {
|
||||
let mode = match skill.meta.context.as_deref() {
|
||||
Some("fork") => "[子代理执行] ",
|
||||
_ => "",
|
||||
};
|
||||
let wtu = match &skill.meta.when_to_use {
|
||||
Some(w) => format!(" - {}", w),
|
||||
None => String::new(),
|
||||
};
|
||||
desc.push_str(&format!(
|
||||
"- {}: {}{}{}\n",
|
||||
skill.meta.name, mode, skill.meta.description, wtu
|
||||
));
|
||||
}
|
||||
desc
|
||||
}
|
||||
|
||||
// ── 文件监听 (Hot Reload) ──
|
||||
|
||||
/// 启动文件监听器,在 skills 目录变更时自动刷新缓存。
|
||||
///
|
||||
/// 返回一个 `JoinHandle`,调用方可以 `await` 它(通常运行到程序退出)。
|
||||
/// 内部使用 debounce:300ms 内的连续变更合并为一次刷新。
|
||||
#[cfg(not(test))]
|
||||
pub fn start_watcher(self_arc: Arc<RwLock<SkillRegistry>>) -> std::thread::JoinHandle<()> {
|
||||
use notify::{RecursiveMode, Watcher};
|
||||
use std::time::Duration;
|
||||
|
||||
let skills_dir = match self_arc.read() {
|
||||
Ok(r) => r.skills_dir.clone(),
|
||||
Err(_) => {
|
||||
warn!("[SkillRegistry] RwLock 异常,无法启动文件监视器");
|
||||
return std::thread::spawn(|| {});
|
||||
}
|
||||
};
|
||||
|
||||
std::thread::spawn(move || {
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
|
||||
let mut watcher =
|
||||
match notify::recommended_watcher(move |res: notify::Result<notify::Event>| {
|
||||
if let Ok(event) = res {
|
||||
// 只关心 SKILL.md 相关的变更
|
||||
let is_skill_change = event
|
||||
.paths
|
||||
.iter()
|
||||
.any(|p| p.file_name().map(|n| n == "SKILL.md").unwrap_or(false));
|
||||
if is_skill_change {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
}
|
||||
}) {
|
||||
Ok(w) => w,
|
||||
Err(e) => {
|
||||
warn!("[SkillRegistry] 无法创建文件监听器: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = watcher.watch(&skills_dir, RecursiveMode::Recursive) {
|
||||
warn!("[SkillRegistry] 无法监听 skills 目录: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
info!("[SkillRegistry] 文件监听已启动: {}", skills_dir.display());
|
||||
|
||||
// 300ms debounce:收集快速连续的事件
|
||||
while let Ok(()) = rx.recv() {
|
||||
// 等待 debounce 窗口
|
||||
while rx.recv_timeout(Duration::from_millis(300)).is_ok() {}
|
||||
info!("[SkillRegistry] 检测到 skill 文件变更,自动刷新");
|
||||
if let Ok(mut registry) = self_arc.write() {
|
||||
let reg: &mut SkillRegistry = &mut registry;
|
||||
reg.refresh();
|
||||
}
|
||||
}
|
||||
// Channel closed, watcher dropped
|
||||
})
|
||||
}
|
||||
|
||||
/// 文件监听器的空实现(测试模式下不启动线程)
|
||||
#[cfg(test)]
|
||||
pub fn start_watcher(_self_arc: Arc<RwLock<SkillRegistry>>) -> std::thread::JoinHandle<()> {
|
||||
std::thread::spawn(|| {})
|
||||
}
|
||||
|
||||
// ── 条件 Skill (Paths-based Activation) ──
|
||||
|
||||
/// 根据访问的文件路径激活匹配的条件 skill。
|
||||
///
|
||||
/// 条件 skill 在其 `paths` frontmatter 中声明了 glob 模式。
|
||||
/// 当 Agent 访问(Read/Edit/Grep)匹配文件时调用此方法将其激活。
|
||||
///
|
||||
/// 返回新激活的 skill 名称列表。
|
||||
pub fn activate_conditional_for_paths(&mut self, file_paths: &[&str]) -> Vec<String> {
|
||||
let mut activated = Vec::new();
|
||||
|
||||
for skill in &mut self.skills {
|
||||
// 只处理有 paths 限制且当前未激活的
|
||||
if skill.meta.paths.is_empty() || !skill.meta.disable_model_invocation {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 检查是否有任何文件路径匹配
|
||||
let matches = file_paths.iter().any(|fp| {
|
||||
skill
|
||||
.meta
|
||||
.paths
|
||||
.iter()
|
||||
.any(|pattern| glob_match_simple(pattern, fp))
|
||||
});
|
||||
|
||||
if matches {
|
||||
skill.meta.disable_model_invocation = false;
|
||||
activated.push(skill.meta.name.clone());
|
||||
info!(
|
||||
"[SkillRegistry] 条件 skill '{}' 已激活 (paths: {:?})",
|
||||
skill.meta.name, skill.meta.paths
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
activated
|
||||
}
|
||||
|
||||
/// 获取与给定文件路径匹配的所有 skill(用于 LLM 上下文提示)。
|
||||
pub fn matching_skills_for_paths(&self, file_paths: &[&str]) -> Vec<SkillMeta> {
|
||||
self.skills
|
||||
.iter()
|
||||
.filter(|s| {
|
||||
!s.meta.paths.is_empty()
|
||||
&& file_paths.iter().any(|fp| {
|
||||
s.meta
|
||||
.paths
|
||||
.iter()
|
||||
.any(|pattern| glob_match_simple(pattern, fp))
|
||||
})
|
||||
})
|
||||
.map(|s| s.meta.clone())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
// ── Glob 匹配 (简化实现,避免引入完整 glob 库的运行时开销) ──
|
||||
|
||||
/// 简化的 glob 模式匹配。
|
||||
///
|
||||
/// 支持的语法:
|
||||
/// - `*` 匹配任意非 '/' 字符序列
|
||||
/// - `**` 匹配任意字符(含 '/')
|
||||
/// - `?` 匹配单个非 '/' 字符
|
||||
/// - 其他字符按字面匹配
|
||||
fn glob_match_simple(pattern: &str, path: &str) -> bool {
|
||||
// 标准化:统一使用 '/' 作为路径分隔符
|
||||
let pattern = pattern.replace('\\', "/");
|
||||
let path = path.replace('\\', "/");
|
||||
|
||||
// 使用 glob crate 进行匹配
|
||||
// 降级方案:简单的后缀/前缀匹配
|
||||
if pattern.contains('*') || pattern.contains('?') {
|
||||
// 尝试使用 glob crate
|
||||
match glob::Pattern::new(&pattern) {
|
||||
Ok(pat) => pat.matches(&path),
|
||||
Err(_) => {
|
||||
// 降级:简单包含匹配
|
||||
let simple = pattern.replace(['*', '?'], "");
|
||||
path.contains(&simple)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 无通配符:精确匹配文件名或后缀
|
||||
path == pattern || path.ends_with(&format!("/{}", pattern))
|
||||
}
|
||||
}
|
||||
|
||||
// ── File I/O Helpers ───────────────────────────────────────────────────────
|
||||
|
||||
/// 获取目录的修改时间
|
||||
fn dir_modified_time(path: &Path) -> Option<SystemTime> {
|
||||
std::fs::metadata(path).ok().and_then(|m| m.modified().ok())
|
||||
}
|
||||
|
||||
/// 扫描 skills 目录,查找所有 `{name}/SKILL.md` 子目录。
|
||||
fn discover_skill_dirs(skills_dir: &Path) -> Vec<(String, PathBuf)> {
|
||||
let mut skills = Vec::new();
|
||||
|
||||
let entries = match std::fs::read_dir(skills_dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return skills,
|
||||
};
|
||||
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let skill_md = path.join("SKILL.md");
|
||||
if skill_md.exists() {
|
||||
let name = path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
skills.push((name, skill_md));
|
||||
}
|
||||
}
|
||||
|
||||
skills
|
||||
}
|
||||
|
||||
/// 解析 SKILL.md 文件:提取 YAML frontmatter 和 Markdown 正文。
|
||||
///
|
||||
/// Frontmatter 格式:
|
||||
/// ---
|
||||
/// key: value
|
||||
/// ---
|
||||
/// 正文
|
||||
fn parse_skill_file(raw: &str) -> Result<(SkillFrontmatter, String), String> {
|
||||
let content = raw.trim();
|
||||
|
||||
if let Some(rest) = content.strip_prefix("---") {
|
||||
// 查找闭合的 ---
|
||||
if let Some((fm_text, body_text)) = rest.split_once("\n---") {
|
||||
let frontmatter: SkillFrontmatter = serde_yaml::from_str(fm_text)
|
||||
.map_err(|e| format!("YAML frontmatter 解析失败: {}", e))?;
|
||||
let body = body_text.trim().to_string();
|
||||
return Ok((frontmatter, body));
|
||||
}
|
||||
}
|
||||
|
||||
// 没有 frontmatter,整个内容作为正文
|
||||
Ok((SkillFrontmatter::default(), content.to_string()))
|
||||
}
|
||||
|
||||
/// 从文件路径加载完整的 Skill
|
||||
fn load_skill_from_path(skill_md_path: &Path, dir_name: &str) -> Result<Skill, String> {
|
||||
let raw = std::fs::read_to_string(skill_md_path).map_err(|e| format!("无法读取文件: {}", e))?;
|
||||
|
||||
let (frontmatter, body) = parse_skill_file(&raw)?;
|
||||
|
||||
// 校验并记录警告
|
||||
let warnings = frontmatter.validate(dir_name);
|
||||
for w in &warnings {
|
||||
warn!("[Skills] {}", w);
|
||||
}
|
||||
|
||||
let name = frontmatter.name.unwrap_or_else(|| dir_name.to_string());
|
||||
let description = frontmatter
|
||||
.description
|
||||
.unwrap_or_else(|| "(无描述)".to_string());
|
||||
let context = frontmatter.context.filter(|c| !c.is_empty());
|
||||
let allowed_tools = frontmatter.allowed_tools.unwrap_or_default();
|
||||
let when_to_use = frontmatter.when_to_use.filter(|w| !w.is_empty());
|
||||
let disable_model_invocation = frontmatter.disable_model_invocation.unwrap_or(false);
|
||||
let user_invocable = frontmatter.user_invocable.unwrap_or(true);
|
||||
let paths = frontmatter.paths.unwrap_or_default();
|
||||
let skill_dir = skill_md_path
|
||||
.parent()
|
||||
.unwrap_or(Path::new("."))
|
||||
.to_path_buf();
|
||||
|
||||
info!(
|
||||
"[Skills] 已加载 skill: {} (context={:?}, allowed_tools={:?}, paths={:?})",
|
||||
name, context, allowed_tools, paths
|
||||
);
|
||||
|
||||
Ok(Skill {
|
||||
meta: SkillMeta {
|
||||
name,
|
||||
description,
|
||||
context,
|
||||
allowed_tools,
|
||||
when_to_use,
|
||||
disable_model_invocation,
|
||||
user_invocable,
|
||||
paths,
|
||||
},
|
||||
body,
|
||||
skill_dir,
|
||||
})
|
||||
}
|
||||
|
||||
// ── Backward-Compatible Public API ─────────────────────────────────────────
|
||||
//
|
||||
// 这些函数保留用于外部调用(如 health check、CLI 工具),
|
||||
// 核心路径(Runtime、LoadSkillTool)应使用 SkillRegistry。
|
||||
|
||||
/// 从 skills 目录加载单个 skill(不经过缓存,直接读文件)。
|
||||
pub fn load_skill_direct(skills_dir: &Path, skill_name: &str) -> Option<Skill> {
|
||||
let file_path = skills_dir.join(skill_name).join("SKILL.md");
|
||||
match load_skill_from_path(&file_path, skill_name) {
|
||||
Ok(skill) => Some(skill),
|
||||
Err(e) => {
|
||||
warn!("[Skills] 直接加载 skill '{}' 失败: {}", skill_name, e);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Variable Substitution ──────────────────────────────────────────────────
|
||||
|
||||
/// 在 skill 正文中执行变量替换。
|
||||
///
|
||||
/// 支持的变量:
|
||||
/// - `${SKILL_DIR}` → skill 所在目录的绝对路径
|
||||
/// - `${SESSION_ID}` → 当前会话 ID
|
||||
pub fn substitute_variables(body: &str, skill_dir: &Path, session_id: Option<&str>) -> String {
|
||||
let mut result = body.to_string();
|
||||
|
||||
// ${SKILL_DIR} — skill 所在目录
|
||||
if let Some(dir_str) = skill_dir.to_str() {
|
||||
result = result.replace("${SKILL_DIR}", dir_str);
|
||||
}
|
||||
|
||||
// ${SESSION_ID}
|
||||
if let Some(sid) = session_id {
|
||||
result = result.replace("${SESSION_ID}", sid);
|
||||
} else {
|
||||
result = result.replace("${SESSION_ID}", "");
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
// ── Tests ──────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── Frontmatter 解析 ──
|
||||
|
||||
#[test]
|
||||
fn test_parse_basic_frontmatter() {
|
||||
let raw = "---\nname: test-skill\ndescription: A test skill\n---\n\n# Body\nSome content";
|
||||
let (fm, body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.name.unwrap(), "test-skill");
|
||||
assert_eq!(fm.description.unwrap(), "A test skill");
|
||||
assert!(body.contains("# Body"));
|
||||
assert!(body.contains("Some content"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_with_list() {
|
||||
let raw = "---\nname: test\ndescription: Test\nallowed-tools:\n- bash\n- read\n---\nBody";
|
||||
let (fm, body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.name.unwrap(), "test");
|
||||
let tools = fm.allowed_tools.unwrap();
|
||||
assert_eq!(tools, vec!["bash", "read"]);
|
||||
assert_eq!(body, "Body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_no_fm() {
|
||||
let raw = "# Just a header\nSome content";
|
||||
let (fm, body) = parse_skill_file(raw).unwrap();
|
||||
assert!(fm.name.is_none());
|
||||
assert_eq!(body, raw);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_context_fork() {
|
||||
let raw = "---\nname: heavy\ndescription: Heavy skill\ncontext: fork\n---\nBody";
|
||||
let (fm, _body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.context.unwrap(), "fork");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_boolean_fields() {
|
||||
let raw = "---\nname: test\ndescription: Test\ndisable-model-invocation: true\nuser-invocable: false\n---\nBody";
|
||||
let (fm, _body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.disable_model_invocation, Some(true));
|
||||
assert_eq!(fm.user_invocable, Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_paths() {
|
||||
let raw = "---\nname: test\ndescription: Test\npaths:\n- \"*.rs\"\n- \"*.md\"\n---\nBody";
|
||||
let (fm, _body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.paths.unwrap(), vec!["*.rs", "*.md"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_when_to_use() {
|
||||
let raw = "---\nname: test\ndescription: Test\nwhen_to_use: When user asks about testing\n---\nBody";
|
||||
let (fm, _body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.when_to_use.unwrap(), "When user asks about testing");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_frontmatter_all_fields() {
|
||||
let raw = r#"---
|
||||
name: full-skill
|
||||
description: A comprehensive skill
|
||||
context: fork
|
||||
allowed-tools:
|
||||
- bash
|
||||
- read
|
||||
model: sonnet
|
||||
argument-hint: "<file>"
|
||||
when_to_use: When doing comprehensive tasks
|
||||
disable-model-invocation: false
|
||||
user-invocable: true
|
||||
version: "1.0"
|
||||
paths:
|
||||
- "*.rs"
|
||||
- "*.toml"
|
||||
agent: code-reviewer
|
||||
effort: high
|
||||
---
|
||||
Body content here"#;
|
||||
let (fm, body) = parse_skill_file(raw).unwrap();
|
||||
assert_eq!(fm.name.unwrap(), "full-skill");
|
||||
assert_eq!(fm.description.unwrap(), "A comprehensive skill");
|
||||
assert_eq!(fm.context.unwrap(), "fork");
|
||||
assert_eq!(fm.allowed_tools.unwrap(), vec!["bash", "read"]);
|
||||
assert_eq!(fm.model.unwrap(), "sonnet");
|
||||
assert_eq!(fm.argument_hint.unwrap(), "<file>");
|
||||
assert_eq!(fm.when_to_use.unwrap(), "When doing comprehensive tasks");
|
||||
assert_eq!(fm.disable_model_invocation, Some(false));
|
||||
assert_eq!(fm.user_invocable, Some(true));
|
||||
assert_eq!(fm.version.unwrap(), "1.0");
|
||||
assert_eq!(fm.paths.unwrap(), vec!["*.rs", "*.toml"]);
|
||||
assert_eq!(fm.agent.unwrap(), "code-reviewer");
|
||||
assert_eq!(fm.effort.unwrap(), "high");
|
||||
assert_eq!(body, "Body content here");
|
||||
}
|
||||
|
||||
// ── Validation ──
|
||||
|
||||
#[test]
|
||||
fn test_validate_missing_description() {
|
||||
let fm = SkillFrontmatter {
|
||||
name: Some("test".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let warnings = fm.validate("test");
|
||||
assert!(warnings.iter().any(|w| w.contains("description")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_invalid_context() {
|
||||
let fm = SkillFrontmatter {
|
||||
name: Some("test".into()),
|
||||
description: Some("desc".into()),
|
||||
context: Some("invalid".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let warnings = fm.validate("test");
|
||||
assert!(warnings.iter().any(|w| w.contains("context")));
|
||||
}
|
||||
|
||||
// ── Variable Substitution ──
|
||||
|
||||
#[test]
|
||||
fn test_substitute_skill_dir() {
|
||||
let body = "Base: ${SKILL_DIR}/data";
|
||||
let result = substitute_variables(body, Path::new("/home/user/skills/myskill"), None);
|
||||
assert_eq!(result, "Base: /home/user/skills/myskill/data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_substitute_session_id() {
|
||||
let body = "Session: ${SESSION_ID}";
|
||||
let result = substitute_variables(body, Path::new("."), Some("abc123"));
|
||||
assert_eq!(result, "Session: abc123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_substitute_both() {
|
||||
let body = "Dir: ${SKILL_DIR}\nSession: ${SESSION_ID}";
|
||||
let result = substitute_variables(body, Path::new("/skills/test"), Some("sess-1"));
|
||||
assert!(result.contains("Dir: /skills/test"));
|
||||
assert!(result.contains("Session: sess-1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_substitute_no_session_id() {
|
||||
// 无 session 时占位符被清空
|
||||
let body = "Session: ${SESSION_ID}";
|
||||
let result = substitute_variables(body, Path::new("."), None);
|
||||
assert_eq!(result, "Session: ");
|
||||
}
|
||||
|
||||
// ── SkillRegistry ──
|
||||
|
||||
#[test]
|
||||
fn test_registry_new_empty() {
|
||||
let registry = SkillRegistry::new(PathBuf::from("/nonexistent"));
|
||||
assert!(registry.is_empty());
|
||||
assert!(registry.build_reminder().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_usage_stats() {
|
||||
let mut registry = SkillRegistry::new(PathBuf::from("/nonexistent"));
|
||||
assert_eq!(registry.usage_score("test"), 0.0);
|
||||
|
||||
registry.record_usage("test");
|
||||
assert!(registry.usage_score("test") > 0.0);
|
||||
|
||||
let stats = registry.usage_stats();
|
||||
assert_eq!(stats.get("test").unwrap().invoke_count, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_reminder_empty() {
|
||||
let registry = SkillRegistry::new(PathBuf::from("/nonexistent"));
|
||||
assert!(registry.build_reminder().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_tool_description_empty() {
|
||||
let registry = SkillRegistry::new(PathBuf::from("/nonexistent"));
|
||||
assert!(registry.build_tool_description().contains("没有可用的技能"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,446 @@
|
||||
// src/agent/subagent.rs
|
||||
//
|
||||
// 子代理运行器 — 上下文隔离子代理(参考 Claude Code s04 Subagents)。
|
||||
//
|
||||
// 父代理通过 delegate_research 工具将子任务委托给子代理执行。
|
||||
// 子代理拥有:
|
||||
// - 全新的 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;
|
||||
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>>,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带完整 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,
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义 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,
|
||||
}
|
||||
}
|
||||
|
||||
/// 运行子代理的 ReAct 循环,返回最终文本摘要。
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `system_prompt` - 子代理的系统提示词
|
||||
/// * `research_prompt` - 要执行的研究任务描述
|
||||
/// * `max_steps` - 子代理最大推理步数(默认 5)
|
||||
/// * `hook_registry` - 可选的 HookRegistry(用于触发子代理生命周期事件)
|
||||
pub async fn run(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
research_prompt: &str,
|
||||
max_steps: usize,
|
||||
) -> ToolOutput {
|
||||
let subagent_name = "delegate_research";
|
||||
|
||||
// OnSubagentStart hook
|
||||
if let Some(ref registry) = self.hook_registry {
|
||||
registry
|
||||
.run_on_subagent_start(&SubagentStartContext {
|
||||
parent_session_id: String::new(),
|
||||
subagent_name: subagent_name.to_string(),
|
||||
prompt: research_prompt.to_string(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// 执行实际工作并捕获结果,以便触发 OnSubagentStop hook
|
||||
let result = self
|
||||
.run_inner(system_prompt, research_prompt, max_steps)
|
||||
.await;
|
||||
let (is_error, result_summary) = if result.is_error {
|
||||
(true, result.content.clone())
|
||||
} else {
|
||||
(false, result.content.chars().take(200).collect())
|
||||
};
|
||||
|
||||
if let Some(ref registry) = self.hook_registry {
|
||||
registry
|
||||
.run_on_subagent_stop(&SubagentStopContext {
|
||||
parent_session_id: String::new(),
|
||||
subagent_name: subagent_name.to_string(),
|
||||
result_summary,
|
||||
steps: max_steps,
|
||||
is_error,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// 实际执行逻辑(提取为内部方法以便 hook 包装)
|
||||
async fn run_inner(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
research_prompt: &str,
|
||||
max_steps: usize,
|
||||
) -> 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(
|
||||
&mut messages,
|
||||
llm,
|
||||
self.config.context_char_limit,
|
||||
"subagent",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// LLM 流式调用
|
||||
let mut stream_rx = match llm.chat_stream(&messages, &tool_defs).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_tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;
|
||||
|
||||
while let Some(event) = stream_rx.recv().await {
|
||||
match event {
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
accumulated_tool_calls = Some(tool_calls);
|
||||
}
|
||||
StreamEvent::Done => break,
|
||||
StreamEvent::Error(e) => {
|
||||
warn!("[SubAgent] 流式错误: {}", e);
|
||||
return ToolOutput::error(format!("子代理流式错误: {}", e));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 无工具调用 = 最终回答
|
||||
let tool_calls = match accumulated_tool_calls {
|
||||
Some(ref tc) if !tc.is_empty() => tc.clone(),
|
||||
_ => {
|
||||
// 转发最终文本到父代理
|
||||
if let Some(ref tx) = self.progress_tx {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: format!(
|
||||
"[子代理] {}",
|
||||
accumulated_content.chars().take(200).collect::<String>()
|
||||
),
|
||||
});
|
||||
}
|
||||
let content_len = accumulated_content.len();
|
||||
info!("[SubAgent] 子代理完成,返回 {} 字符摘要", content_len);
|
||||
return ToolOutput::success(
|
||||
accumulated_content,
|
||||
serde_json::json!({
|
||||
"steps": step,
|
||||
"content_length": content_len
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// 构建 assistant 消息
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
if accumulated_content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_content.clone())
|
||||
},
|
||||
None,
|
||||
Some(tool_calls.clone()),
|
||||
);
|
||||
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 {
|
||||
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),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
|
||||
pre_result.final_args
|
||||
} else {
|
||||
args.clone()
|
||||
};
|
||||
|
||||
// Permission check
|
||||
if self.permission_checker.is_denied(tool_name) {
|
||||
warn!("[SubAgent] 权限检查拒绝工具: {}", tool_name);
|
||||
let tool_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!("工具 {} 在子代理上下文中不可用(权限不足)", tool_name),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 执行工具
|
||||
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 {
|
||||
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);
|
||||
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).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 }))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
// src/agent/task_board.rs
|
||||
//
|
||||
// 共享任务看板 — 跨代理任务可见性与依赖图解析。
|
||||
// 参考 learn-claude-code s12 Task System + s17 Autonomous Agents。
|
||||
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::info;
|
||||
|
||||
/// 任务摘要
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TaskSummary {
|
||||
pub session_id: String,
|
||||
pub task_id: String,
|
||||
pub content: String,
|
||||
pub status: String,
|
||||
pub owner: Option<String>,
|
||||
pub can_start: bool,
|
||||
pub blocked_by: Vec<String>,
|
||||
}
|
||||
|
||||
/// 共享任务看板
|
||||
pub struct TaskBoard {
|
||||
db: SqlitePool,
|
||||
}
|
||||
|
||||
impl TaskBoard {
|
||||
pub fn new(db: SqlitePool) -> Self {
|
||||
TaskBoard { db }
|
||||
}
|
||||
|
||||
/// 检查任务是否可以开始(所有 blockedBy 依赖已完成)
|
||||
pub async fn can_start(&self, session_id: &str, task_id: &str) -> anyhow::Result<bool> {
|
||||
let blocked_by: Option<String> = sqlx::query_scalar(
|
||||
"SELECT blocked_by FROM agent_tasks WHERE session_id = ? AND task_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&self.db)
|
||||
.await?
|
||||
.flatten();
|
||||
|
||||
let blocked: Vec<String> =
|
||||
serde_json::from_str(&blocked_by.unwrap_or_default()).unwrap_or_default();
|
||||
if blocked.is_empty() {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
for dep_id in &blocked {
|
||||
let dep_status: Option<String> = sqlx::query_scalar(
|
||||
"SELECT status FROM agent_tasks WHERE session_id = ? AND task_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(dep_id)
|
||||
.fetch_optional(&self.db)
|
||||
.await?
|
||||
.flatten();
|
||||
|
||||
if dep_status.as_deref() != Some("completed") {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// 原子认领任务(乐观并发控制)
|
||||
pub async fn claim_task(
|
||||
&self,
|
||||
session_id: &str,
|
||||
task_id: &str,
|
||||
claimant: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let rows = sqlx::query(
|
||||
"UPDATE agent_tasks SET owner = ?, status = 'in_progress' \
|
||||
WHERE session_id = ? AND task_id = ? \
|
||||
AND (owner IS NULL OR owner = '' OR status = 'pending')",
|
||||
)
|
||||
.bind(claimant)
|
||||
.bind(session_id)
|
||||
.bind(task_id)
|
||||
.execute(&self.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(rows > 0)
|
||||
}
|
||||
|
||||
/// 列出所有可认领的待处理任务(跨 session)
|
||||
pub async fn list_available_tasks(&self, limit: usize) -> anyhow::Result<Vec<TaskSummary>> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
let rows: Vec<(
|
||||
String,
|
||||
String,
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
)> = sqlx::query_as(
|
||||
"SELECT session_id, task_id, content, status, blocked_by, owner \
|
||||
FROM agent_tasks \
|
||||
WHERE status = 'pending' AND (owner IS NULL OR owner = '') \
|
||||
ORDER BY created_at ASC LIMIT ?",
|
||||
)
|
||||
.bind(limit as i32)
|
||||
.fetch_all(&self.db)
|
||||
.await?;
|
||||
|
||||
let mut summaries = Vec::new();
|
||||
for (session_id, task_id, content, status, blocked_by, owner) in rows {
|
||||
let blocked: Vec<String> =
|
||||
serde_json::from_str(&blocked_by.unwrap_or_default()).unwrap_or_default();
|
||||
let can_start = self.can_start(&session_id, &task_id).await.unwrap_or(false);
|
||||
summaries.push(TaskSummary {
|
||||
session_id,
|
||||
task_id,
|
||||
content,
|
||||
status,
|
||||
owner,
|
||||
can_start,
|
||||
blocked_by: blocked,
|
||||
});
|
||||
}
|
||||
|
||||
info!(
|
||||
"[TaskBoard] 查询到 {} 个可认领任务 (limit={})",
|
||||
summaries.len(),
|
||||
limit
|
||||
);
|
||||
Ok(summaries)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_can_start_no_deps() {
|
||||
// can_start 需要数据库连接,此处仅验证结构
|
||||
// 实际集成测试应在 tests/ 目录中
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
// src/agent/team/config.rs
|
||||
//
|
||||
// 团队配置类型与持久化 (.team/{session_id}/config.json)。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 团队成员状态
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum MemberStatus {
|
||||
Spawning,
|
||||
Working,
|
||||
Idle,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// 单个成员配置
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MemberConfig {
|
||||
pub name: String,
|
||||
pub role: String,
|
||||
pub system_prompt: String,
|
||||
pub status: MemberStatus,
|
||||
}
|
||||
|
||||
/// 团队配置
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TeamConfig {
|
||||
pub session_id: String,
|
||||
pub lead_name: String,
|
||||
pub members: Vec<MemberConfig>,
|
||||
}
|
||||
|
||||
impl TeamConfig {
|
||||
/// 创建新的团队配置
|
||||
pub fn new(session_id: &str) -> Self {
|
||||
TeamConfig {
|
||||
session_id: session_id.to_string(),
|
||||
lead_name: "lead".to_string(),
|
||||
members: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 添加成员
|
||||
pub fn add_member(&mut self, name: &str, role: &str, system_prompt: &str) {
|
||||
self.members.push(MemberConfig {
|
||||
name: name.to_string(),
|
||||
role: role.to_string(),
|
||||
system_prompt: system_prompt.to_string(),
|
||||
status: MemberStatus::Spawning,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
// src/agent/team/inbox.rs
|
||||
//
|
||||
// 团队消息邮箱系统。
|
||||
// 使用 .team/{session_id}/inbox/{agent_name}.jsonl 作为 append-only 消息文件。
|
||||
|
||||
use chrono::Utc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// 团队消息类型
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TeamMessageType {
|
||||
Task,
|
||||
Result,
|
||||
Question,
|
||||
Answer,
|
||||
Status,
|
||||
}
|
||||
|
||||
impl TeamMessageType {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
TeamMessageType::Task => "task",
|
||||
TeamMessageType::Result => "result",
|
||||
TeamMessageType::Question => "question",
|
||||
TeamMessageType::Answer => "answer",
|
||||
TeamMessageType::Status => "status",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 团队消息
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TeamMessage {
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
pub content: String,
|
||||
pub msg_type: TeamMessageType,
|
||||
pub timestamp: String,
|
||||
}
|
||||
|
||||
impl TeamMessage {
|
||||
pub fn new(from: &str, to: &str, content: &str, msg_type: TeamMessageType) -> Self {
|
||||
TeamMessage {
|
||||
from: from.to_string(),
|
||||
to: to.to_string(),
|
||||
content: content.to_string(),
|
||||
msg_type,
|
||||
timestamp: Utc::now().to_rfc3339(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取团队目录路径
|
||||
pub fn team_dir(session_id: &str) -> PathBuf {
|
||||
PathBuf::from(".team").join(session_id)
|
||||
}
|
||||
|
||||
/// 获取指定 agent 的收件箱路径
|
||||
pub fn inbox_path(team_dir: &Path, agent_name: &str) -> PathBuf {
|
||||
team_dir.join("inbox").join(format!("{}.jsonl", agent_name))
|
||||
}
|
||||
|
||||
/// 向收件箱追加一条消息
|
||||
pub fn append_message(team_dir: &Path, agent_name: &str, msg: &TeamMessage) -> std::io::Result<()> {
|
||||
let inbox = inbox_path(team_dir, agent_name);
|
||||
if let Some(parent) = inbox.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let line = serde_json::to_string(msg).unwrap_or_default();
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&inbox)?;
|
||||
file.write_all(line.as_bytes())?;
|
||||
file.write_all(b"\n")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 读取并清空收件箱
|
||||
pub fn drain_inbox(team_dir: &Path, agent_name: &str) -> Vec<TeamMessage> {
|
||||
let inbox = inbox_path(team_dir, agent_name);
|
||||
if !inbox.exists() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let content = match std::fs::read_to_string(&inbox) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let messages: Vec<TeamMessage> = content
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.filter_map(|l| serde_json::from_str(l).ok())
|
||||
.collect();
|
||||
|
||||
// 清空文件
|
||||
let _ = std::fs::write(&inbox, "");
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
/// 检查收件箱中是否有未读消息
|
||||
pub fn has_pending(team_dir: &Path, agent_name: &str) -> bool {
|
||||
let inbox = inbox_path(team_dir, agent_name);
|
||||
inbox.exists() && inbox.metadata().map(|m| m.len() > 0).unwrap_or(false)
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
// src/agent/team/manager.rs
|
||||
//
|
||||
// 团队管理器:spawn/stop/send/broadcast/list。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::agent::runtime::AgentConfig;
|
||||
use crate::api::AppState;
|
||||
|
||||
use super::config::{MemberStatus, TeamConfig};
|
||||
use super::inbox::{self, TeamMessage, TeamMessageType};
|
||||
use super::teammate;
|
||||
|
||||
/// 队友运行时句柄
|
||||
#[derive(Clone)]
|
||||
pub struct TeamMemberHandle {
|
||||
pub name: String,
|
||||
pub role: String,
|
||||
pub status: Arc<Mutex<MemberStatus>>,
|
||||
cancelled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl TeamMemberHandle {
|
||||
/// 请求停止该队友
|
||||
pub fn request_stop(&self) {
|
||||
self.cancelled.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
/// 团队管理器
|
||||
///
|
||||
/// # 锁顺序约定 (CRITICAL)
|
||||
///
|
||||
/// 本模块中存在两个 tokio::sync::Mutex 的嵌套获取:
|
||||
/// 1. `TeamManager.handles` (外层)
|
||||
/// 2. `TeamMemberHandle.status` (内层)
|
||||
///
|
||||
/// 任何代码如果先获取 `status` 再获取 `handles` 将导致死锁。
|
||||
/// 所有新增代码必须遵守 `handles → status` 的顺序。
|
||||
/// 参考: `list_members()` 的实现作为正确顺序的示例。
|
||||
pub struct TeamManager {
|
||||
pub config: TeamConfig,
|
||||
team_dir: std::path::PathBuf,
|
||||
handles: Arc<Mutex<HashMap<String, TeamMemberHandle>>>,
|
||||
app_state: Arc<AppState>,
|
||||
}
|
||||
|
||||
impl TeamManager {
|
||||
/// 创建新的团队管理器
|
||||
pub fn new(app_state: Arc<AppState>, session_id: &str) -> Self {
|
||||
let team_dir = inbox::team_dir(session_id);
|
||||
TeamManager {
|
||||
config: TeamConfig::new(session_id),
|
||||
team_dir,
|
||||
handles: Arc::new(Mutex::new(HashMap::new())),
|
||||
app_state,
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取团队目录
|
||||
pub fn team_dir(&self) -> &std::path::Path {
|
||||
&self.team_dir
|
||||
}
|
||||
|
||||
/// 生成队友的系统提示词
|
||||
fn build_teammate_system_prompt(role: &str) -> String {
|
||||
format!(
|
||||
"你是一位专业的天体物理学研究助手,在一个研究团队中工作。你的角色是:{}。\n\
|
||||
\n\
|
||||
你可以使用文献搜索、下载、RAG 检索等工具完成任务。\n\
|
||||
你通过团队收件箱接收任务分配,完成后通过消息汇报结果。\n\
|
||||
\n\
|
||||
核心原则:\n\
|
||||
1. 收到任务后立即开始工作,不要等待确认。\n\
|
||||
2. 完成任务后向 lead 发送 Result 类型的消息汇报。\n\
|
||||
3. 只使用与你的角色相关的工具。\n\
|
||||
4. 用中文输出结果,引用具体文献来源。",
|
||||
role
|
||||
)
|
||||
}
|
||||
|
||||
/// 生成队友的工作任务提示
|
||||
fn build_teammate_task_prompt(role: &str) -> String {
|
||||
format!(
|
||||
"你已加入天体物理研究团队,角色:{}。\n\
|
||||
请在收件箱中等待来自 lead 分配的任务。\n\
|
||||
收到任务后使用你的工具完成,然后将结果发送回 lead。",
|
||||
role
|
||||
)
|
||||
}
|
||||
|
||||
/// 启动一个队友
|
||||
pub async fn spawn(&self, name: &str, role: &str) -> TeamMemberHandle {
|
||||
info!("[TeamManager] 启动队友: {} ({})", name, role);
|
||||
|
||||
let cancelled = Arc::new(AtomicBool::new(false));
|
||||
let cancelled_clone = cancelled.clone();
|
||||
let status = Arc::new(Mutex::new(MemberStatus::Spawning));
|
||||
let status_clone = status.clone();
|
||||
let name_clone = name.to_string();
|
||||
|
||||
let app_state = self.app_state.clone();
|
||||
let team_dir = self.team_dir.clone();
|
||||
let role_owned = role.to_string();
|
||||
let system_prompt = Self::build_teammate_system_prompt(&role_owned);
|
||||
let task_prompt = Self::build_teammate_task_prompt(&role_owned);
|
||||
let agent_config = AgentConfig::from_env_optional();
|
||||
|
||||
// 后台启动队友 ReAct 循环
|
||||
tokio::spawn(async move {
|
||||
teammate::run_teammate_loop(
|
||||
app_state,
|
||||
team_dir,
|
||||
name_clone,
|
||||
role_owned,
|
||||
system_prompt,
|
||||
task_prompt,
|
||||
agent_config,
|
||||
status_clone,
|
||||
cancelled_clone,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
let handle = TeamMemberHandle {
|
||||
name: name.to_string(),
|
||||
role: role.to_string(),
|
||||
status,
|
||||
cancelled,
|
||||
};
|
||||
|
||||
let mut handles = self.handles.lock().await;
|
||||
handles.insert(name.to_string(), handle.clone());
|
||||
handle
|
||||
}
|
||||
|
||||
/// 停止一个队友
|
||||
pub async fn stop(&self, name: &str) {
|
||||
info!("[TeamManager] 停止队友: {}", name);
|
||||
let handles = self.handles.lock().await;
|
||||
if let Some(handle) = handles.get(name) {
|
||||
handle.request_stop();
|
||||
}
|
||||
}
|
||||
|
||||
/// 停止所有队友
|
||||
pub async fn stop_all(&self) {
|
||||
info!("[TeamManager] 停止所有队友");
|
||||
let handles = self.handles.lock().await;
|
||||
for handle in handles.values() {
|
||||
handle.request_stop();
|
||||
}
|
||||
}
|
||||
|
||||
/// 发送消息给指定队友
|
||||
pub fn send_message(&self, from: &str, to: &str, content: &str, msg_type: TeamMessageType) {
|
||||
let msg = TeamMessage::new(from, to, content, msg_type);
|
||||
if let Err(e) = inbox::append_message(&self.team_dir, to, &msg) {
|
||||
warn!("[TeamManager] 发送消息失败: {}", e);
|
||||
} else {
|
||||
info!(
|
||||
"[TeamManager] {} -> {}: {}",
|
||||
from,
|
||||
to,
|
||||
&content.chars().take(80).collect::<String>()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 广播消息给所有队友
|
||||
pub fn broadcast(&self, from: &str, content: &str) {
|
||||
info!(
|
||||
"[TeamManager] 广播消息: {}",
|
||||
content.chars().take(80).collect::<String>()
|
||||
);
|
||||
let msg = TeamMessage::new(from, "all", content, TeamMessageType::Status);
|
||||
for member in &self.config.members {
|
||||
if member.name != from {
|
||||
let _ = inbox::append_message(&self.team_dir, &member.name, &msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查并清空指定 agent 的收件箱
|
||||
pub fn check_inbox(&self, agent_name: &str) -> Vec<TeamMessage> {
|
||||
inbox::drain_inbox(&self.team_dir, agent_name)
|
||||
}
|
||||
|
||||
/// 列出所有队友状态
|
||||
///
|
||||
/// 正确锁顺序示例:先获取 `handles`,再获取每个 `status`。
|
||||
pub async fn list_members(&self) -> Vec<(String, String, MemberStatus)> {
|
||||
let handles = self.handles.lock().await;
|
||||
let mut members = Vec::new();
|
||||
for (name, handle) in handles.iter() {
|
||||
let status = handle.status.lock().await.clone();
|
||||
members.push((name.clone(), handle.role.clone(), status));
|
||||
}
|
||||
members
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
// src/agent/team/mod.rs
|
||||
//
|
||||
// 多智能体团队协作模块(参考 Claude Code s09 Agent Teams)。
|
||||
//
|
||||
// 通过文件邮箱 (.team/{session_id}/inbox/*.jsonl) 实现
|
||||
// Lead Agent 与多个 Teammate Agent 之间的消息传递与任务协调。
|
||||
|
||||
pub mod config;
|
||||
pub mod inbox;
|
||||
pub mod manager;
|
||||
pub mod teammate;
|
||||
@@ -0,0 +1,270 @@
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
// src/agent/terminal.rs
|
||||
//
|
||||
// 智能体循环终止信号。
|
||||
// 参考 Claude Code query.ts 的 Terminal 类型设计:
|
||||
// 用结构化枚举替代隐式的 break / Err(...) 退出,
|
||||
// 使调用方可以精确知道循环为何结束。
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
/// Agent 循环终止原因
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "reason", content = "detail")]
|
||||
pub enum TurnTerminal {
|
||||
/// 正常完成 - Agent 给出了最终回答
|
||||
Completed {
|
||||
session_id: String,
|
||||
total_steps: usize,
|
||||
},
|
||||
|
||||
/// 达到最大推理步数
|
||||
MaxStepsReached {
|
||||
session_id: String,
|
||||
steps: usize,
|
||||
max_steps: usize,
|
||||
},
|
||||
|
||||
/// 用户手动中止
|
||||
CancelledByUser { session_id: String, at_step: usize },
|
||||
|
||||
/// 检测到工具死循环
|
||||
DuplicateCallDetected {
|
||||
session_id: String,
|
||||
tool_name: String,
|
||||
at_step: usize,
|
||||
},
|
||||
|
||||
/// 大模型流式调用失败
|
||||
ModelStreamError {
|
||||
session_id: String,
|
||||
message: String,
|
||||
at_step: usize,
|
||||
},
|
||||
|
||||
/// 大模型返回错误(非流式)
|
||||
ModelError { session_id: String, message: String },
|
||||
}
|
||||
|
||||
impl TurnTerminal {
|
||||
/// 是否为正常完成
|
||||
pub fn is_completed(&self) -> bool {
|
||||
matches!(self, TurnTerminal::Completed { .. })
|
||||
}
|
||||
|
||||
/// 获取关联的 session_id
|
||||
pub fn session_id(&self) -> &str {
|
||||
match self {
|
||||
TurnTerminal::Completed { session_id, .. }
|
||||
| TurnTerminal::MaxStepsReached { session_id, .. }
|
||||
| TurnTerminal::CancelledByUser { session_id, .. }
|
||||
| TurnTerminal::DuplicateCallDetected { session_id, .. }
|
||||
| TurnTerminal::ModelStreamError { session_id, .. }
|
||||
| TurnTerminal::ModelError { session_id, .. } => session_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// 人类可读的终止描述
|
||||
pub fn description(&self) -> &str {
|
||||
match self {
|
||||
TurnTerminal::Completed { .. } => "正常完成",
|
||||
TurnTerminal::MaxStepsReached { .. } => "达到最大步数",
|
||||
TurnTerminal::CancelledByUser { .. } => "用户手动中止",
|
||||
TurnTerminal::DuplicateCallDetected { .. } => "检测到死循环",
|
||||
TurnTerminal::ModelStreamError { .. } => "模型流式错误",
|
||||
TurnTerminal::ModelError { .. } => "模型错误",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TurnTerminal {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
TurnTerminal::Completed {
|
||||
session_id,
|
||||
total_steps,
|
||||
} => {
|
||||
write!(f, "[{}] 正常完成 ({} steps)", session_id, total_steps)
|
||||
}
|
||||
TurnTerminal::MaxStepsReached {
|
||||
session_id,
|
||||
steps,
|
||||
max_steps,
|
||||
} => {
|
||||
write!(f, "[{}] 达到最大步数 ({}/{})", session_id, steps, max_steps)
|
||||
}
|
||||
TurnTerminal::CancelledByUser {
|
||||
session_id,
|
||||
at_step,
|
||||
} => {
|
||||
write!(f, "[{}] 用户在第 {} 步手动中止", session_id, at_step)
|
||||
}
|
||||
TurnTerminal::DuplicateCallDetected {
|
||||
session_id,
|
||||
tool_name,
|
||||
at_step,
|
||||
} => {
|
||||
write!(
|
||||
f,
|
||||
"[{}] 检测到 {} 死循环 (step {})",
|
||||
session_id, tool_name, at_step
|
||||
)
|
||||
}
|
||||
TurnTerminal::ModelStreamError {
|
||||
session_id,
|
||||
message,
|
||||
at_step,
|
||||
} => {
|
||||
write!(f, "[{}] 流式错误 step {}: {}", session_id, at_step, message)
|
||||
}
|
||||
TurnTerminal::ModelError {
|
||||
session_id,
|
||||
message,
|
||||
} => {
|
||||
write!(f, "[{}] 模型错误: {}", session_id, message)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,659 +0,0 @@
|
||||
// src/agent/tools.rs
|
||||
//
|
||||
// 科研智能体工具集定义与实现。
|
||||
// 每个工具遵循 AgentTool trait,向大模型声明 JSON Schema 参数定义,
|
||||
// 并在 execute 中调用已有的服务层完成实际业务操作。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, error};
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::ToolDefinition;
|
||||
|
||||
/// 工具执行上下文,封装全局共享状态
|
||||
pub struct ToolContext {
|
||||
pub app_state: Arc<AppState>,
|
||||
}
|
||||
|
||||
/// 工具执行结果
|
||||
#[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!({}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 智能体工具 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;
|
||||
}
|
||||
|
||||
/// 工具注册表,管理所有可用工具
|
||||
pub struct ToolRegistry {
|
||||
tools: Vec<Box<dyn AgentTool>>,
|
||||
}
|
||||
impl ToolRegistry {
|
||||
/// 创建默认工具注册表(包含全部科研工具)
|
||||
pub fn new() -> Self {
|
||||
let tools: Vec<Box<dyn AgentTool>> = vec![
|
||||
Box::new(SearchPapersTool),
|
||||
Box::new(GetPaperMetadataTool),
|
||||
Box::new(DownloadPaperTool),
|
||||
Box::new(ParsePaperTool),
|
||||
Box::new(GetPaperContentTool),
|
||||
Box::new(RagSearchTool),
|
||||
Box::new(QueryTargetTool),
|
||||
];
|
||||
ToolRegistry { tools }
|
||||
}
|
||||
|
||||
/// 根据名称查找工具
|
||||
pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
|
||||
self.tools.iter().find(|t| t.name() == name).map(|t| t.as_ref())
|
||||
}
|
||||
|
||||
/// 生成所有工具的 ToolDefinition 列表(用于发送给 LLM)
|
||||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||||
self.tools.iter().map(|t| {
|
||||
ToolDefinition::new(t.name(), t.description(), t.parameters())
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// 截断文本到指定最大字符数
|
||||
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())
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 1. SearchPapersTool ──────────────────────────
|
||||
|
||||
/// 文献搜索工具:调用 ADS/arXiv 进行跨库检索
|
||||
pub struct SearchPapersTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SearchPapersTool {
|
||||
fn name(&self) -> &str { "search_papers" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"搜索天文学文献。支持 NASA ADS 和 arXiv 跨平台联合检索,结果自动合并去重,并关联本地馆藏状态与引用关系网。输入关键词或高级检索式,返回匹配的文献列表。\
|
||||
适用于:查找相关文献、了解研究领域现状、获取特定主题的论文。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "搜索关键词或高级检索式。支持语法:\
|
||||
1. 字段限定:au:\"作者\" 或 author:\"作者\"、ti:\"标题\" 或 title:\"标题\"、abs:\"摘要关键字\";\
|
||||
2. 年份限定:year:2020(单年)或 year:2020-2025(年份区间);\
|
||||
3. 逻辑运算:支持 AND、OR、NOT 逻辑组合及括号分组,如 '(ti:subdwarf OR ti:\"white dwarf\") AND year:2020-2025';\
|
||||
4. 短语匹配:用双引号 \"\" 包含精确匹配短语,如 '\"Gaia BH1\"'。\
|
||||
所有的中文标点符号(如“”(),;)在后台均会自动清洗转换。"
|
||||
},
|
||||
"rows": {
|
||||
"type": "integer",
|
||||
"description": "返回结果数量,默认5,最大20",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let query = match args.get("query").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'query'"),
|
||||
};
|
||||
let rows = args.get("rows").and_then(|r| r.as_i64()).unwrap_or(5).min(20) as i32;
|
||||
|
||||
info!("[SearchPapersTool] 执行文献搜索: query='{}', rows={}", query, rows);
|
||||
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::search::search_papers(state, &query, "all", 0, rows, "relevance").await {
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success("未找到匹配的文献。请尝试调整搜索关键词。", json!({ "count": 0 }));
|
||||
}
|
||||
|
||||
// 格式化结果(保留详细信息给 LLM,但不含摘要且不作截断)
|
||||
let display_results: Vec<serde_json::Value> = results.iter().map(|p| {
|
||||
let first_author = p.authors.first().cloned().unwrap_or_else(|| "未知".to_string());
|
||||
json!({
|
||||
"bibcode": p.bibcode,
|
||||
"title": p.title,
|
||||
"first_author": first_author,
|
||||
"year": p.year,
|
||||
"citation_count": p.citation_count,
|
||||
})
|
||||
}).collect();
|
||||
|
||||
let content = display_results.iter().enumerate().map(|(i, r)| {
|
||||
format!(
|
||||
"{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次",
|
||||
i + 1,
|
||||
r["bibcode"].as_str().unwrap_or(""),
|
||||
r["title"].as_str().unwrap_or(""),
|
||||
r["year"].as_str().unwrap_or(""),
|
||||
r["first_author"].as_str().unwrap_or("未知"),
|
||||
r["citation_count"].as_i64().unwrap_or(0)
|
||||
)
|
||||
}).collect::<Vec<_>>().join("\n\n");
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({
|
||||
"count": results.len(),
|
||||
"papers": display_results
|
||||
})
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[SearchPapersTool] 检索失败: {}", e);
|
||||
ToolOutput::error(format!("文献检索失败: {}", e))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 1b. GetPaperMetadataTool ──────────────────────────
|
||||
|
||||
/// 获取文献元数据工具:获取指定文献的完整元数据(包含完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等)
|
||||
pub struct GetPaperMetadataTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperMetadataTool {
|
||||
fn name(&self) -> &str { "get_paper_metadata" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"获取指定文献的完整元数据信息(包括完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等)。\
|
||||
适用于:需要查看某篇文献的详细信息、阅读完整摘要以评估文献相关性。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode(如 '2024ApJ...960..123A')、DOI(如 '10.3847/1538-4357/ad0c5a')或 arXiv ID(如 '2401.12345')"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperMetadataTool] 获取文献元数据: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode).await {
|
||||
Ok(paper) => {
|
||||
let content = format!(
|
||||
"文献元数据 [{}]:\n\
|
||||
标题: {}\n\
|
||||
作者: {}\n\
|
||||
年份: {}\n\
|
||||
期刊: {}\n\
|
||||
关键字: {}\n\
|
||||
DOI: {}\n\
|
||||
arXiv ID: {}\n\
|
||||
引用数: {} 次\n\
|
||||
参考文献数: {} 次\n\
|
||||
文献类型: {}\n\
|
||||
已下载: {}\n\
|
||||
已解析为 Markdown: {}\n\
|
||||
摘要:\n{}",
|
||||
paper.bibcode,
|
||||
paper.title,
|
||||
paper.authors.join(", "),
|
||||
paper.year,
|
||||
paper.pub_journal,
|
||||
paper.keywords.join(", "),
|
||||
paper.doi,
|
||||
paper.arxiv_id,
|
||||
paper.citation_count,
|
||||
paper.reference_count,
|
||||
paper.doctype,
|
||||
paper.is_downloaded,
|
||||
paper.has_markdown,
|
||||
paper.abstract_text
|
||||
);
|
||||
ToolOutput::success(content, json!(paper))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("获取文献 {} 元数据失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2. GetPaperContentTool ──────────────────────────
|
||||
|
||||
/// 获取文献内容工具:仅从本地读取并获取已解析的文献 Markdown 全文内容
|
||||
pub struct GetPaperContentTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperContentTool {
|
||||
fn name(&self) -> &str { "get_paper_content" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"读取并在本地库中获取已解析的文献 Markdown 完整文本内容。\
|
||||
注意:本工具仅能读取已在数据库注册且已解析的文献,不会自动触发下载或解析。若文献未下载或未解析,本工具会返回详细指引,提示先依次调用 download_paper 和 parse_paper。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode(如 '2024ApJ...960..123A')、DOI(如 '10.3847/1538-4357/ad0c5a')或 arXiv ID(如 '2401.12345')"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperContentTool] 获取文献内容: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
let paths = crate::api::helpers::check_paper_paths_in_db(&state.db, &state.config.library_dir, &bibcode).await;
|
||||
let md_opt = match paths {
|
||||
Ok(Some((_, _, md_opt, _))) => md_opt,
|
||||
Ok(None) => return ToolOutput::error(format!("获取文献内容失败:该文献未在本地数据库中注册,请先使用 search_papers 搜索该文献。")),
|
||||
Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
|
||||
};
|
||||
|
||||
let md_rel = match md_opt {
|
||||
Some(rel) => rel,
|
||||
None => return ToolOutput::error(format!("获取文献内容失败:该文献尚未完成结构化解析。如果未下载,请先调用 download_paper;如果已下载,请先调用 parse_paper 进行解析。")),
|
||||
};
|
||||
|
||||
let md_abs = state.config.library_dir.join(&md_rel);
|
||||
if !md_abs.exists() {
|
||||
return ToolOutput::error(format!("获取文献内容失败:文献本地 Markdown 文件已丢失,请重新调用 parse_paper 进行解析。"));
|
||||
}
|
||||
|
||||
match std::fs::read_to_string(&md_abs) {
|
||||
Ok(content) => ToolOutput::success(
|
||||
content.clone(),
|
||||
json!({ "bibcode": bibcode, "chars": content.len() })
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("获取文献内容失败,读取本地文件错误: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2a. DownloadPaperTool ──────────────────────────
|
||||
|
||||
/// 下载文献全文资源工具:仅下载文献全文资源(PDF/HTML)至本地图书馆
|
||||
pub struct DownloadPaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for DownloadPaperTool {
|
||||
fn name(&self) -> &str { "download_paper" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"下载指定文献的全文资源(PDF 或 HTML)至本地图书馆,为后续的结构化解析做好准备。\
|
||||
适用于:需要阅读或分析新搜寻到的、尚未下载的文献。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode(如 '2024ApJ...960..123A')、DOI(如 '10.3847/1538-4357/ad0c5a')或 arXiv ID(如 '2401.12345')"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新下载(即使本地已下载该文献)"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[DownloadPaperTool] 下载文献全文资源: {}, 强制重下: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match state.downloader.download_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&bibcode,
|
||||
force,
|
||||
)
|
||||
.await {
|
||||
Ok(paper) => ToolOutput::success(
|
||||
format!("文献 {} 全文资源下载成功。格式 - PDF: {}, HTML: {}", bibcode, paper.has_pdf, paper.has_html),
|
||||
json!({ "bibcode": bibcode, "has_pdf": paper.has_pdf, "has_html": paper.has_html })
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("文献 {} 下载失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2b. ParsePaperTool ──────────────────────────
|
||||
|
||||
/// 结构化解析文献内容工具:仅对已下载的物理资源进行结构化解析生成 Markdown
|
||||
pub struct ParsePaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for ParsePaperTool {
|
||||
fn name(&self) -> &str { "parse_paper" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将指定文献本地已下载的 HTML 或 PDF 资源解析为结构化的 Markdown 文本,并保存至本地 Markdown 文件夹。\
|
||||
注意:调用此工具前必须确保文献已被成功下载(已执行 download_paper)。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode(如 '2024ApJ...960..123A')、DOI(如 '10.3847/1538-4357/ad0c5a')或 arXiv ID(如 '2401.12345')"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新解析(即使本地已解析过该文献)"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[ParsePaperTool] 结构化解析文献内容: {}, 强制重析: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::parser::parse_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&state.qiniu,
|
||||
&state.config,
|
||||
&bibcode,
|
||||
force,
|
||||
)
|
||||
.await {
|
||||
Ok(markdown) => ToolOutput::success(
|
||||
format!("文献 {} 结构化解析成功。解析后 Markdown 字符总数: {}", bibcode, markdown.len()),
|
||||
json!({ "bibcode": bibcode, "chars": markdown.len() })
|
||||
),
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("请先下载") {
|
||||
ToolOutput::error(format!("文献 {} 解析失败:未检测到已下载的本地资源文件,请先调用 download_paper 工具进行下载。", bibcode))
|
||||
} else {
|
||||
ToolOutput::error(format!("文献 {} 解析失败: {}", bibcode, msg))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// ────────────────────────── 4. RagSearchTool ──────────────────────────
|
||||
|
||||
/// RAG 向量检索工具:基于语义相似度检索文献切片
|
||||
pub struct RagSearchTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for RagSearchTool {
|
||||
fn name(&self) -> &str { "rag_search" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在已向量化的文献库中进行语义检索。输入自然语言问题,返回最相关的文献片段。\
|
||||
适用于:跨多篇文献查找特定信息、回答需要综合多个来源的问题。要求文献已完成向量化(embed)。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "用于语义检索的自然语言问题"
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"description": "返回最相关的片段数量,默认5",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["question"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let question = match args.get("question").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'question'"),
|
||||
};
|
||||
let top_k = args.get("top_k").and_then(|k| k.as_u64()).unwrap_or(5) as usize;
|
||||
|
||||
info!("[RagSearchTool] 执行语义检索: question='{}', top_k={}", question, top_k);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::rag::retrieve(&state.db, &state.embedding, &question, top_k).await {
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success(
|
||||
"未找到相关的文献片段。文献库中可能尚无向量化数据,请先对目标文献执行向量化操作。",
|
||||
json!({ "count": 0 })
|
||||
);
|
||||
}
|
||||
|
||||
let content = results.iter().enumerate().map(|(i, r)| {
|
||||
format!(
|
||||
"[片段 {} | 来源: {} | 段落: {} | 相似度距离: {:.4}]\n{}",
|
||||
i + 1, r.bibcode, r.paragraph_index, r.distance, r.content
|
||||
)
|
||||
}).collect::<Vec<_>>().join("\n\n---\n\n");
|
||||
|
||||
let sources: Vec<serde_json::Value> = results.iter().map(|r| {
|
||||
json!({
|
||||
"bibcode": r.bibcode,
|
||||
"paragraph_index": r.paragraph_index,
|
||||
"distance": r.distance,
|
||||
"preview": truncate_content(&r.content, 100)
|
||||
})
|
||||
}).collect();
|
||||
|
||||
ToolOutput::success(
|
||||
truncate_content(&content, 4000),
|
||||
json!({ "count": results.len(), "sources": sources })
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("RAG 语义检索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 5. QueryTargetTool ──────────────────────────
|
||||
|
||||
/// 天体信息查询工具:通过 CDS Sesame 查询天体物理属性
|
||||
pub struct QueryTargetTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for QueryTargetTool {
|
||||
fn name(&self) -> &str { "query_target" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"查询天体的基本物理属性信息。输入天体名称,返回坐标 (RA/Dec)、视星等、光谱型、视差等属性。\
|
||||
数据来源为 CDS SIMBAD/Sesame 名称解析服务。适用于:获取天体基本参数、验证天体身份。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_name": {
|
||||
"type": "string",
|
||||
"description": "天体名称,如 'NGC 6752', 'GD 358', 'HD 209458', 'M 31' 等"
|
||||
}
|
||||
},
|
||||
"required": ["object_name"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let object_name = match args.get("object_name").and_then(|n| n.as_str()) {
|
||||
Some(n) => n.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'object_name'"),
|
||||
};
|
||||
|
||||
info!("[QueryTargetTool] 查询天体信息: {}", object_name);
|
||||
let state = &ctx.app_state;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
match crate::services::target::query_target_cached(&state.db, &object_name, None, &client).await {
|
||||
Ok(info) => {
|
||||
let content = format!(
|
||||
"天体: {}\nRA: {}\nDec: {}\n视差: {}\n光谱型: {}\nV星等: {}\n别名: {}",
|
||||
info.target_name,
|
||||
info.ra.as_deref().unwrap_or("未知"),
|
||||
info.dec.as_deref().unwrap_or("未知"),
|
||||
info.parallax.map(|p| format!("{:.4} mas", p)).unwrap_or_else(|| "未知".to_string()),
|
||||
info.spectral_type.as_deref().unwrap_or("未知"),
|
||||
info.v_magnitude.map(|v| format!("{:.2}", v)).unwrap_or_else(|| "未知".to_string()),
|
||||
if info.aliases.is_empty() { "无".to_string() } else { info.aliases.join(", ") }
|
||||
);
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({
|
||||
"target_name": info.target_name,
|
||||
"ra": info.ra,
|
||||
"dec": info.dec,
|
||||
"parallax": info.parallax,
|
||||
"spectral_type": info.spectral_type,
|
||||
"v_magnitude": info.v_magnitude,
|
||||
"aliases": info.aliases
|
||||
})
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("天体 '{}' 查询失败: {}", object_name, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[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();
|
||||
let defs = registry.definitions();
|
||||
assert_eq!(defs.len(), 7);
|
||||
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"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_registry_get() {
|
||||
let registry = ToolRegistry::new();
|
||||
assert!(registry.get("search_papers").is_some());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
// src/agent/tools/ask_user.rs
|
||||
//
|
||||
// 用户交互工具 — Agent 向用户提问并等待回复。
|
||||
//
|
||||
// 当任务需求不明确时(如缺少参数、需要选择方案),Agent 调用此工具
|
||||
// 阻止 ReAct 循环,向用户展示问题,等待用户回复后继续执行。
|
||||
//
|
||||
// 实现方式:
|
||||
// - 使用 oneshot 通道向 SSE 层发送问题
|
||||
// - 阻塞等待用户通过 API 端点提交答案
|
||||
// - 超时后返回错误(默认 5 分钟)
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::info;
|
||||
|
||||
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
/// 提交给用户的问题
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserQuestion {
|
||||
/// 问题 ID(用于前端关联回答)
|
||||
pub question_id: String,
|
||||
/// 完整问题文本
|
||||
pub question: String,
|
||||
/// 短标签(显示为 chip/tag)
|
||||
pub header: String,
|
||||
/// 预定义选项列表
|
||||
pub options: Vec<UserOption>,
|
||||
/// 是否允许多选
|
||||
pub multi_select: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserOption {
|
||||
pub label: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// 用户的回答
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserAnswer {
|
||||
pub question_id: String,
|
||||
pub answers: Vec<String>,
|
||||
pub free_text: Option<String>,
|
||||
}
|
||||
|
||||
/// 向用户提问的工具。
|
||||
///
|
||||
/// 使用 oneshot 通道机制:创建问题 → 通过 AppState 发送 → 阻塞等待 → 返回答案。
|
||||
pub struct AskUserTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for AskUserTool {
|
||||
fn name(&self) -> &str {
|
||||
"ask_user"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"当任务需求不明确、缺少关键参数、或需要在多个方案之间选择时,向用户提问。\
|
||||
支持预定义选项(单选/多选)。会暂停当前任务等待用户回复。\
|
||||
适用场景:确认文献搜索范围、选择分析方案、确定输出格式。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "要询问用户的问题,应清晰具体。如:'我应该搜索哪个天区的数据?'"
|
||||
},
|
||||
"header": {
|
||||
"type": "string",
|
||||
"description": "问题的简短标签,最多 12 字。如:'搜索范围'、'分析方案'"
|
||||
},
|
||||
"options": {
|
||||
"type": "array",
|
||||
"description": "预定义选项列表(可选,不提供则允许自由回答)",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {
|
||||
"type": "string",
|
||||
"description": "选项标签,简洁明确。如:'Gaia DR3'"
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": "选项说明。如:'盖亚卫星第三期数据发布,包含 18 亿颗恒星'"
|
||||
}
|
||||
},
|
||||
"required": ["label", "description"]
|
||||
}
|
||||
},
|
||||
"multi_select": {
|
||||
"type": "boolean",
|
||||
"description": "是否允许多选,默认 false",
|
||||
"default": false
|
||||
}
|
||||
},
|
||||
"required": ["question", "header"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 必须阻塞等待用户回复
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
// 子代理(silent 模式)不能向用户提问 — 静默模式下无法交互
|
||||
if ctx.silent {
|
||||
return ToolOutput::error(
|
||||
"ask_user 在子代理/后台上下文中不可用。请基于已有信息继续,或使用其他工具获取所需数据。"
|
||||
);
|
||||
}
|
||||
|
||||
let question_text = match args.get("question").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'question'"),
|
||||
};
|
||||
let header = match args.get("header").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'header'"),
|
||||
};
|
||||
let multi_select = args
|
||||
.get("multi_select")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
let options: Vec<UserOption> = args
|
||||
.get("options")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|opt| {
|
||||
Some(UserOption {
|
||||
label: opt.get("label")?.as_str()?.to_string(),
|
||||
description: opt.get("description")?.as_str()?.to_string(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let question_id = uuid::Uuid::new_v4().to_string();
|
||||
let short_id = question_id[..8].to_string();
|
||||
|
||||
let question = UserQuestion {
|
||||
question_id: short_id.clone(),
|
||||
question: question_text.clone(),
|
||||
header,
|
||||
options,
|
||||
multi_select,
|
||||
};
|
||||
|
||||
info!(
|
||||
"[AskUser] 向用户提问: id={}, header={}, options={}",
|
||||
short_id,
|
||||
&question.header,
|
||||
question.options.len()
|
||||
);
|
||||
|
||||
// 创建 oneshot 通道
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let question_json = serde_json::to_string(&question).unwrap_or_default();
|
||||
|
||||
// 将通道发送端存储到 AppState 的待处理问题列表
|
||||
{
|
||||
let mut pending = match ctx.app_state.pending_questions.lock() {
|
||||
Ok(p) => p,
|
||||
Err(_) => {
|
||||
return ToolOutput::error("待处理问题队列不可用(内部锁异常),请稍后重试");
|
||||
}
|
||||
};
|
||||
pending.insert(
|
||||
short_id.clone(),
|
||||
crate::api::PendingQuestion {
|
||||
question_json: question_json.clone(),
|
||||
answer_tx: tx,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// 通过 SSE 通道发送问题事件(如果存在)
|
||||
if let Some(sse_tx) = &ctx.app_state.sse_broadcast {
|
||||
let _ = sse_tx.send(crate::api::AppEvent::UserQuestion {
|
||||
data: question_json,
|
||||
});
|
||||
}
|
||||
|
||||
// 等待用户回答(5 分钟超时)
|
||||
let timeout = tokio::time::Duration::from_secs(300);
|
||||
match tokio::time::timeout(timeout, rx).await {
|
||||
Ok(Ok(answer)) => {
|
||||
info!(
|
||||
"[AskUser] 收到用户回答: id={}, answers={:?}",
|
||||
short_id, answer.answers
|
||||
);
|
||||
|
||||
// 清理
|
||||
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
|
||||
pending.remove(&short_id);
|
||||
}
|
||||
|
||||
let free_text = answer.free_text.unwrap_or_default();
|
||||
let response = if !answer.answers.is_empty() {
|
||||
let note = if free_text.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("\n补充说明: {}", free_text)
|
||||
};
|
||||
format!("用户回答: {}{}", answer.answers.join(", "), note)
|
||||
} else {
|
||||
free_text.clone()
|
||||
};
|
||||
|
||||
ToolOutput::success(
|
||||
response,
|
||||
json!({
|
||||
"question_id": short_id,
|
||||
"answers": answer.answers,
|
||||
"free_text": free_text
|
||||
}),
|
||||
)
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
// 通道关闭(发送端被 drop)
|
||||
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
|
||||
pending.remove(&short_id);
|
||||
}
|
||||
ToolOutput::error("用户取消了回答")
|
||||
}
|
||||
Err(_) => {
|
||||
// 超时
|
||||
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
|
||||
pending.remove(&short_id);
|
||||
}
|
||||
ToolOutput::error(format!(
|
||||
"等待用户回答超时 (5 分钟)。问题: {}",
|
||||
question_text
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
// src/agent/tools/astro/mod.rs
|
||||
//
|
||||
// 天文科研工具集 — 文献搜索/下载/解析、RAG 检索、天体目标查询、研究笔记。
|
||||
|
||||
pub mod note;
|
||||
pub mod paper;
|
||||
pub mod rag;
|
||||
pub mod search;
|
||||
pub mod target;
|
||||
|
||||
pub use note::SaveNoteTool;
|
||||
pub use paper::{DownloadPaperTool, GetPaperContentTool, ParsePaperTool};
|
||||
pub use rag::RagSearchTool;
|
||||
pub use search::{GetPaperMetadataTool, SearchPapersTool};
|
||||
pub use target::QueryTargetTool;
|
||||
@@ -0,0 +1,99 @@
|
||||
// src/agent/tools/note.rs — 研究笔记保存工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
/// 研究笔记保存工具:将 Agent 的研究发现保存为 Markdown 笔记
|
||||
pub struct SaveNoteTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SaveNoteTool {
|
||||
fn name(&self) -> &str {
|
||||
"save_note"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将研究中间结果或最终结论保存为 Markdown 格式的笔记文件。适用于:保存文献综述、记录研究发现、导出分析结果。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"title": {
|
||||
"type": "string",
|
||||
"description": "笔记标题(将作为文件名的一部分)"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "笔记正文内容,支持 Markdown 格式(包括 LaTeX 数学公式)"
|
||||
}
|
||||
},
|
||||
"required": ["title", "content"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 写文件有副作用,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let title = match args.get("title").and_then(|t| t.as_str()) {
|
||||
Some(t) => t.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'title'"),
|
||||
};
|
||||
let content = match args.get("content").and_then(|c| c.as_str()) {
|
||||
Some(c) => c.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'content'"),
|
||||
};
|
||||
|
||||
info!("[SaveNote] 保存笔记: title='{}'", title);
|
||||
|
||||
let safe_filename: String = title
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect::<String>()
|
||||
.trim()
|
||||
.replace(' ', "_");
|
||||
let filename = format!("{}.md", safe_filename);
|
||||
|
||||
let notes_dir = ctx.app_state.config.library_dir.join("notes");
|
||||
if let Err(e) = std::fs::create_dir_all(¬es_dir) {
|
||||
return ToolOutput::error(format!("无法创建笔记目录: {}", e));
|
||||
}
|
||||
|
||||
let filepath = notes_dir.join(&filename);
|
||||
let now = chrono::Local::now();
|
||||
let full_content = format!(
|
||||
"---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}",
|
||||
title,
|
||||
now.format("%Y-%m-%d %H:%M:%S"),
|
||||
content
|
||||
);
|
||||
|
||||
match std::fs::write(&filepath, &full_content) {
|
||||
Ok(_) => {
|
||||
info!("[SaveNote] 笔记已保存: {}", filepath.display());
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"笔记已保存到 {} ({} 字符)",
|
||||
filepath.display(),
|
||||
content.len()
|
||||
),
|
||||
json!({ "filename": filename, "path": filepath.to_string_lossy().to_string(), "size": content.len() }),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("保存笔记失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
// src/agent/tools/paper.rs — 文献下载、解析、内容读取工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
// ── GetPaperContentTool ──
|
||||
|
||||
/// 获取文献内容工具:仅从本地读取已解析的 Markdown 全文
|
||||
pub struct GetPaperContentTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperContentTool {
|
||||
fn name(&self) -> &str {
|
||||
"get_paper_content"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"读取并在本地库中获取已解析的文献 Markdown 完整文本内容。\
|
||||
注意:本工具仅能读取已在数据库注册且已解析的文献,不会自动触发下载或解析。\
|
||||
若文献未下载或未解析,本工具会返回详细指引。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读取操作,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperContent] 获取文献内容: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
let paths = crate::api::helpers::check_paper_paths_in_db(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&bibcode,
|
||||
)
|
||||
.await;
|
||||
let md_opt = match paths {
|
||||
Ok(Some((_, _, md_opt, _))) => md_opt,
|
||||
Ok(None) => return ToolOutput::error(
|
||||
"获取文献内容失败:该文献未在本地数据库中注册,请先使用 search_papers 搜索该文献。",
|
||||
),
|
||||
Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
|
||||
};
|
||||
|
||||
let md_rel = match md_opt {
|
||||
Some(rel) => rel,
|
||||
None => {
|
||||
return ToolOutput::error(
|
||||
"获取文献内容失败:该文献尚未完成结构化解析。如果未下载,请先调用 download_paper;如果已下载,请先调用 parse_paper 进行解析。",
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let md_abs = state.config.library_dir.join(&md_rel);
|
||||
if !md_abs.exists() {
|
||||
return ToolOutput::error(
|
||||
"获取文献内容失败:文献本地 Markdown 文件已丢失,请重新调用 parse_paper 进行解析。",
|
||||
);
|
||||
}
|
||||
|
||||
match std::fs::read_to_string(&md_abs) {
|
||||
Ok(content) => ToolOutput::success(
|
||||
content.clone(),
|
||||
json!({ "bibcode": bibcode, "chars": content.len() }),
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("获取文献内容失败,读取本地文件错误: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── DownloadPaperTool ──
|
||||
|
||||
/// 下载文献全文资源工具
|
||||
pub struct DownloadPaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for DownloadPaperTool {
|
||||
fn name(&self) -> &str {
|
||||
"download_paper"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"下载指定文献的全文资源(PDF 或 HTML)至本地图书馆,为后续的结构化解析做好准备。\
|
||||
适用于:需要阅读或分析新搜寻到的、尚未下载的文献。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新下载(即使本地已下载该文献)"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 下载有副作用,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
/// 下载失败时应中止兄弟并行执行(避免继续处理同一文献)
|
||||
fn causes_sibling_abort(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[DownloadPaper] 下载文献: {}, 强制重下: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match state
|
||||
.downloader
|
||||
.download_paper_service(&state.db, &state.config.library_dir, &bibcode, force)
|
||||
.await
|
||||
{
|
||||
Ok(paper) => ToolOutput::success(
|
||||
format!(
|
||||
"文献 {} 全文资源下载成功。格式 - PDF: {}, HTML: {}",
|
||||
bibcode, paper.has_pdf, paper.has_html
|
||||
),
|
||||
json!({ "bibcode": bibcode, "has_pdf": paper.has_pdf, "has_html": paper.has_html }),
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("文献 {} 下载失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── ParsePaperTool ──
|
||||
|
||||
/// 结构化解析文献内容工具
|
||||
pub struct ParsePaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for ParsePaperTool {
|
||||
fn name(&self) -> &str {
|
||||
"parse_paper"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将指定文献本地已下载的 HTML 或 PDF 资源解析为结构化的 Markdown 文本,并保存至本地 Markdown 文件夹。\
|
||||
注意:调用此工具前必须确保文献已被成功下载(已执行 download_paper)。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新解析(即使本地已解析过该文献)"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 解析有副作用(写文件),中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
/// 解析失败时应中止兄弟并行执行
|
||||
fn causes_sibling_abort(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[ParsePaper] 解析文献: {}, 强制重析: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::parser::parse_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&state.qiniu,
|
||||
&state.config,
|
||||
&bibcode,
|
||||
force,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(markdown) => ToolOutput::success(
|
||||
format!(
|
||||
"文献 {} 结构化解析成功。解析后 Markdown 字符总数: {}",
|
||||
bibcode,
|
||||
markdown.len()
|
||||
),
|
||||
json!({ "bibcode": bibcode, "chars": markdown.len() }),
|
||||
),
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("请先下载") {
|
||||
ToolOutput::error(format!(
|
||||
"文献 {} 解析失败:未检测到已下载的本地资源文件,请先调用 download_paper 工具进行下载。",
|
||||
bibcode
|
||||
))
|
||||
} else {
|
||||
ToolOutput::error(format!("文献 {} 解析失败: {}", bibcode, msg))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
// src/agent/tools/rag.rs — RAG 语义检索工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{truncate_content, AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// RAG 向量检索工具:基于语义相似度检索文献切片
|
||||
pub struct RagSearchTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for RagSearchTool {
|
||||
fn name(&self) -> &str {
|
||||
"rag_search"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在已向量化的文献库中进行语义检索。输入自然语言问题,返回最相关的文献片段。\
|
||||
适用于:跨多篇文献查找特定信息、回答需要综合多个来源的问题。要求文献已完成向量化(embed)。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "用于语义检索的自然语言问题"
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"description": "返回最相关的片段数量,默认5",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["question"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读取向量数据库,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let question = match args.get("question").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'question'"),
|
||||
};
|
||||
let top_k = args.get("top_k").and_then(|k| k.as_u64()).unwrap_or(5) as usize;
|
||||
|
||||
info!(
|
||||
"[RagSearch] 语义检索: question='{}', top_k={}",
|
||||
question, top_k
|
||||
);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::rag::retrieve(&state.db, &state.embedding, &question, top_k).await {
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success(
|
||||
"未找到相关的文献片段。文献库中可能尚无向量化数据,请先对目标文献执行向量化操作。",
|
||||
json!({ "count": 0 }),
|
||||
);
|
||||
}
|
||||
|
||||
let content = results
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, r)| {
|
||||
format!(
|
||||
"[片段 {} | 来源: {} | 段落: {} | 相似度距离: {:.4}]\n{}",
|
||||
i + 1,
|
||||
r.bibcode,
|
||||
r.paragraph_index,
|
||||
r.distance,
|
||||
r.content
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n---\n\n");
|
||||
|
||||
let sources: Vec<serde_json::Value> = results
|
||||
.iter()
|
||||
.map(|r| {
|
||||
json!({
|
||||
"bibcode": r.bibcode,
|
||||
"paragraph_index": r.paragraph_index,
|
||||
"distance": r.distance,
|
||||
"preview": truncate_content(&r.content, 100)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
ToolOutput::success(
|
||||
truncate_content(&content, 4000),
|
||||
json!({ "count": results.len(), "sources": sources }),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("RAG 语义检索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
// src/agent/tools/search.rs — 文献搜索与元数据工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
// ── SearchPapersTool ──
|
||||
|
||||
/// 文献搜索工具:调用 ADS/arXiv 进行跨库检索
|
||||
pub struct SearchPapersTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SearchPapersTool {
|
||||
fn name(&self) -> &str {
|
||||
"search_papers"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"搜索天文学文献。支持 NASA ADS 和 arXiv 跨平台联合检索,结果自动合并去重,并关联本地馆藏状态与引用关系网。输入关键词或高级检索式,返回匹配的文献列表。\
|
||||
适用于:查找相关文献、了解研究领域现状、获取特定主题的论文。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "搜索关键词或高级检索式。支持语法:\
|
||||
1. 字段限定:au:\"作者\" 或 author:\"作者\"、ti:\"标题\" 或 title:\"标题\"、abs:\"摘要关键字\";\
|
||||
2. 年份限定:year:2020(单年)或 year:2020-2025(年份区间);\
|
||||
3. 逻辑运算:支持 AND、OR、NOT 逻辑组合及括号分组,如 '(ti:subdwarf OR ti:\"white dwarf\") AND year:2020-2025';\
|
||||
4. 短语匹配:用双引号 \"\" 包含精确匹配短语,如 '\"Gaia BH1\"'。\
|
||||
所有的中文标点符号(如\"\"(),;)在后台均会自动清洗转换。"
|
||||
},
|
||||
"rows": {
|
||||
"type": "integer",
|
||||
"description": "返回结果数量,默认5,最大20",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读取外部 API,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let query = match args.get("query").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'query'"),
|
||||
};
|
||||
let rows = args
|
||||
.get("rows")
|
||||
.and_then(|r| r.as_i64())
|
||||
.unwrap_or(5)
|
||||
.min(20) as i32;
|
||||
|
||||
info!(
|
||||
"[SearchPapers] 执行文献搜索: query='{}', rows={}",
|
||||
query, rows
|
||||
);
|
||||
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::search::search_papers(state, &query, "all", 0, rows, "relevance")
|
||||
.await
|
||||
{
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success(
|
||||
"未找到匹配的文献。请尝试调整搜索关键词。",
|
||||
json!({ "count": 0 }),
|
||||
);
|
||||
}
|
||||
|
||||
let display_results: Vec<serde_json::Value> = results
|
||||
.iter()
|
||||
.map(|p| {
|
||||
let first_author = p
|
||||
.authors
|
||||
.first()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "未知".to_string());
|
||||
json!({
|
||||
"bibcode": p.bibcode,
|
||||
"title": p.title,
|
||||
"first_author": first_author,
|
||||
"year": p.year,
|
||||
"citation_count": p.citation_count,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let content = display_results
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, r)| {
|
||||
format!(
|
||||
"{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次",
|
||||
i + 1,
|
||||
r["bibcode"].as_str().unwrap_or(""),
|
||||
r["title"].as_str().unwrap_or(""),
|
||||
r["year"].as_str().unwrap_or(""),
|
||||
r["first_author"].as_str().unwrap_or("未知"),
|
||||
r["citation_count"].as_i64().unwrap_or(0)
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({ "count": results.len(), "papers": display_results }),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[SearchPapers] 检索失败: {}", e);
|
||||
ToolOutput::error(format!("文献检索失败: {}", e))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── GetPaperMetadataTool ──
|
||||
|
||||
/// 获取文献元数据工具:获取指定文献的完整元数据
|
||||
pub struct GetPaperMetadataTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperMetadataTool {
|
||||
fn name(&self) -> &str {
|
||||
"get_paper_metadata"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"获取指定文献的完整元数据信息(包括完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等)。\
|
||||
适用于:需要查看某篇文献的详细信息、阅读完整摘要以评估文献相关性。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符,支持 ADS Bibcode(如 '2024ApJ...960..123A')、DOI 或 arXiv ID"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读取数据库,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperMetadata] 获取文献元数据: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
|
||||
.await
|
||||
{
|
||||
Ok(paper) => {
|
||||
let content = format!(
|
||||
"文献元数据 [{}]:\n 标题: {}\n 作者: {}\n 年份: {}\n 期刊: {}\n 关键字: {}\n DOI: {}\n arXiv ID: {}\n 引用数: {} 次\n 参考文献数: {} 次\n 文献类型: {}\n 已下载: {}\n 已解析为 Markdown: {}\n 摘要:\n{}",
|
||||
paper.bibcode, paper.title, paper.authors.join(", "), paper.year,
|
||||
paper.pub_journal, paper.keywords.join(", "), paper.doi, paper.arxiv_id,
|
||||
paper.citation_count, paper.reference_count, paper.doctype,
|
||||
paper.is_downloaded, paper.has_markdown, paper.abstract_text
|
||||
);
|
||||
ToolOutput::success(content, json!(paper))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("获取文献 {} 元数据失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
// src/agent/tools/target.rs — 天体物理查询工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 天体信息查询工具:通过 CDS Sesame 查询天体物理属性
|
||||
pub struct QueryTargetTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for QueryTargetTool {
|
||||
fn name(&self) -> &str {
|
||||
"query_target"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"查询天体的基本物理属性信息。输入天体名称,返回坐标 (RA/Dec)、视星等、光谱型、视差等属性。\
|
||||
数据来源为 CDS SIMBAD/Sesame 名称解析服务。适用于:获取天体基本参数、验证天体身份。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_name": {
|
||||
"type": "string",
|
||||
"description": "天体名称,如 'NGC 6752', 'GD 358', 'HD 209458', 'M 31' 等"
|
||||
}
|
||||
},
|
||||
"required": ["object_name"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读取外部天体数据库,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let object_name = match args.get("object_name").and_then(|n| n.as_str()) {
|
||||
Some(n) => n.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'object_name'"),
|
||||
};
|
||||
|
||||
info!("[QueryTarget] 查询天体: {}", object_name);
|
||||
let state = &ctx.app_state;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
match crate::services::target::query_target_cached(&state.db, &object_name, None, &client)
|
||||
.await
|
||||
{
|
||||
Ok(info) => {
|
||||
let content = format!(
|
||||
"天体: {}\nRA: {}\nDec: {}\n视差: {}\n光谱型: {}\nV星等: {}\n别名: {}",
|
||||
info.target_name,
|
||||
info.ra.as_deref().unwrap_or("未知"),
|
||||
info.dec.as_deref().unwrap_or("未知"),
|
||||
info.parallax
|
||||
.map(|p| format!("{:.4} mas", p))
|
||||
.unwrap_or_else(|| "未知".to_string()),
|
||||
info.spectral_type.as_deref().unwrap_or("未知"),
|
||||
info.v_magnitude
|
||||
.map(|v| format!("{:.2}", v))
|
||||
.unwrap_or_else(|| "未知".to_string()),
|
||||
if info.aliases.is_empty() {
|
||||
"无".to_string()
|
||||
} else {
|
||||
info.aliases.join(", ")
|
||||
}
|
||||
);
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({
|
||||
"target_name": info.target_name,
|
||||
"ra": info.ra, "dec": info.dec,
|
||||
"parallax": info.parallax,
|
||||
"spectral_type": info.spectral_type,
|
||||
"v_magnitude": info.v_magnitude,
|
||||
"aliases": info.aliases
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("天体 '{}' 查询失败: {}", object_name, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
// 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] = &["download_paper", "parse_paper"];
|
||||
|
||||
/// 后台任务启动工具
|
||||
pub struct BgTaskRunTool {
|
||||
queue: Arc<BgNotificationQueue>,
|
||||
}
|
||||
|
||||
impl BgTaskRunTool {
|
||||
pub fn new(queue: Arc<BgNotificationQueue>) -> Self {
|
||||
BgTaskRunTool { queue }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for BgTaskRunTool {
|
||||
fn name(&self) -> &str {
|
||||
"bg_task_run"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在后台异步执行慢速工具(download_paper, parse_paper)。\
|
||||
返回任务ID后立即让 LLM 继续思考,后台完成的结果会在下一轮对话中自动通知。\
|
||||
适用于:下载PDF、解析论文等耗时操作。使用 bg_task_check 查询任务状态。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tool_name": {
|
||||
"type": "string",
|
||||
"description": "要在后台执行的工具名称(download_paper 或 parse_paper)",
|
||||
"enum": ["download_paper", "parse_paper"]
|
||||
},
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符(ADS bibcode)"
|
||||
}
|
||||
},
|
||||
"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(", ")
|
||||
));
|
||||
}
|
||||
|
||||
info!(
|
||||
"[BgTaskRun] 启动后台任务: tool={}, bibcode={}",
|
||||
tool_name, bibcode
|
||||
);
|
||||
|
||||
let handle = background::spawn_background_task(
|
||||
ctx.app_state.clone(),
|
||||
self.queue.clone(),
|
||||
tool_name.clone(),
|
||||
bibcode.clone(),
|
||||
)
|
||||
.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"
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// 后台任务查询工具
|
||||
pub struct BgTaskCheckTool {
|
||||
queue: Arc<BgNotificationQueue>,
|
||||
}
|
||||
|
||||
impl BgTaskCheckTool {
|
||||
pub fn new(queue: Arc<BgNotificationQueue>) -> Self {
|
||||
BgTaskCheckTool { queue }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for BgTaskCheckTool {
|
||||
fn name(&self) -> &str {
|
||||
"bg_task_check"
|
||||
}
|
||||
|
||||
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 }))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
// src/agent/tools/compress.rs — 手动上下文压缩工具
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 手动上下文压缩工具:LLM 可主动调用以压缩对话历史
|
||||
pub struct CompressTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for CompressTool {
|
||||
fn name(&self) -> &str {
|
||||
"compress_context"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"手动压缩对话上下文。当你发现对话历史过长、token 消耗过大时,主动调用此工具进行压缩以释放空间。\
|
||||
压缩后历史对话将被摘要替代,但关键信息不会丢失。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
/// 设置标志位是幂等操作,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
// 实际压缩由 AgentRuntime 通过 pending_manual_compress 标志位处理
|
||||
ToolOutput::success(
|
||||
"上下文压缩标记已设置。当前对话历史将在下一轮 LLM 调用前被压缩。",
|
||||
json!({ "action": "compress" }),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::time::Duration;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
// ── run_bash ────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Bash 命令执行工具。
|
||||
///
|
||||
/// 允许 Agent 执行只读或数据处理的 Shell 命令。
|
||||
pub struct RunBashTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for RunBashTool {
|
||||
fn name(&self) -> &str {
|
||||
"run_bash"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"执行一个 Bash 命令并返回 stdout 和 stderr。\
|
||||
适用场景:运行 skill 目录中的 scripts/*.py、文件处理、数据提取。\
|
||||
限制:最大执行时间 60 秒,输出自动截断至 4000 字符。\
|
||||
禁止交互式命令和破坏性操作。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要执行的命令。如 'wc -l file.txt' 或 'python script.py'"
|
||||
},
|
||||
"working_dir": {
|
||||
"type": "string",
|
||||
"description": "工作目录(默认项目根目录)"
|
||||
},
|
||||
"timeout_secs": {
|
||||
"type": "integer",
|
||||
"description": "超时时间(秒),默认 60,最大 120",
|
||||
"default": 60
|
||||
}
|
||||
},
|
||||
"required": ["command"]
|
||||
})
|
||||
}
|
||||
|
||||
/// Bash 执行不应被中断(可能有写操作)
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let command = match args.get("command").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'command'"),
|
||||
};
|
||||
|
||||
// 命令安全校验
|
||||
if let Some(rejection) = validate_bash_command(command) {
|
||||
return ToolOutput::error(rejection);
|
||||
}
|
||||
|
||||
let timeout_secs: u64 = args
|
||||
.get("timeout_secs")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|v| (v.max(1) as u64).min(120))
|
||||
.unwrap_or(60);
|
||||
|
||||
// 工作目录
|
||||
let working_dir = match args.get("working_dir").and_then(|v| v.as_str()) {
|
||||
Some(dir_str) => {
|
||||
if has_path_traversal(dir_str) {
|
||||
return ToolOutput::error("工作目录路径包含非法字符");
|
||||
}
|
||||
match resolve_path(dir_str) {
|
||||
Some(p) if is_path_allowed(&p, ctx) => p,
|
||||
Some(_) => return ToolOutput::error("无权访问指定的工作目录"),
|
||||
None => return ToolOutput::error("无法解析工作目录路径"),
|
||||
}
|
||||
}
|
||||
None => match std::env::current_dir() {
|
||||
Ok(d) => d,
|
||||
Err(e) => return ToolOutput::error(format!("无法获取当前目录: {}", e)),
|
||||
},
|
||||
};
|
||||
|
||||
info!(
|
||||
"[RunBash] executing: {} (cwd: {}, timeout: {}s)",
|
||||
command,
|
||||
working_dir.display(),
|
||||
timeout_secs
|
||||
);
|
||||
|
||||
// 执行命令
|
||||
let timeout = Duration::from_secs(timeout_secs);
|
||||
let result = tokio::time::timeout(
|
||||
timeout,
|
||||
tokio::process::Command::new("bash")
|
||||
.arg("-c")
|
||||
.arg(command)
|
||||
.current_dir(&working_dir)
|
||||
.output(),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
let exit_code = output.status.code().unwrap_or(-1);
|
||||
|
||||
let mut content = if exit_code == 0 {
|
||||
String::new()
|
||||
} else {
|
||||
format!("[退出码: {}]\n", exit_code)
|
||||
};
|
||||
|
||||
if !stdout.trim().is_empty() {
|
||||
content.push_str(&stdout);
|
||||
}
|
||||
if !stderr.trim().is_empty() {
|
||||
if !content.is_empty() {
|
||||
content.push('\n');
|
||||
}
|
||||
content.push_str(&format!("[stderr]\n{}", stderr));
|
||||
}
|
||||
if content.is_empty() {
|
||||
content = "(无输出)".to_string();
|
||||
}
|
||||
|
||||
let output_truncated = crate::agent::tools::truncate_content(&content, 4000);
|
||||
|
||||
info!(
|
||||
"[RunBash] completed with exit code {} ({} chars output)",
|
||||
exit_code,
|
||||
content.len()
|
||||
);
|
||||
|
||||
ToolOutput::success(
|
||||
output_truncated,
|
||||
json!({
|
||||
"exit_code": exit_code,
|
||||
"stdout_length": stdout.len(),
|
||||
"stderr_length": stderr.len()
|
||||
}),
|
||||
)
|
||||
}
|
||||
Ok(Err(e)) => ToolOutput::error(format!("命令执行失败: {}", e)),
|
||||
Err(_) => ToolOutput::error(format!("命令执行超时 ({}s): {}", timeout_secs, command)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 验证 Bash 命令安全性(黑名单 + 启发式检查)。
|
||||
/// 返回 `Some(reason)` 表示拒绝,`None` 表示允许。
|
||||
/// 安全命令白名单:这些命令的第一个单词匹配时,自动允许(仍需路径沙箱检查)。
|
||||
#[allow(dead_code)]
|
||||
const SAFE_COMMANDS: &[&str] = &[
|
||||
"ls", "cat", "head", "tail", "find", "grep", "wc", "echo", "pwd", "sort", "uniq", "cut", "tr",
|
||||
"awk", "sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename",
|
||||
"dirname", "realpath", "readlink", "xargs", "tee", "date", "sleep", "true", "false",
|
||||
];
|
||||
|
||||
/// 检查命令是否属于安全白名单(第一个单词匹配即可)。
|
||||
#[allow(dead_code)]
|
||||
fn is_safe_command(command: &str) -> bool {
|
||||
let first_word = command.split_whitespace().next().unwrap_or("");
|
||||
SAFE_COMMANDS.contains(&first_word)
|
||||
}
|
||||
|
||||
fn validate_bash_command(command: &str) -> Option<String> {
|
||||
let trimmed = command.trim();
|
||||
|
||||
// 空命令
|
||||
if trimmed.is_empty() || trimmed == "bash" || trimmed == "bash -c" {
|
||||
return Some("不允许执行空命令".to_string());
|
||||
}
|
||||
|
||||
// 禁止交互式/破坏性命令
|
||||
let interactive_patterns = [
|
||||
"sudo ",
|
||||
"su ",
|
||||
"passwd",
|
||||
"ssh ",
|
||||
"telnet ",
|
||||
"login",
|
||||
"less ",
|
||||
"more ",
|
||||
"vim ",
|
||||
"vi ",
|
||||
"nano ",
|
||||
"emacs ",
|
||||
"top",
|
||||
"htop",
|
||||
"watch ",
|
||||
"tail -f",
|
||||
"rm -rf /",
|
||||
"mkfs.",
|
||||
"dd if=",
|
||||
"chmod 777",
|
||||
"> /dev/",
|
||||
];
|
||||
|
||||
let lower = trimmed.to_lowercase();
|
||||
for pattern in &interactive_patterns {
|
||||
if lower.contains(&pattern.to_lowercase()) {
|
||||
return Some(format!("不允许执行 '{}' 类命令", pattern));
|
||||
}
|
||||
}
|
||||
|
||||
// 允许通过
|
||||
None
|
||||
}
|
||||
|
||||
/// 检查命令是否需要用户权限确认。
|
||||
/// 安全白名单中的命令不需确认,其他命令需要。
|
||||
#[allow(dead_code)]
|
||||
pub fn bash_needs_permission(command: &str) -> bool {
|
||||
!is_safe_command(command)
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
// ── file_edit helpers ──────────────────────────────────────────────────────
|
||||
|
||||
/// 将弯引号(curly quotes)转换为直引号(straight quotes)。
|
||||
/// LLM 输出的是直引号,但文件中可能使用弯引号,需要统一后匹配。
|
||||
fn normalize_quotes(s: &str) -> String {
|
||||
s.replace(['\u{2018}', '\u{2019}'], "'") // right single curly
|
||||
.replace(['\u{201c}', '\u{201d}'], "\"") // right double curly
|
||||
}
|
||||
|
||||
/// 在文件内容中查找 old_string,兼容引号差异。
|
||||
///
|
||||
/// 先尝试精确匹配,失败后用引号规范化再试。
|
||||
/// 返回文件中的实际字符串(用于替换),如果没找到则返回 None。
|
||||
fn find_actual_string<'a>(file_content: &'a str, old_string: &str) -> Option<&'a str> {
|
||||
// 1. 精确匹配
|
||||
if let Some(pos) = file_content.find(old_string) {
|
||||
return Some(&file_content[pos..pos + old_string.len()]);
|
||||
}
|
||||
|
||||
// 2. 引号规范化后匹配。
|
||||
// 弯引号(3 bytes) 和直引号(1 byte) 的字节长度不同,所以不能直接用
|
||||
// normalized_file 中的字节位置去切 original file。改用字符位置对齐。
|
||||
let normalized_search = normalize_quotes(old_string);
|
||||
let normalized_file = normalize_quotes(file_content);
|
||||
|
||||
// 找到匹配在归一化后的文件中的字符偏移
|
||||
let norm_char_pos = normalized_file.find(&normalized_search)?;
|
||||
// 统计归一化文件中 norm_char_pos 字节对应的字符数
|
||||
let char_start: usize = normalized_file[..norm_char_pos].chars().count();
|
||||
let char_len: usize = old_string.chars().count();
|
||||
|
||||
// 在原文件中找到对应的字节范围
|
||||
let orig_chars: Vec<(usize, char)> = file_content.char_indices().collect();
|
||||
if char_start + char_len > orig_chars.len() {
|
||||
return None;
|
||||
}
|
||||
let start_byte = orig_chars[char_start].0;
|
||||
let end_byte = if char_start + char_len < orig_chars.len() {
|
||||
orig_chars[char_start + char_len].0
|
||||
} else {
|
||||
file_content.len()
|
||||
};
|
||||
|
||||
Some(&file_content[start_byte..end_byte])
|
||||
}
|
||||
|
||||
/// 当 old_string 通过引号规范化才匹配成功时,
|
||||
/// 对 new_string 施加相同的弯引号风格,保持文件风格一致。
|
||||
fn preserve_quote_style(old_string: &str, actual_old: &str, new_string: &str) -> String {
|
||||
if old_string == actual_old {
|
||||
return new_string.to_string();
|
||||
}
|
||||
|
||||
let has_curly_single = actual_old.contains('\u{2018}') || actual_old.contains('\u{2019}');
|
||||
let has_curly_double = actual_old.contains('\u{201c}') || actual_old.contains('\u{201d}');
|
||||
|
||||
if !has_curly_single && !has_curly_double {
|
||||
return new_string.to_string();
|
||||
}
|
||||
|
||||
let mut result = new_string.to_string();
|
||||
|
||||
if has_curly_double {
|
||||
result = apply_curly_double_quotes(&result);
|
||||
}
|
||||
if has_curly_single {
|
||||
result = apply_curly_single_quotes(&result);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// 将直双引号替换为弯双引号(根据上下文判断开/闭)
|
||||
fn apply_curly_double_quotes(s: &str) -> String {
|
||||
let chars: Vec<char> = s.chars().collect();
|
||||
let mut result = String::with_capacity(s.len());
|
||||
for (i, &ch) in chars.iter().enumerate() {
|
||||
if ch == '"' {
|
||||
if is_opening_context(&chars, i) {
|
||||
result.push('\u{201c}'); // left double curly
|
||||
} else {
|
||||
result.push('\u{201d}'); // right double curly
|
||||
}
|
||||
} else {
|
||||
result.push(ch);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// 将直单引号替换为弯单引号(跳过缩略形式如 don't, it's)
|
||||
fn apply_curly_single_quotes(s: &str) -> String {
|
||||
let chars: Vec<char> = s.chars().collect();
|
||||
let mut result = String::with_capacity(s.len());
|
||||
for (i, &ch) in chars.iter().enumerate() {
|
||||
if ch == '\'' {
|
||||
let prev_is_letter = i > 0 && chars[i - 1].is_alphabetic();
|
||||
let next_is_letter = i + 1 < chars.len() && chars[i + 1].is_alphabetic();
|
||||
if prev_is_letter && next_is_letter {
|
||||
// 缩略形式 (don't, it's) — 使用右弯单引号
|
||||
result.push('\u{2019}');
|
||||
} else if is_opening_context(&chars, i) {
|
||||
result.push('\u{2018}'); // left single curly
|
||||
} else {
|
||||
result.push('\u{2019}'); // right single curly
|
||||
}
|
||||
} else {
|
||||
result.push(ch);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn is_opening_context(chars: &[char], index: usize) -> bool {
|
||||
if index == 0 {
|
||||
return true;
|
||||
}
|
||||
let prev = chars[index - 1];
|
||||
matches!(prev, ' ' | '\t' | '\n' | '\r' | '(' | '[' | '{')
|
||||
}
|
||||
|
||||
/// 去掉每行末尾的空白字符(保留换行符)。
|
||||
/// 对非 .md/.mdx 文件使用,因为 Markdown 中两个尾随空格表示硬换行。
|
||||
fn strip_trailing_whitespace(s: &str) -> String {
|
||||
let mut result = String::with_capacity(s.len());
|
||||
for line in s.lines() {
|
||||
let trimmed = line.trim_end();
|
||||
result.push_str(trimmed);
|
||||
result.push('\n');
|
||||
}
|
||||
// 如果原字符串不以换行结尾,去掉我们添加的换行
|
||||
if !s.ends_with('\n') && !result.is_empty() {
|
||||
result.pop();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
// ── file_edit ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// 文件编辑工具(精确字符串替换)。
|
||||
///
|
||||
/// 在已有文件中查找 old_string 并替换为 new_string。
|
||||
/// old_string 必须唯一匹配(防止误改)。
|
||||
pub struct FileEditTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for FileEditTool {
|
||||
fn name(&self) -> &str {
|
||||
"file_edit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在文件中执行精确的字符串替换。查找 old_string 并替换为 new_string。\
|
||||
old_string 在文件中必须唯一(仅出现一次),以防止意外破坏其他内容。\
|
||||
适用场景:修改脚本参数、更新配置值、在 Markdown 笔记中追加或修正内容。\
|
||||
注意:仅替换匹配片段,文件其余部分保持不变。如需完整覆写请使用 file_write。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "要编辑的文件路径"
|
||||
},
|
||||
"old_string": {
|
||||
"type": "string",
|
||||
"description": "要被替换的原字符串(必须唯一匹配)"
|
||||
},
|
||||
"new_string": {
|
||||
"type": "string",
|
||||
"description": "替换后的新字符串"
|
||||
},
|
||||
"replace_all": {
|
||||
"type": "boolean",
|
||||
"description": "替换所有匹配(默认 false,要求 old_string 唯一)",
|
||||
"default": false
|
||||
}
|
||||
},
|
||||
"required": ["file_path", "old_string", "new_string"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let path_str = match args.get("file_path").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'file_path'"),
|
||||
};
|
||||
let old_string = match args.get("old_string").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'old_string'"),
|
||||
};
|
||||
let new_string = match args.get("new_string").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'new_string'"),
|
||||
};
|
||||
let replace_all = args
|
||||
.get("replace_all")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_path_traversal(path_str) {
|
||||
return ToolOutput::error("路径包含非法字符(.. 或 ~),拒绝访问");
|
||||
}
|
||||
|
||||
let file_path = match resolve_path(path_str) {
|
||||
Some(p) => p,
|
||||
None => return ToolOutput::error("无法解析文件路径"),
|
||||
};
|
||||
|
||||
if !is_path_allowed(&file_path, ctx) {
|
||||
return ToolOutput::error(format!("无权编辑路径 '{}'", path_str));
|
||||
}
|
||||
|
||||
if !file_path.exists() {
|
||||
return ToolOutput::error(format!("文件不存在: {}", file_path.display()));
|
||||
}
|
||||
|
||||
if old_string.is_empty() {
|
||||
return ToolOutput::error("old_string 不能为空");
|
||||
}
|
||||
|
||||
if old_string == new_string {
|
||||
return ToolOutput::success(
|
||||
"未做任何更改:old_string 与 new_string 完全相同。",
|
||||
json!({ "file_path": file_path.to_string_lossy(), "changed": false }),
|
||||
);
|
||||
}
|
||||
|
||||
let original = match std::fs::read_to_string(&file_path) {
|
||||
Ok(c) => c,
|
||||
Err(e) => return ToolOutput::error(format!("读取文件失败: {}", e)),
|
||||
};
|
||||
|
||||
// 对非 Markdown 文件去除 LLM 输出中常见的尾随空白
|
||||
let is_markdown = file_path
|
||||
.extension()
|
||||
.map(|e| e == "md" || e == "mdx")
|
||||
.unwrap_or(false);
|
||||
let (old_string, new_string) = if is_markdown {
|
||||
(old_string, new_string)
|
||||
} else {
|
||||
(
|
||||
strip_trailing_whitespace(&old_string),
|
||||
strip_trailing_whitespace(&new_string),
|
||||
)
|
||||
};
|
||||
|
||||
// 查找文件中的实际字符串(兼容引号差异)
|
||||
let actual_old = match find_actual_string(&original, &old_string) {
|
||||
Some(s) => s.to_string(),
|
||||
None => {
|
||||
return ToolOutput::error(format!(
|
||||
"在文件中未找到要替换的字符串。\n查找内容: {}\n提示: 请确认原文内容完全一致(含空格/换行),\
|
||||
或使用 read_file 重新读取文件确认当前内容。",
|
||||
old_string
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 保持文件的引号风格
|
||||
let actual_new = preserve_quote_style(&old_string, &actual_old, &new_string);
|
||||
|
||||
// 检查匹配次数
|
||||
let match_count = original.matches(&actual_old).count();
|
||||
|
||||
if !replace_all && match_count > 1 {
|
||||
// 提取上下文帮助 LLM 定位
|
||||
let mut ctx_lines: Vec<String> = Vec::new();
|
||||
for (line_no, line) in original.lines().enumerate() {
|
||||
if line.contains(&actual_old) {
|
||||
ctx_lines.push(format!(" L{}: {}", line_no + 1, line.trim()));
|
||||
}
|
||||
if ctx_lines.len() >= 5 {
|
||||
break;
|
||||
} // 最多显示 5 处
|
||||
}
|
||||
return ToolOutput::error(format!(
|
||||
"'{}' 在文件中出现了 {} 处,无法确定要修改哪一个。\n\
|
||||
出现位置:\n{}\n\
|
||||
请包含更多上下文(前后行)使 old_string 唯一,
|
||||
或设置 replace_all: true 以替换全部匹配。",
|
||||
actual_old,
|
||||
match_count,
|
||||
ctx_lines.join("\n")
|
||||
));
|
||||
}
|
||||
|
||||
let modified = if replace_all {
|
||||
original.replace(&actual_old, &actual_new)
|
||||
} else {
|
||||
original.replacen(&actual_old, &actual_new, 1)
|
||||
};
|
||||
|
||||
if modified == original {
|
||||
return ToolOutput::success(
|
||||
"文件内容未发生变化(new_string 与 old_string 相同)。",
|
||||
json!({ "file_path": file_path.to_string_lossy(), "changed": false }),
|
||||
);
|
||||
}
|
||||
|
||||
match std::fs::write(&file_path, &modified) {
|
||||
Ok(()) => {
|
||||
let count = if replace_all { match_count } else { 1 };
|
||||
info!("[FileEdit] {} 处替换: {}", count, file_path.display());
|
||||
ToolOutput::success(
|
||||
format!("文件 {} 已更新({} 处替换)。", file_path.display(), count),
|
||||
json!({
|
||||
"file_path": file_path.to_string_lossy(),
|
||||
"replacements": count,
|
||||
"changed": true
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("写入失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
// src/agent/tools/filesystem/glob.rs
|
||||
//
|
||||
// glob_files 工具 — 基于 glob 模式查找文件。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::truncate_content;
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 基于 glob 模式查找文件工具。
|
||||
pub struct GlobFilesTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GlobFilesTool {
|
||||
fn name(&self) -> &str {
|
||||
"glob_files"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"使用 glob 模式查找匹配的文件。支持通配符 * 和 **。返回匹配的文件路径列表。路径必须在允许的沙箱范围内。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {
|
||||
"type": "string",
|
||||
"description": "glob 匹配模式(如 '**/*.rs', 'src/**/*.md')"
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "搜索起始路径,默认为当前工作目录",
|
||||
"default": "."
|
||||
}
|
||||
},
|
||||
"required": ["pattern"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let pattern = match args.get("pattern").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'pattern'"),
|
||||
};
|
||||
|
||||
let search_path_str = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
|
||||
|
||||
if has_path_traversal(search_path_str) {
|
||||
return ToolOutput::error("路径包含不安全字符");
|
||||
}
|
||||
|
||||
let search_path = match resolve_path(search_path_str) {
|
||||
Some(p) => p,
|
||||
None => return ToolOutput::error(format!("无法解析路径: {}", search_path_str)),
|
||||
};
|
||||
|
||||
if !is_path_allowed(&search_path, ctx) {
|
||||
return ToolOutput::error("路径不在允许的沙箱范围内");
|
||||
}
|
||||
|
||||
let glob_pattern = search_path.join(pattern);
|
||||
let pattern_str = glob_pattern.to_string_lossy().to_string();
|
||||
|
||||
match glob::glob(&pattern_str) {
|
||||
Ok(paths) => {
|
||||
let results: Vec<String> = paths
|
||||
.flatten()
|
||||
.filter(|p| p.is_file())
|
||||
.filter_map(|p| {
|
||||
let relative = p.strip_prefix(&search_path).ok()?;
|
||||
Some(relative.display().to_string())
|
||||
})
|
||||
.take(200)
|
||||
.collect();
|
||||
|
||||
if results.is_empty() {
|
||||
ToolOutput::success(
|
||||
format!("未找到匹配 '{}' 的文件", pattern),
|
||||
json!({"pattern": pattern, "matches": 0}),
|
||||
)
|
||||
} else {
|
||||
let count = results.len();
|
||||
let output = results.join("\n");
|
||||
let truncated = truncate_content(&output, 4000);
|
||||
ToolOutput::success(truncated, json!({"pattern": pattern, "matches": count}))
|
||||
}
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("glob 模式无效: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
// src/agent/tools/filesystem/grep.rs
|
||||
//
|
||||
// grep_files 工具 — 在文件中搜索匹配模式的行(类似 grep 命令)。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::truncate_content;
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 在文件中搜索匹配模式的行(类似 grep 命令)。
|
||||
pub struct GrepFilesTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GrepFilesTool {
|
||||
fn name(&self) -> &str {
|
||||
"grep_files"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在指定目录或文件中搜索匹配正则表达式的行。返回匹配行及其上下文(前后各 2 行)。适用于在代码库或文献中搜索特定模式。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {
|
||||
"type": "string",
|
||||
"description": "要搜索的正则表达式模式"
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "搜索路径(文件或目录),默认为当前工作目录",
|
||||
"default": "."
|
||||
},
|
||||
"include": {
|
||||
"type": "string",
|
||||
"description": "文件过滤 glob 模式(如 '*.rs', '*.md'),默认所有文本文件",
|
||||
"default": null
|
||||
},
|
||||
"max_results": {
|
||||
"type": "integer",
|
||||
"description": "最大结果数,默认 50",
|
||||
"default": 50
|
||||
}
|
||||
},
|
||||
"required": ["pattern"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let pattern = match args.get("pattern").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'pattern'"),
|
||||
};
|
||||
|
||||
let search_path_str = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
|
||||
|
||||
if has_path_traversal(search_path_str) {
|
||||
return ToolOutput::error("路径包含不安全字符");
|
||||
}
|
||||
|
||||
let search_path = match resolve_path(search_path_str) {
|
||||
Some(p) => p,
|
||||
None => return ToolOutput::error(format!("无法解析路径: {}", search_path_str)),
|
||||
};
|
||||
|
||||
if !is_path_allowed(&search_path, ctx) {
|
||||
return ToolOutput::error("路径不在允许的沙箱范围内");
|
||||
}
|
||||
|
||||
let include_pattern = args.get("include").and_then(|v| v.as_str());
|
||||
let max_results = args
|
||||
.get("max_results")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(50) as usize;
|
||||
|
||||
let re = match regex::Regex::new(pattern) {
|
||||
Ok(r) => r,
|
||||
Err(e) => return ToolOutput::error(format!("正则表达式无效: {}", e)),
|
||||
};
|
||||
|
||||
let mut results = Vec::new();
|
||||
let mut count = 0;
|
||||
|
||||
if search_path.is_file() {
|
||||
count += Self::search_file(&search_path, &re, &mut results);
|
||||
} else if search_path.is_dir() {
|
||||
count += Self::search_dir(
|
||||
&search_path,
|
||||
&re,
|
||||
include_pattern,
|
||||
max_results,
|
||||
&mut results,
|
||||
);
|
||||
}
|
||||
|
||||
if results.is_empty() {
|
||||
ToolOutput::success(
|
||||
format!("未找到匹配 '{}' 的结果 (搜索了 {} 个位置)", pattern, count),
|
||||
json!({"pattern": pattern, "matches": 0, "files_searched": count}),
|
||||
)
|
||||
} else {
|
||||
let output = results.join("\n---\n");
|
||||
let truncated = truncate_content(&output, 4000);
|
||||
ToolOutput::success(
|
||||
truncated,
|
||||
json!({"pattern": pattern, "matches": results.len(), "files_searched": count}),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GrepFilesTool {
|
||||
fn search_file(path: &Path, re: ®ex::Regex, results: &mut Vec<String>) -> usize {
|
||||
let content = match std::fs::read_to_string(path) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return 0,
|
||||
};
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
let mut matched = false;
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
if re.is_match(line) {
|
||||
if !matched {
|
||||
results.push(format!("📄 {}:", path.display()));
|
||||
matched = true;
|
||||
}
|
||||
let ctx_start = i.saturating_sub(2);
|
||||
let ctx_end = (i + 3).min(lines.len());
|
||||
for (j, line) in lines.iter().enumerate().take(ctx_end).skip(ctx_start) {
|
||||
let marker = if j == i { ">" } else { " " };
|
||||
results.push(format!(" {} {:4}: {}", marker, j + 1, line));
|
||||
}
|
||||
results.push(String::new());
|
||||
}
|
||||
}
|
||||
1
|
||||
}
|
||||
|
||||
fn search_dir(
|
||||
dir: &Path,
|
||||
re: ®ex::Regex,
|
||||
include: Option<&str>,
|
||||
max_results: usize,
|
||||
results: &mut Vec<String>,
|
||||
) -> usize {
|
||||
let mut count = 0;
|
||||
let text_extensions = [
|
||||
"rs", "py", "js", "ts", "jsx", "tsx", "html", "css", "md", "txt", "json", "yaml",
|
||||
"yml", "toml", "sh", "sql", "c", "cpp", "h", "hpp", "java", "go", "rb", "php", "swift",
|
||||
];
|
||||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
if results.len() >= max_results {
|
||||
break;
|
||||
}
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
let dir_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if !dir_name.starts_with('.')
|
||||
&& dir_name != "target"
|
||||
&& dir_name != "node_modules"
|
||||
{
|
||||
count += Self::search_dir(&path, re, include, max_results, results);
|
||||
}
|
||||
} else if path.is_file() {
|
||||
if let Some(inc) = include {
|
||||
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if !glob_match(inc, file_name) {
|
||||
continue;
|
||||
}
|
||||
} else {
|
||||
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
|
||||
if !text_extensions.contains(&ext) && !ext.is_empty() {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
count += Self::search_file(&path, re, results);
|
||||
}
|
||||
}
|
||||
}
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
fn glob_match(pattern: &str, name: &str) -> bool {
|
||||
if pattern == "*" {
|
||||
return true;
|
||||
}
|
||||
if pattern.starts_with("*.") {
|
||||
return name.ends_with(&pattern[1..]);
|
||||
}
|
||||
name == pattern
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
// src/agent/tools/filesystem/mod.rs
|
||||
//
|
||||
// 文件系统访问工具集 — 每个工具独立子模块,mod.rs 仅做 re-export。
|
||||
//
|
||||
// 安全约束由 security 子模块统一提供:
|
||||
// 1. 路径沙箱:只允许访问 library_dir、skills_dir 及项目根目录
|
||||
// 2. 路径穿越防护:拒绝含 ".." 的路径
|
||||
// 3. Bash 超时 + 输出截断
|
||||
|
||||
mod bash;
|
||||
mod edit;
|
||||
mod glob;
|
||||
mod grep;
|
||||
mod read;
|
||||
pub mod security;
|
||||
mod write;
|
||||
|
||||
pub use bash::RunBashTool;
|
||||
pub use edit::FileEditTool;
|
||||
pub use glob::GlobFilesTool;
|
||||
pub use grep::GrepFilesTool;
|
||||
pub use read::ReadFileTool;
|
||||
pub use write::FileWriteTool;
|
||||
@@ -0,0 +1,165 @@
|
||||
// src/agent/tools/filesystem/read.rs
|
||||
//
|
||||
// read_file 工具 — 读取文件内容。
|
||||
// P1 增强:文件状态缓存去重(参考 Claude Code FileReadTool + FileStateCache)。
|
||||
// 读取前检查缓存和 mtime,相同则返回 FILE_UNCHANGED_STUB。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::runtime::file_cache::{self, FileState};
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 读取文件内容工具。
|
||||
pub struct ReadFileTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for ReadFileTool {
|
||||
fn name(&self) -> &str {
|
||||
"read_file"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"读取指定路径的文件内容。支持文本文件和代码文件。会自动截断过长的内容。路径必须在允许的沙箱范围内。如果文件自上次读取后未修改,会返回占位消息以节省上下文。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "要读取的文件路径(绝对路径或相对于当前工作目录的路径)"
|
||||
},
|
||||
"max_lines": {
|
||||
"type": "integer",
|
||||
"description": "最大读取行数,默认全部",
|
||||
"default": null
|
||||
}
|
||||
},
|
||||
"required": ["file_path"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'file_path'"),
|
||||
};
|
||||
|
||||
if has_path_traversal(file_path_str) {
|
||||
return ToolOutput::error("路径包含不安全字符 (.. 或 ~)");
|
||||
}
|
||||
|
||||
let path = match resolve_path(file_path_str) {
|
||||
Some(p) => p,
|
||||
None => return ToolOutput::error(format!("无法解析路径: {}", file_path_str)),
|
||||
};
|
||||
|
||||
if !is_path_allowed(&path, ctx) {
|
||||
return ToolOutput::error(format!("路径不在允许的沙箱范围内: {}", file_path_str));
|
||||
}
|
||||
|
||||
let max_lines = args
|
||||
.get("max_lines")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(i64::MAX) as usize;
|
||||
|
||||
let offset = 1usize; // 始终从第一行开始(与 Claude Code 不同,我们暂不暴露 offset 参数)
|
||||
let limit = if max_lines == usize::MAX {
|
||||
None
|
||||
} else {
|
||||
Some(max_lines)
|
||||
};
|
||||
|
||||
// ── 文件状态缓存去重 ──
|
||||
// 检查缓存中是否有此文件的记录,且 mtime 未变更。
|
||||
if let Ok(mut cache) = ctx.read_file_state.lock() {
|
||||
let display_path = path.to_string_lossy().to_string();
|
||||
|
||||
if let Some(cached) = cache.get(&display_path) {
|
||||
// 检查 offset/limit 是否匹配
|
||||
let range_match = cached.offset == offset && cached.limit == limit;
|
||||
|
||||
if range_match {
|
||||
// 获取磁盘上的当前 mtime 与缓存的 timestamp 比较
|
||||
let disk_mtime = file_cache::get_file_mtime(&display_path);
|
||||
if let Some(mtime) = disk_mtime {
|
||||
if mtime == cached.timestamp {
|
||||
info!(
|
||||
"[ReadFile] 文件未修改,返回 stub: {} (mtime={})",
|
||||
file_path_str, mtime
|
||||
);
|
||||
return ToolOutput::success(
|
||||
file_cache::FILE_UNCHANGED_STUB.to_string(),
|
||||
json!({
|
||||
"file_path": file_path_str,
|
||||
"dedup": true,
|
||||
"mtime": mtime,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// ── 去重检查结束 ──
|
||||
|
||||
match std::fs::read_to_string(&path) {
|
||||
Ok(content) => {
|
||||
let total_lines = content.lines().count();
|
||||
let truncated = if max_lines < total_lines {
|
||||
content
|
||||
.lines()
|
||||
.take(max_lines)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
+ &format!(
|
||||
"\n\n[... 已截断,共 {} 行,显示前 {} 行 ...]",
|
||||
total_lines, max_lines
|
||||
)
|
||||
} else {
|
||||
content.clone()
|
||||
};
|
||||
|
||||
// 截断到单次输出上限
|
||||
let truncated_content = crate::agent::tools::truncate_content(&truncated, 4000);
|
||||
let content_len = truncated_content.len();
|
||||
info!(
|
||||
"[ReadFile] 读取 {}: {} 字符 ({} 行)",
|
||||
file_path_str, content_len, total_lines
|
||||
);
|
||||
|
||||
// ── 更新文件状态缓存 ──
|
||||
if let Ok(mut cache) = ctx.read_file_state.lock() {
|
||||
let display_path = path.to_string_lossy().to_string();
|
||||
let mtime = file_cache::get_file_mtime(&display_path).unwrap_or(0);
|
||||
cache.set(
|
||||
&display_path,
|
||||
FileState {
|
||||
content: content.clone(),
|
||||
timestamp: mtime,
|
||||
offset,
|
||||
limit,
|
||||
},
|
||||
);
|
||||
}
|
||||
// ── 缓存更新结束 ──
|
||||
|
||||
ToolOutput::success(
|
||||
truncated_content,
|
||||
json!({"file_path": file_path_str, "total_lines": total_lines}),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("读取文件失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
// src/agent/tools/filesystem/security.rs
|
||||
//
|
||||
// 文件系统操作的路径安全检查。
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::agent::tools::ToolContext;
|
||||
|
||||
/// 检查路径是否在允许的沙箱范围内。
|
||||
pub fn is_path_allowed(path: &Path, ctx: &ToolContext) -> bool {
|
||||
let config = &ctx.app_state.config;
|
||||
let canonical = match path.canonicalize() {
|
||||
Ok(p) => p,
|
||||
Err(_) => match path.parent() {
|
||||
Some(parent) => match parent.canonicalize() {
|
||||
Ok(p) => p,
|
||||
Err(_) => return false,
|
||||
},
|
||||
None => return false,
|
||||
},
|
||||
};
|
||||
let allowed_roots = [
|
||||
config.library_dir.canonicalize().ok(),
|
||||
config.skills_dir.canonicalize().ok(),
|
||||
std::env::current_dir().ok(),
|
||||
];
|
||||
for root in allowed_roots.iter().flatten() {
|
||||
if canonical.starts_with(root) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 检查路径字符串是否包含穿越尝试
|
||||
pub fn has_path_traversal(path_str: &str) -> bool {
|
||||
path_str.contains("..") || path_str.contains('~')
|
||||
}
|
||||
|
||||
/// 规范化用户提供的路径
|
||||
pub fn resolve_path(path_str: &str) -> Option<PathBuf> {
|
||||
let path = Path::new(path_str);
|
||||
if path.is_absolute() {
|
||||
Some(path.to_path_buf())
|
||||
} else {
|
||||
std::env::current_dir().ok().map(|cwd| cwd.join(path))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
// src/agent/tools/filesystem/write.rs
|
||||
//
|
||||
// file_write 工具 — 将内容写入文件。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::filesystem::security::{
|
||||
has_path_traversal, is_path_allowed, resolve_path,
|
||||
};
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
/// 写入文件内容工具。
|
||||
pub struct FileWriteTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for FileWriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"file_write"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将内容写入指定路径的文件。如果文件已存在则覆盖。路径必须在允许的沙箱范围内。适用于保存研究结果、生成报告等。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "要写入的文件路径"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "要写入的文件内容"
|
||||
}
|
||||
},
|
||||
"required": ["file_path", "content"]
|
||||
})
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'file_path'"),
|
||||
};
|
||||
let content = match args.get("content").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'content'"),
|
||||
};
|
||||
|
||||
if has_path_traversal(file_path_str) {
|
||||
return ToolOutput::error("路径包含不安全字符 (.. 或 ~)");
|
||||
}
|
||||
|
||||
let path = match resolve_path(file_path_str) {
|
||||
Some(p) => p,
|
||||
None => return ToolOutput::error(format!("无法解析路径: {}", file_path_str)),
|
||||
};
|
||||
|
||||
if !is_path_allowed(&path, ctx) {
|
||||
return ToolOutput::error(format!("路径不在允许的沙箱范围内: {}", file_path_str));
|
||||
}
|
||||
|
||||
if let Some(parent) = path.parent() {
|
||||
if let Err(e) = std::fs::create_dir_all(parent) {
|
||||
return ToolOutput::error(format!("创建父目录失败: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
match std::fs::write(&path, &content) {
|
||||
Ok(_) => {
|
||||
let line_count = content.lines().count();
|
||||
info!("[FileWrite] 写入 {}: {} 行", file_path_str, line_count);
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"成功写入文件: {} ({} 行, {} 字符)",
|
||||
file_path_str,
|
||||
line_count,
|
||||
content.len()
|
||||
),
|
||||
json!({"file_path": file_path_str, "lines": line_count, "bytes": content.len()}),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("写入文件失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
// src/agent/tools/memory.rs
|
||||
//
|
||||
// save_memory 工具 — 让 Agent 可以将重要信息持久化到项目记忆系统。
|
||||
// 参考 Claude Code memdir 设计。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::info;
|
||||
|
||||
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
use crate::agent::memory::dedup;
|
||||
use crate::agent::memory::types::MemoryType;
|
||||
use crate::agent::memory::MemoryManager;
|
||||
|
||||
pub struct SaveMemoryTool {
|
||||
memory_manager: Arc<Mutex<MemoryManager>>,
|
||||
}
|
||||
|
||||
impl SaveMemoryTool {
|
||||
pub fn new(memory_manager: Arc<Mutex<MemoryManager>>) -> Self {
|
||||
SaveMemoryTool { memory_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SaveMemoryTool {
|
||||
fn name(&self) -> &str {
|
||||
"save_memory"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将重要信息保存到项目记忆系统。记忆会跨会话持久化,在后续会话中自动加载。\
|
||||
用于保存:用户偏好、研究方法论、项目进展、重要发现、外部参考。\n\n\
|
||||
不应保存:代码模式、架构详情、Git 历史、调试方案、已记录在 CLAUDE.md 中的内容、\
|
||||
临时任务状态。即使用户要求保存以上内容,请先询问哪些部分是非预期的。\n\n\
|
||||
系统会自动检测内容重复和低质量输入。保存前先检查是否有可更新的现有记忆 — 不要写重复项。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"slug": {
|
||||
"type": "string",
|
||||
"description": "记忆标识符(短横线命名,如 'user-prefs')"
|
||||
},
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "记忆标题"
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": "简短描述(用于决定何时加载此记忆)"
|
||||
},
|
||||
"memory_type": {
|
||||
"type": "string",
|
||||
"enum": ["user", "feedback", "project", "reference"],
|
||||
"description": "记忆类型:user=用户偏好, feedback=用户反馈, project=项目进展, reference=外部参考"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "记忆内容(Markdown 格式)"
|
||||
}
|
||||
},
|
||||
"required": ["slug", "name", "description", "memory_type", "content"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
false // 写入操作,不并发安全
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block // 写入操作不可中断
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let slug = args["slug"].as_str().unwrap_or("");
|
||||
let name = args["name"].as_str().unwrap_or("");
|
||||
let description = args["description"].as_str().unwrap_or("");
|
||||
let memory_type_str = args["memory_type"].as_str().unwrap_or("user");
|
||||
let content = args["content"].as_str().unwrap_or("");
|
||||
|
||||
if slug.is_empty() || name.is_empty() || content.is_empty() {
|
||||
return ToolOutput::error("slug, name, content 均为必填项");
|
||||
}
|
||||
|
||||
// Validate slug format (kebab-case)
|
||||
if slug.contains(' ') || slug.contains('/') || slug.contains('\\') {
|
||||
return ToolOutput::error("slug 不能包含空格、斜杠或反斜杠");
|
||||
}
|
||||
|
||||
let memory_type = match MemoryType::from_str(memory_type_str) {
|
||||
Some(t) => t,
|
||||
None => {
|
||||
return ToolOutput::error(format!(
|
||||
"无效的 memory_type: {}。有效值: user, feedback, project, reference",
|
||||
memory_type_str
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let mut mgr = self.memory_manager.lock().await;
|
||||
|
||||
// ── 写入时门控 ──
|
||||
|
||||
// 1. 内容质量检查(仅警告,不拒绝)
|
||||
let quality = dedup::check_content_quality(content);
|
||||
let quality_warning = match &quality {
|
||||
dedup::QualityCheck::TooShort(n) => {
|
||||
Some(format!("⚠️ 内容偏短 ({} 字符),建议展开说明", n))
|
||||
}
|
||||
dedup::QualityCheck::TransientState => {
|
||||
Some("⚠️ 检测到瞬时状态描述,建议仅保存长期有价值的信息".to_string())
|
||||
}
|
||||
dedup::QualityCheck::VagueLanguage(w) => {
|
||||
Some(format!("⚠️ 检测到模糊语言 '{}',建议使用明确表述", w))
|
||||
}
|
||||
dedup::QualityCheck::CodePattern => {
|
||||
Some("⚠️ 检测到代码片段 — 代码模式不应保存为记忆".to_string())
|
||||
}
|
||||
dedup::QualityCheck::Accept => None,
|
||||
};
|
||||
|
||||
// 2. Jaccard 内容重复检测
|
||||
let duplicate_slug = dedup::find_duplicate_by_content(content, mgr.entries(), 0.70);
|
||||
|
||||
// 检查 slug 是否已存在
|
||||
let is_update = dedup::slug_exists(mgr.memory_dir(), slug);
|
||||
|
||||
match mgr.save_memory(slug, name, description, memory_type, content) {
|
||||
Ok(_) => {
|
||||
info!("[SaveMemory] 已保存记忆: {} ({})", name, slug);
|
||||
// 标记主代理已写入,抑制本会话的自动提取
|
||||
mgr.mark_main_agent_wrote();
|
||||
// 构建现有记忆清单供 LLM 参考
|
||||
let manifest = dedup::build_manifest_preview(mgr.entries());
|
||||
let action = if is_update {
|
||||
"🔄 已更新"
|
||||
} else {
|
||||
"✅ 已保存"
|
||||
};
|
||||
let mut message = format!(
|
||||
"{} 记忆: {} ({}) — 类型: {}",
|
||||
action, name, slug, memory_type_str
|
||||
);
|
||||
|
||||
// 附加质量警告
|
||||
let has_quality_warning = quality_warning.is_some();
|
||||
if let Some(w) = &quality_warning {
|
||||
message.push_str(&format!("\n\n{}", w));
|
||||
}
|
||||
|
||||
// 附加重复检测信息
|
||||
if let Some(ref dup_slug) = duplicate_slug {
|
||||
message.push_str(&format!(
|
||||
"\n\n💡 检测到与现有记忆 `{}` 内容接近(≥70% 重叠),请考虑更新该文件而非创建新的。",
|
||||
dup_slug
|
||||
));
|
||||
}
|
||||
|
||||
message.push_str(&format!("\n\n{}", manifest));
|
||||
|
||||
if !is_update && duplicate_slug.is_none() {
|
||||
message
|
||||
.push_str("\n\n💡 提示:如有其他重要信息需保存,请继续使用 save_memory。");
|
||||
}
|
||||
ToolOutput::success(
|
||||
message,
|
||||
json!({
|
||||
"slug": slug,
|
||||
"name": name,
|
||||
"memory_type": memory_type_str,
|
||||
"is_update": is_update,
|
||||
"quality_check": has_quality_warning,
|
||||
"duplicate_detected": duplicate_slug.is_some()
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("保存记忆失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,417 @@
|
||||
// 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::DelegateResearchTool;
|
||||
pub use team::{CheckTeamInboxTool, SendTeammateMessageTool, SpawnTeammateTool, TeamBroadcastTool};
|
||||
pub use todo::persist_tasks;
|
||||
pub use todo::TodoWriteTool;
|
||||
|
||||
/// 工具执行上下文,封装全局共享状态
|
||||
pub struct ToolContext {
|
||||
pub app_state: Arc<AppState>,
|
||||
/// 静默模式:子代理运行时为 true,跳过用户权限提示
|
||||
pub silent: bool,
|
||||
/// 文件状态缓存(跨工具调用共享,用于 Read 去重)
|
||||
pub read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
}
|
||||
|
||||
impl ToolContext {
|
||||
/// 创建标准上下文
|
||||
pub fn new(app_state: Arc<AppState>) -> Self {
|
||||
ToolContext {
|
||||
app_state,
|
||||
silent: false,
|
||||
read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带共享文件缓存的上下文(用于 AgentRuntime 保持同一 cache 实例)
|
||||
pub fn with_file_cache(
|
||||
app_state: Arc<AppState>,
|
||||
read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
) -> Self {
|
||||
ToolContext {
|
||||
app_state,
|
||||
silent: false,
|
||||
read_file_state,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建静默上下文(子代理使用)
|
||||
pub fn silent(app_state: Arc<AppState>) -> Self {
|
||||
ToolContext {
|
||||
app_state,
|
||||
silent: true,
|
||||
read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::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)]
|
||||
pub enum PermissionRule {
|
||||
/// 不可覆盖的拒绝
|
||||
Deny { tool_name: String, reason: String },
|
||||
/// 允许
|
||||
Allow { tool_name: String },
|
||||
/// 需要用户确认
|
||||
Ask { tool_name: String, message: String },
|
||||
}
|
||||
|
||||
/// 智能体工具 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() 使用。
|
||||
pub struct ToolRegistry {
|
||||
tools: std::collections::HashMap<String, Box<dyn AgentTool>>,
|
||||
ordered_names: Vec<String>,
|
||||
}
|
||||
|
||||
// ── 工具注册辅助函数(消除重复代码) ──
|
||||
|
||||
/// 注册所有基础研究工具(文件、文献、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(DelegateResearchTool::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(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建默认工具注册表(包含全部科研工具,不含后台工具)
|
||||
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(),
|
||||
};
|
||||
add_base_tools(&mut registry, skill_registry);
|
||||
if let Some(q) = queue {
|
||||
add_background_tools(&mut registry, q);
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
/// 动态添加工具(用于需要共享状态的工具,如 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);
|
||||
}
|
||||
|
||||
/// 替换已存在的工具(保持名称在 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);
|
||||
}
|
||||
|
||||
/// 根据名称查找工具 (O(1))
|
||||
pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
|
||||
self.tools.get(name).map(|t| t.as_ref())
|
||||
}
|
||||
|
||||
/// 生成所有工具的 ToolDefinition 列表(用于发送给 LLM)。
|
||||
/// 按名称字母序排序以保证跨调用的稳定性,提升 prompt cache 命中率。
|
||||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||||
let mut defs: Vec<_> = self
|
||||
.tools
|
||||
.values()
|
||||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||||
.collect();
|
||||
defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
|
||||
defs
|
||||
}
|
||||
}
|
||||
|
||||
/// 截断文本到指定最大字符数
|
||||
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 == "delegate_research"));
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
// src/agent/tools/persist.rs
|
||||
//
|
||||
// 工具结果持久化到磁盘。
|
||||
// 参考 Claude Code toolResultStorage.ts 设计。
|
||||
//
|
||||
// 当工具输出超过配置的字符限制时,将完整内容写入磁盘文件,
|
||||
// 返回一个 <persisted-output> 占位符给模型,模型可通过 read_file 工具读取完整内容。
|
||||
// 使用独占创建(create_new)保证幂等——同一 tool_call_id 不会被重复写入。
|
||||
//
|
||||
// 目录结构:
|
||||
// {library_dir}/tool-results/{tool_call_id}.txt
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use tracing::info;
|
||||
|
||||
/// 将工具输出持久化到磁盘(如果超过大小限制)。
|
||||
///
|
||||
/// 返回 (最终内容, 持久化文件路径)。
|
||||
/// 如果内容未超过限制,直接返回原内容且持久化路径为 None。
|
||||
pub fn maybe_persist_tool_result(
|
||||
content: &str,
|
||||
tool_call_id: &str,
|
||||
max_chars: usize,
|
||||
tool_results_dir: &Path,
|
||||
) -> (String, Option<PathBuf>) {
|
||||
if content.len() <= max_chars {
|
||||
return (content.to_string(), None);
|
||||
}
|
||||
|
||||
// 创建目录(幂等)
|
||||
if let Err(_e) = std::fs::create_dir_all(tool_results_dir) {
|
||||
// 无法创建目录则直接截断(不持久化)
|
||||
let truncated: String = content.chars().take(max_chars).collect();
|
||||
return (
|
||||
format!(
|
||||
"{}...\n[输出已截断,原始长度: {} 字符]",
|
||||
truncated,
|
||||
content.len()
|
||||
),
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
let file_path = tool_results_dir.join(format!("{}.txt", tool_call_id));
|
||||
|
||||
// 独占创建——只在文件不存在时写入,保证幂等性
|
||||
let written = match std::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.open(&file_path)
|
||||
{
|
||||
Ok(_) => {
|
||||
// 文件创建成功,写入内容
|
||||
match std::fs::write(&file_path, content) {
|
||||
Ok(_) => {
|
||||
info!(
|
||||
"[Persist] 工具结果已持久化: {} ({} 字符)",
|
||||
file_path.display(),
|
||||
content.len()
|
||||
);
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("[Persist] 写入失败: {}", e);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
|
||||
// 文件已存在,无需重复写入
|
||||
info!("[Persist] 工具结果已存在,跳过: {}", file_path.display());
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("[Persist] 创建文件失败: {}", e);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if !written {
|
||||
// 持久化失败,回退到截断
|
||||
let truncated: String = content.chars().take(max_chars).collect();
|
||||
return (
|
||||
format!(
|
||||
"{}...\n[输出已截断,原始长度: {} 字符]",
|
||||
truncated,
|
||||
content.len()
|
||||
),
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
// 预览(在第一个换行处截断,避免 mid-line cut)
|
||||
let preview_limit = 500.min(max_chars);
|
||||
let preview: String =
|
||||
if let Some(newline_pos) = content[..preview_limit.min(content.len())].rfind('\n') {
|
||||
content[..newline_pos].to_string()
|
||||
} else {
|
||||
content.chars().take(preview_limit).collect()
|
||||
};
|
||||
|
||||
let stub = format!(
|
||||
"<persisted-output>\n\
|
||||
path: {}\n\
|
||||
size: {} chars\n\
|
||||
preview: |\n {}\n\n\
|
||||
完整输出已持久化到磁盘。使用 read_file 工具以 path 参数读取完整内容。\n\
|
||||
</persisted-output>",
|
||||
file_path.display(),
|
||||
content.len(),
|
||||
preview.replace('\n', "\n "),
|
||||
);
|
||||
|
||||
(stub, Some(file_path))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn temp_dir() -> (PathBuf, impl Drop) {
|
||||
let dir = std::env::temp_dir().join(format!("astro_test_{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let dir_clone = dir.clone();
|
||||
struct Cleanup(PathBuf);
|
||||
impl Drop for Cleanup {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.0);
|
||||
}
|
||||
}
|
||||
(dir, Cleanup(dir_clone))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_small_content_not_persisted() {
|
||||
let (dir, _cleanup) = temp_dir();
|
||||
let (content, path) = maybe_persist_tool_result("small result", "call_1", 4000, &dir);
|
||||
assert_eq!(content, "small result");
|
||||
assert!(path.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_content_persisted() {
|
||||
let (dir, _cleanup) = temp_dir();
|
||||
let large = "x".repeat(5000);
|
||||
let (content, path) = maybe_persist_tool_result(&large, "call_2", 100, &dir);
|
||||
assert!(content.contains("<persisted-output>"));
|
||||
assert!(content.contains("call_2.txt"));
|
||||
assert!(path.is_some());
|
||||
let file_path = path.unwrap();
|
||||
assert!(file_path.exists());
|
||||
let written = std::fs::read_to_string(&file_path).unwrap();
|
||||
assert_eq!(written, large);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_idempotent_write() {
|
||||
let (dir, _cleanup) = temp_dir();
|
||||
let large1 = "a".repeat(5000);
|
||||
let large2 = "b".repeat(5000);
|
||||
|
||||
let (_, path1) = maybe_persist_tool_result(&large1, "call_3", 100, &dir);
|
||||
let (content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir);
|
||||
|
||||
assert!(path1.is_some());
|
||||
assert!(path2.is_some());
|
||||
let written = std::fs::read_to_string(path1.unwrap()).unwrap();
|
||||
assert_eq!(written, large1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_preview_at_newline_boundary() {
|
||||
let (dir, _cleanup) = temp_dir();
|
||||
let large = format!("Short line\n{}", "x".repeat(5000));
|
||||
let (content, _) = maybe_persist_tool_result(&large, "call_4", 100, &dir);
|
||||
let preview_start = content.find("preview:").unwrap();
|
||||
let preview_section = &content[preview_start..];
|
||||
assert!(preview_section.contains("Short line"));
|
||||
assert!(!preview_section.contains("xxx"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
// src/agent/tools/skill.rs — 技能加载工具(Layer 2 按需加载)
|
||||
//
|
||||
// 参考 Claude Code src/tools/SkillTool/SkillTool.ts 设计:
|
||||
// - description() 动态生成,列出所有可用 skill 及描述
|
||||
// - 支持 inline 模式(直接返回 skill 内容)和 fork 模式(子代理执行)
|
||||
// - 支持 allowed-tools 白名单返回
|
||||
// - 支持变量替换(${SKILL_DIR}, ${SESSION_ID})
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
use crate::agent::skills::{substitute_variables, SkillRegistry};
|
||||
|
||||
/// 技能加载工具
|
||||
///
|
||||
/// 持有 SkillRegistry 引用以实现:
|
||||
/// - 缓存读取(避免重复磁盘 I/O)
|
||||
/// - 动态 description 生成
|
||||
/// - 使用统计
|
||||
/// - fork 模式下的子代理执行
|
||||
pub struct LoadSkillTool {
|
||||
registry: Arc<RwLock<SkillRegistry>>,
|
||||
}
|
||||
|
||||
impl LoadSkillTool {
|
||||
pub fn new(registry: Arc<RwLock<SkillRegistry>>) -> Self {
|
||||
LoadSkillTool { registry }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for LoadSkillTool {
|
||||
fn name(&self) -> &str {
|
||||
"load_skill"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
// 注意:description() 返回 &str,但我们需要动态内容。
|
||||
// 实际使用 ToolRegistry 时,此方法的返回值作为 base description,
|
||||
// 详细的 skill 列表通过 system-reminder 注入。
|
||||
// 如果 trait 允许,应改为返回 String。当前保持与 trait 兼容。
|
||||
"加载指定领域技能的完整内容。可用技能列表已在系统提示词中列出。当需要某个技能的详细指引时调用此工具。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"skill_name": {
|
||||
"type": "string",
|
||||
"description": "要加载的技能名称。可用技能列表见系统提示词。"
|
||||
}
|
||||
},
|
||||
"required": ["skill_name"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯读文件,并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// fork 模式下的 skill 不应被中断
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Cancel
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let skill_name = match args.get("skill_name").and_then(|s| s.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'skill_name'"),
|
||||
};
|
||||
|
||||
info!("[LoadSkill] 加载 skill: {}", skill_name);
|
||||
|
||||
// 从缓存注册表读取(单语句:guard 自动 drop,不跨越 await)
|
||||
let skill = match self.registry.read() {
|
||||
Ok(reg) => reg.get_skill(&skill_name).cloned(),
|
||||
Err(e) => {
|
||||
warn!("[LoadSkill] RwLock poisoned: {:?}", e);
|
||||
return ToolOutput::error("技能注册表不可用(内部锁异常),请稍后重试");
|
||||
}
|
||||
};
|
||||
let skill = match skill {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
return ToolOutput::error(format!(
|
||||
"技能 '{}' 未找到。可用技能见系统提示词中的列表。",
|
||||
skill_name
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 记录使用统计(单语句:write guard 自动 drop)
|
||||
if let Ok(mut reg) = self.registry.write() {
|
||||
reg.record_usage(&skill_name);
|
||||
} else {
|
||||
warn!("[LoadSkill] 无法记录使用统计(RwLock poisoned)");
|
||||
}
|
||||
|
||||
// 变量替换
|
||||
let session_id = if ctx
|
||||
.app_state
|
||||
.config
|
||||
.database_url
|
||||
.contains("session") { "current" } else { "" };
|
||||
|
||||
let body = substitute_variables(
|
||||
&skill.body,
|
||||
&skill.skill_dir,
|
||||
if session_id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(session_id)
|
||||
},
|
||||
);
|
||||
|
||||
// 构建 skill 目录的绝对路径(用于 "Base directory" 前缀)
|
||||
let skill_dir_path = skill
|
||||
.skill_dir
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| skill.skill_dir.clone());
|
||||
let base_dir_note = format!(
|
||||
"Base directory for this skill: {}\n\n",
|
||||
skill_dir_path.display()
|
||||
);
|
||||
|
||||
// 检查 context 模式
|
||||
let is_fork = skill.meta.context.as_deref() == Some("fork");
|
||||
|
||||
if is_fork {
|
||||
// ── Fork 模式:使用子代理执行 skill ──
|
||||
info!(
|
||||
"[LoadSkill] Skill '{}' 标记为 fork 模式,启动子代理执行",
|
||||
skill_name
|
||||
);
|
||||
|
||||
let max_steps = args
|
||||
.get("max_steps")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(5)
|
||||
.min(10) as usize;
|
||||
|
||||
let runner = crate::agent::subagent::SubAgentRunner::new(ctx.app_state.clone());
|
||||
let system_prompt = format!(
|
||||
"你是一位专业的天体物理学研究助手。\n\n\
|
||||
你正在按照以下技能指引执行任务。\n\
|
||||
技能目录: {}\n\n{}{}\n\n\
|
||||
请严格遵循上述指引完成任务,使用可用工具收集和分析信息。\
|
||||
你可以使用 Read/Grep/Bash 工具访问技能目录中的文件。\
|
||||
完成后给出最终结果。",
|
||||
skill_dir_path.display(),
|
||||
base_dir_note,
|
||||
body
|
||||
);
|
||||
|
||||
// 子代理任务描述取自 skill description 或 body 前 200 字符
|
||||
let research_prompt = format!(
|
||||
"按照 '{}' 技能的指引完成任务:{}",
|
||||
skill.meta.name, skill.meta.description
|
||||
);
|
||||
|
||||
let result = runner
|
||||
.run(&system_prompt, &research_prompt, max_steps)
|
||||
.await;
|
||||
|
||||
if result.is_error {
|
||||
ToolOutput::error(format!(
|
||||
"技能 '{}' 子代理执行失败: {}",
|
||||
skill_name, result.content
|
||||
))
|
||||
} else {
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"[子代理执行结果 - 技能: {} ({})]\n\n{}",
|
||||
skill.meta.name, skill.meta.description, result.content
|
||||
),
|
||||
json!({
|
||||
"skill_name": skill.meta.name,
|
||||
"description": skill.meta.description,
|
||||
"context": "fork",
|
||||
"execution_mode": "subagent",
|
||||
"body_length": skill.body.len(),
|
||||
"allowed_tools": skill.meta.allowed_tools,
|
||||
}),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// ── Inline 模式:直接返回 skill 内容 ──
|
||||
let allowed_tools_note = if skill.meta.allowed_tools.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(
|
||||
"\n\n> **工具白名单**: {}(此技能建议仅使用这些工具)",
|
||||
skill.meta.allowed_tools.join(", ")
|
||||
)
|
||||
};
|
||||
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"# 技能:{} ({})\n\n{}{}{}",
|
||||
skill.meta.name,
|
||||
skill.meta.description,
|
||||
base_dir_note,
|
||||
body,
|
||||
allowed_tools_note
|
||||
),
|
||||
json!({
|
||||
"skill_name": skill.meta.name,
|
||||
"description": skill.meta.description,
|
||||
"context": skill.meta.context,
|
||||
"execution_mode": "inline",
|
||||
"body_length": skill.body.len(),
|
||||
"allowed_tools": skill.meta.allowed_tools,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
// src/agent/tools/subagent.rs — 子代理委托工具 (delegate_research)
|
||||
//
|
||||
// 参考 Claude Code s04 Subagents 设计。
|
||||
// LLM 通过此工具将子任务委托给上下文隔离的子代理执行。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
use tracing::info;
|
||||
|
||||
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
use crate::agent::hooks::HookRegistry;
|
||||
use crate::agent::runtime::permission::PermissionChecker;
|
||||
use crate::agent::runtime::AgentStreamEvent;
|
||||
use crate::agent::subagent::SubAgentRunner;
|
||||
|
||||
/// 子代理委托工具
|
||||
pub struct DelegateResearchTool {
|
||||
hook_registry: Option<Arc<HookRegistry>>,
|
||||
permission_checker: Arc<PermissionChecker>,
|
||||
progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
|
||||
}
|
||||
|
||||
impl Default for DelegateResearchTool {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl DelegateResearchTool {
|
||||
/// 创建不带 hooks 的工具实例(向后兼容)
|
||||
pub fn new() -> Self {
|
||||
DelegateResearchTool {
|
||||
hook_registry: None,
|
||||
permission_checker: Arc::new(PermissionChecker::new()),
|
||||
progress_tx: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带完整 hooks/permissions/progress 的工具实例
|
||||
pub fn new_with_hooks(
|
||||
hook_registry: Option<Arc<HookRegistry>>,
|
||||
permission_checker: Arc<PermissionChecker>,
|
||||
progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
|
||||
) -> Self {
|
||||
DelegateResearchTool {
|
||||
hook_registry,
|
||||
permission_checker,
|
||||
progress_tx,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for DelegateResearchTool {
|
||||
fn name(&self) -> &str {
|
||||
"delegate_research"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将子研究任务委托给独立的子代理执行。子代理拥有完整工具访问权限(文献搜索、下载、RAG检索等),\
|
||||
但只有最终文本摘要会返回给父代理,中间工具调用不会污染父上下文。\
|
||||
适用于:文献综述、多步数据收集、独立子问题研究等可以独立完成的子任务。\
|
||||
重要:delegate_research 返回后,你仍应基于其结果继续分析和回答用户问题。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"research_prompt": {
|
||||
"type": "string",
|
||||
"description": "要委托给子代理执行的完整研究任务描述。应包含具体的搜索目标、需要收集的信息、期望的输出格式。"
|
||||
},
|
||||
"max_steps": {
|
||||
"type": "integer",
|
||||
"description": "子代理最大推理步数,默认5,最大10",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["research_prompt"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 子代理可能执行写操作,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let research_prompt = match args.get("research_prompt").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'research_prompt'"),
|
||||
};
|
||||
|
||||
let max_steps = args
|
||||
.get("max_steps")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(5)
|
||||
.min(10) as usize;
|
||||
|
||||
info!(
|
||||
"[DelegateResearch] 启动子代理: prompt_len={}, max_steps={}",
|
||||
research_prompt.len(),
|
||||
max_steps
|
||||
);
|
||||
|
||||
let system_prompt = "你是一位专业的天体物理学研究助手,在一个独立的子任务上下文中工作。\
|
||||
你可以使用文献搜索、下载、RAG检索等工具。\
|
||||
请高效完成任务,然后直接给出最终答案。不要进行不必要的重复操作。\
|
||||
用中文回答,引用具体文献来源。";
|
||||
|
||||
let runner = SubAgentRunner::new_with_hooks(
|
||||
ctx.app_state.clone(),
|
||||
self.hook_registry.clone(),
|
||||
self.permission_checker.clone(),
|
||||
self.progress_tx.clone(),
|
||||
);
|
||||
let result = runner.run(system_prompt, &research_prompt, max_steps).await;
|
||||
|
||||
if result.is_error {
|
||||
ToolOutput::error(format!("子代理执行失败: {}", result.content))
|
||||
} else {
|
||||
// 包装子代理结果,标注来源
|
||||
ToolOutput::success(
|
||||
format!("[子代理研究结果]\n\n{}", result.content),
|
||||
result.metadata,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
// src/agent/tools/team.rs — 团队协作工具
|
||||
//
|
||||
// 4 个团队工具:
|
||||
// spawn_teammate — 启动一个队友 agent
|
||||
// send_teammate_message — 发送消息给指定队友
|
||||
// team_broadcast — 广播消息给所有队友
|
||||
// check_team_inbox — 检查收件箱
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
use crate::agent::team::inbox::TeamMessageType;
|
||||
use crate::agent::team::manager::TeamManager;
|
||||
|
||||
/// 启动队友工具
|
||||
pub struct SpawnTeammateTool {
|
||||
team_manager: Arc<Mutex<Option<TeamManager>>>,
|
||||
}
|
||||
|
||||
impl SpawnTeammateTool {
|
||||
pub fn new(team_manager: Arc<Mutex<Option<TeamManager>>>) -> Self {
|
||||
SpawnTeammateTool { team_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SpawnTeammateTool {
|
||||
fn name(&self) -> &str {
|
||||
"spawn_teammate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"启动一个队友 agent。队友拥有独立的上下文,可以并行执行文献搜索、论文下载等任务。\
|
||||
通过 send_teammate_message 向队友发送任务,通过 check_team_inbox 检查结果。\
|
||||
适用于:需要多线并行的文献调研任务。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "队友名称(如 searcher, reader)"
|
||||
},
|
||||
"role": {
|
||||
"type": "string",
|
||||
"description": "队友角色描述(如 ADS文献搜索专家, 论文全文阅读专家)"
|
||||
}
|
||||
},
|
||||
"required": ["name", "role"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 启动队友有副作用,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let name = match args.get("name").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 'name' 参数"),
|
||||
};
|
||||
let role = match args.get("role").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 'role' 参数"),
|
||||
};
|
||||
|
||||
// NOTE: team_manager 锁在此 await 期间保持持有。
|
||||
// spawn 操作通常很快(只是创建子代理会话),
|
||||
// 如果需要降低锁持有时间,可以将 TeamManager 改为内部使用 Arc。
|
||||
let tm_lock = self.team_manager.lock().await;
|
||||
match tm_lock.as_ref() {
|
||||
Some(tm) => {
|
||||
let handle = tm.spawn(&name, &role).await;
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"✅ 队友已启动: {} ({})\n使用 send_teammate_message 发送任务。",
|
||||
handle.name, handle.role
|
||||
),
|
||||
json!({ "name": handle.name, "role": handle.role }),
|
||||
)
|
||||
}
|
||||
None => ToolOutput::error("团队管理器未初始化"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 发送消息工具
|
||||
pub struct SendTeammateMessageTool {
|
||||
team_manager: Arc<Mutex<Option<TeamManager>>>,
|
||||
}
|
||||
|
||||
impl SendTeammateMessageTool {
|
||||
pub fn new(team_manager: Arc<Mutex<Option<TeamManager>>>) -> Self {
|
||||
SendTeammateMessageTool { team_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SendTeammateMessageTool {
|
||||
fn name(&self) -> &str {
|
||||
"send_teammate_message"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"向指定队友发送消息(任务分配、问题等)。消息会投递到队友的收件箱,\
|
||||
队友在处理循环中自动读取。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"to": {
|
||||
"type": "string",
|
||||
"description": "收件队友名称"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "消息内容(任务描述、问题等)"
|
||||
}
|
||||
},
|
||||
"required": ["to", "content"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 消息投递有副作用,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let to = match args.get("to").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 'to' 参数"),
|
||||
};
|
||||
let content = match args.get("content").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 'content' 参数"),
|
||||
};
|
||||
|
||||
let tm_lock = self.team_manager.lock().await;
|
||||
match tm_lock.as_ref() {
|
||||
Some(tm) => {
|
||||
tm.send_message("lead", &to, &content, TeamMessageType::Task);
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"📤 消息已发送给 {}: {}",
|
||||
to,
|
||||
&content.chars().take(100).collect::<String>()
|
||||
),
|
||||
json!({ "to": to }),
|
||||
)
|
||||
}
|
||||
None => ToolOutput::error("团队管理器未初始化"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 广播消息工具
|
||||
pub struct TeamBroadcastTool {
|
||||
team_manager: Arc<Mutex<Option<TeamManager>>>,
|
||||
}
|
||||
|
||||
impl TeamBroadcastTool {
|
||||
pub fn new(team_manager: Arc<Mutex<Option<TeamManager>>>) -> Self {
|
||||
TeamBroadcastTool { team_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for TeamBroadcastTool {
|
||||
fn name(&self) -> &str {
|
||||
"team_broadcast"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"向所有队友广播消息。适用于状态同步、全局指令等。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "广播内容"
|
||||
}
|
||||
},
|
||||
"required": ["content"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 广播有副作用,中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let content = match args.get("content").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 'content' 参数"),
|
||||
};
|
||||
|
||||
let tm_lock = self.team_manager.lock().await;
|
||||
match tm_lock.as_ref() {
|
||||
Some(tm) => {
|
||||
tm.broadcast("lead", &content);
|
||||
ToolOutput::success("📢 已广播消息给所有队友。", json!({}))
|
||||
}
|
||||
None => ToolOutput::error("团队管理器未初始化"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查收件箱工具
|
||||
pub struct CheckTeamInboxTool {
|
||||
team_manager: Arc<Mutex<Option<TeamManager>>>,
|
||||
}
|
||||
|
||||
impl CheckTeamInboxTool {
|
||||
pub fn new(team_manager: Arc<Mutex<Option<TeamManager>>>) -> Self {
|
||||
CheckTeamInboxTool { team_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for CheckTeamInboxTool {
|
||||
fn name(&self) -> &str {
|
||||
"check_team_inbox"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"检查收件箱,获取队友发来的消息。读取后消息会被清空。不指定 agent 时检查 lead 的收件箱。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"agent_name": {
|
||||
"type": "string",
|
||||
"description": "可选:要检查的 agent 名称,默认为 lead"
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
/// 收件箱检查是破坏性读取(消息会被清空),中断时应阻塞以完成
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let agent = args
|
||||
.get("agent_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("lead");
|
||||
|
||||
let tm_lock = self.team_manager.lock().await;
|
||||
match tm_lock.as_ref() {
|
||||
Some(tm) => {
|
||||
let msgs = tm.check_inbox(agent);
|
||||
if msgs.is_empty() {
|
||||
ToolOutput::success(
|
||||
format!("📭 {} 的收件箱为空。", agent),
|
||||
json!({ "messages": [] }),
|
||||
)
|
||||
} else {
|
||||
let mut result = format!("📬 {} 的收件箱 ({} 条消息):\n\n", agent, msgs.len());
|
||||
for msg in &msgs {
|
||||
result.push_str(&format!(
|
||||
" [{}] 来自 {}: {}\n",
|
||||
msg.msg_type.as_str(),
|
||||
msg.from,
|
||||
&msg.content.chars().take(200).collect::<String>(),
|
||||
));
|
||||
}
|
||||
ToolOutput::success(result, json!({ "count": msgs.len() }))
|
||||
}
|
||||
}
|
||||
None => ToolOutput::error("团队管理器未初始化"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
// src/agent/tools/todo.rs — 任务规划工具 (TodoWrite)
|
||||
//
|
||||
// P1 改进:任务状态持久化到 SQLite agent_tasks 表,
|
||||
// 支持 blockedBy 依赖关系和跨 turn 状态恢复。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
/// 任务规划工具:让 LLM 在开始复杂研究前先制定计划,执行中更新进度。
|
||||
/// 任务状态持久化到 SQLite,支持 DAG 依赖。
|
||||
pub struct TodoWriteTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for TodoWriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"todo_write"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"任务规划工具。在开始复杂研究前列出待办事项,执行中标记进度(每项状态:pending/in_progress/completed)。\
|
||||
一次只能有一个 in_progress 任务。支持任务依赖(blockedBy:依赖的其他任务ID列表)。\
|
||||
适用范围:任何需要多步工具调用的研究任务。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"todos": {
|
||||
"type": "array",
|
||||
"description": "任务列表,每项包含 id(唯一标识)、content(任务描述)、status(pending/in_progress/completed)、blockedBy(可选,依赖的其他任务ID列表)",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": { "type": "string", "description": "任务唯一标识" },
|
||||
"content": { "type": "string", "description": "任务描述" },
|
||||
"status": {
|
||||
"type": "string",
|
||||
"enum": ["pending", "in_progress", "completed"],
|
||||
"description": "任务状态"
|
||||
},
|
||||
"blockedBy": {
|
||||
"type": "array",
|
||||
"items": { "type": "string" },
|
||||
"description": "该任务依赖的其他任务ID列表(这些任务必须先完成)"
|
||||
}
|
||||
},
|
||||
"required": ["id", "content", "status"]
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": ["todos"]
|
||||
})
|
||||
}
|
||||
|
||||
/// 纯格式化输出,无副作用(持久化由 runtime 层处理),并发安全
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let todos = match args.get("todos").and_then(|t| t.as_array()) {
|
||||
Some(t) => t,
|
||||
None => return ToolOutput::error("缺少必需参数 'todos'"),
|
||||
};
|
||||
|
||||
// 从 ToolContext 中无法直接获取 session_id
|
||||
// TodoWrite 工具生成格式化的输出,实际持久化在 runtime 层完成
|
||||
// 这里只做验证和格式化
|
||||
|
||||
let mut formatted = String::from("📋 当前任务计划:\n\n");
|
||||
let mut in_progress_count = 0;
|
||||
|
||||
for todo in todos {
|
||||
let id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let status = todo
|
||||
.get("status")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("pending");
|
||||
|
||||
let icon = match status {
|
||||
"in_progress" => {
|
||||
in_progress_count += 1;
|
||||
"🔄"
|
||||
}
|
||||
"completed" => "✅",
|
||||
_ => "⏳",
|
||||
};
|
||||
|
||||
// 显示依赖关系
|
||||
let blocked_by: Vec<String> = todo
|
||||
.get("blockedBy")
|
||||
.and_then(|b| b.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut line = format!("{} [{}] {}", icon, id, content);
|
||||
if !blocked_by.is_empty() {
|
||||
line.push_str(&format!(" (依赖: {})", blocked_by.join(", ")));
|
||||
}
|
||||
formatted.push_str(&line);
|
||||
formatted.push('\n');
|
||||
}
|
||||
|
||||
// 验证约束
|
||||
if in_progress_count > 1 {
|
||||
formatted.push_str(
|
||||
"\n⚠️ 提醒:你当前有多个 in_progress 任务。请先完成当前任务再开始下一个。",
|
||||
);
|
||||
} else if in_progress_count == 0
|
||||
&& todos
|
||||
.iter()
|
||||
.any(|t| t.get("status").and_then(|v| v.as_str()) == Some("pending"))
|
||||
{
|
||||
formatted
|
||||
.push_str("\n💡 提示:还有待处理任务,请选择一个设为 in_progress 并开始执行。");
|
||||
}
|
||||
|
||||
ToolOutput::success(formatted, json!({ "task_count": todos.len() }))
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 TodoWrite 的任务列表持久化到 agent_tasks 表。
|
||||
///
|
||||
/// 使用 INSERT OR REPLACE 实现 upsert(基于 session_id + task_id 唯一约束)。
|
||||
/// `owner` 参数指定任务的归属 agent(默认 "lead")。
|
||||
pub async fn persist_tasks(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
todos: &[serde_json::Value],
|
||||
owner: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
for todo in todos {
|
||||
let task_id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("unknown");
|
||||
let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let status = todo
|
||||
.get("status")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("pending");
|
||||
|
||||
// 收集 blockedBy 数组
|
||||
let blocked_by: Vec<String> = todo
|
||||
.get("blockedBy")
|
||||
.and_then(|b| b.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let blocked_by_json =
|
||||
serde_json::to_string(&blocked_by).unwrap_or_else(|_| "[]".to_string());
|
||||
|
||||
// 简单 DAG 验证:不能依赖自身
|
||||
if blocked_by.contains(&task_id.to_string()) {
|
||||
warn!(
|
||||
"[TodoWrite] 任务 {} 依赖自身,已跳过 blockedBy 中的自引用",
|
||||
task_id
|
||||
);
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_tasks (session_id, task_id, content, status, blocked_by, owner) \
|
||||
VALUES (?, ?, ?, ?, ?, ?) \
|
||||
ON CONFLICT(session_id, task_id) DO UPDATE SET \
|
||||
content=excluded.content, \
|
||||
status=excluded.status, \
|
||||
blocked_by=excluded.blocked_by, \
|
||||
owner=excluded.owner, \
|
||||
updated_at=CURRENT_TIMESTAMP",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(task_id)
|
||||
.bind(content)
|
||||
.bind(status)
|
||||
.bind(&blocked_by_json)
|
||||
.bind(owner)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
info!(
|
||||
"[TodoWrite] 已持久化 {} 个任务到会话 {}",
|
||||
todos.len(),
|
||||
session_id
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_todo_output_format() {
|
||||
let tool = TodoWriteTool;
|
||||
let args = json!({
|
||||
"todos": [
|
||||
{"id": "1", "content": "搜索文献", "status": "completed"},
|
||||
{"id": "2", "content": "阅读论文", "status": "in_progress", "blockedBy": ["1"]},
|
||||
{"id": "3", "content": "撰写综述", "status": "pending", "blockedBy": ["2"]}
|
||||
]
|
||||
});
|
||||
|
||||
// 验证参数 schema
|
||||
let params = tool.parameters();
|
||||
assert!(params["required"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.contains(&json!("todos")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_todo_multiple_in_progress_warning() {
|
||||
// 直接测试格式化逻辑
|
||||
let todos = json!([
|
||||
{"id": "1", "content": "任务A", "status": "in_progress"},
|
||||
{"id": "2", "content": "任务B", "status": "in_progress"}
|
||||
]);
|
||||
|
||||
let mut in_progress_count = 0;
|
||||
for todo in todos.as_array().unwrap() {
|
||||
if todo.get("status").and_then(|v| v.as_str()) == Some("in_progress") {
|
||||
in_progress_count += 1;
|
||||
}
|
||||
}
|
||||
assert_eq!(in_progress_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_todo_with_blocked_by() {
|
||||
let todos = json!([
|
||||
{"id": "1", "content": "文献检索", "status": "completed"},
|
||||
{"id": "2", "content": "文献分析", "status": "pending", "blockedBy": ["1"]}
|
||||
]);
|
||||
|
||||
let items = todos.as_array().unwrap();
|
||||
let blocked: Vec<String> = items[1]
|
||||
.get("blockedBy")
|
||||
.and_then(|b| b.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
assert_eq!(blocked, vec!["1"]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
// src/agent/trajectory.rs
|
||||
//
|
||||
// Trajectory 数据导出 — 为 RLHF/微调准备结构化训练数据。
|
||||
// 参考 learn-claude-code "Collect trajectory data" 设计。
|
||||
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tracing::info;
|
||||
|
||||
use super::runtime::AgentMetrics;
|
||||
use super::terminal::TurnTerminal;
|
||||
|
||||
/// 单条 trajectory 记录(JSONL 格式)
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TrajectoryRecord {
|
||||
pub timestamp: String,
|
||||
pub session_id: String,
|
||||
pub model: String,
|
||||
pub system_prompt: String,
|
||||
pub messages: Vec<TrajectoryMessage>,
|
||||
pub final_answer: Option<String>,
|
||||
pub metrics: TrajectoryMetrics,
|
||||
pub terminal_reason: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TrajectoryMessage {
|
||||
pub role: String,
|
||||
pub content: Option<String>,
|
||||
pub tool_calls: Option<serde_json::Value>,
|
||||
pub tool_call_id: Option<String>,
|
||||
pub thought: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TrajectoryMetrics {
|
||||
pub total_steps: usize,
|
||||
pub tool_calls: std::collections::HashMap<String, usize>,
|
||||
pub compression_count: usize,
|
||||
pub duplicate_detections: usize,
|
||||
}
|
||||
|
||||
/// Trajectory 导出器
|
||||
pub struct TrajectoryExporter;
|
||||
|
||||
impl TrajectoryExporter {
|
||||
/// 导出指定会话的完整 trajectory 到 JSONL 文件。
|
||||
pub async fn export(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
library_dir: &Path,
|
||||
model: &str,
|
||||
system_prompt: &str,
|
||||
metrics: &AgentMetrics,
|
||||
terminal: Option<&TurnTerminal>,
|
||||
) -> anyhow::Result<PathBuf> {
|
||||
// 从 DB 加载消息
|
||||
let messages = Self::load_messages(db, session_id).await?;
|
||||
|
||||
// 提取最终答案
|
||||
let final_answer = messages
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|m| m.role == "assistant" && m.tool_calls.is_none())
|
||||
.and_then(|m| m.content.clone());
|
||||
|
||||
let terminal_reason = terminal
|
||||
.map(|t| t.description().to_string())
|
||||
.unwrap_or_else(|| "completed".to_string());
|
||||
|
||||
let record = TrajectoryRecord {
|
||||
timestamp: chrono::Utc::now().to_rfc3339(),
|
||||
session_id: session_id.to_string(),
|
||||
model: model.to_string(),
|
||||
system_prompt: system_prompt.to_string(),
|
||||
messages,
|
||||
final_answer,
|
||||
metrics: TrajectoryMetrics {
|
||||
total_steps: metrics.total_steps,
|
||||
tool_calls: metrics.tool_calls.clone(),
|
||||
compression_count: metrics.compression_count,
|
||||
duplicate_detections: metrics.duplicate_detections,
|
||||
},
|
||||
terminal_reason,
|
||||
};
|
||||
|
||||
let dir = library_dir.join("trajectories");
|
||||
fs::create_dir_all(&dir)?;
|
||||
|
||||
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
|
||||
let path = dir.join(format!("{}_{}.jsonl", session_id, timestamp));
|
||||
|
||||
let json = serde_json::to_string(&record)?;
|
||||
fs::write(&path, format!("{}\n", json))?;
|
||||
|
||||
info!("[Trajectory] 已导出: {}", path.display());
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
async fn load_messages(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<TrajectoryMessage>> {
|
||||
let rows = sqlx::query_as::<
|
||||
_,
|
||||
(
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
),
|
||||
>(
|
||||
"SELECT role, content, thought, tool_calls, tool_call_id \
|
||||
FROM agent_messages WHERE session_id = ? \
|
||||
ORDER BY created_at ASC, step_index ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|(role, content, thought, tool_calls, tool_call_id)| {
|
||||
let tool_calls_json = tool_calls
|
||||
.as_deref()
|
||||
.and_then(|tc| serde_json::from_str(tc).ok());
|
||||
TrajectoryMessage {
|
||||
role,
|
||||
content,
|
||||
tool_calls: tool_calls_json,
|
||||
tool_call_id,
|
||||
thought,
|
||||
}
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// 列出所有已导出的 trajectory 文件
|
||||
pub fn list_trajectories(library_dir: &Path) -> Vec<String> {
|
||||
let dir = library_dir.join("trajectories");
|
||||
match fs::read_dir(&dir) {
|
||||
Ok(entries) => entries
|
||||
.flatten()
|
||||
.filter_map(|e| {
|
||||
let name = e.file_name().to_string_lossy().to_string();
|
||||
if name.ends_with(".jsonl") {
|
||||
Some(name)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
Err(_) => Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_list_empty_trajectories() {
|
||||
let dir = PathBuf::from("/tmp/nonexistent_trajectory_dir");
|
||||
let result = TrajectoryExporter::list_trajectories(&dir);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user