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:
fmq
2026-06-17 00:14:02 +08:00
parent b1fb884f21
commit 49784739fa
113 changed files with 20253 additions and 2869 deletions
+137
View File
@@ -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 循环结束");
}
}
+203
View File
@@ -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
}
+721
View File
@@ -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_compactLLM 摘要)
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);
}
}
+266
View File
@@ -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());
}
}
+974
View File
@@ -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
}
}
+132
View File
@@ -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), "");
}
}
+143
View File
@@ -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
);
}
}
+406
View File
@@ -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());
}
}
+236
View File
@@ -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("现有记忆清单"));
}
}
+83
View File
@@ -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("不应保存"));
}
}
+420
View File
@@ -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());
}
}
+351
View File
@@ -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.0LLM 选中即默认置信度)。
/// 仅影响排序顺序,不减少选中的数量。
/// 仅活跃记忆参与衰减排序;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 应在后");
}
}
+330
View File
@@ -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
View File
@@ -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;
-786
View File
@@ -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
}
}
}
+214
View File
@@ -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);
}
}
+85
View File
@@ -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")
))
}
+428
View File
@@ -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);
}
}
+370
View File
@@ -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);
}
});
}
+571
View File
@@ -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"));
}
}
+137
View File
@@ -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
+189
View File
@@ -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, &registry);
// 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, &registry);
// 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, &registry);
// 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);
}
}
+187
View File
@@ -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);
}
}
+144
View File
@@ -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)
}
+161
View File
@@ -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,
}
}
+315
View File
@@ -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;
}
}
+104
View File
@@ -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);
}
}
+332
View File
@@ -0,0 +1,332 @@
// src/agent/runtime/token_budget.rs
//
// Token 预算管理。
// 参考 Claude Code TokenBudget 设计。
// 软限制:接近上限时注入 nudging 消息提醒模型。
// 硬限制:达到上限时触发强制压缩或终止。
/// Token 预算管理器。
///
/// 参考 Claude Code TokenBudget 设计,增加:
/// - 多级渐进式 nudgenear_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);
}
}
+846
View File
@@ -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 — LoadSkillToolLLM 按需调用,注入完整 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>,
}
/// 完整的 SkillLayer 2LLM 调用 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("没有可用的技能"));
}
}
+446
View File
@@ -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 }))
}
}
}
+142
View File
@@ -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/ 目录中
}
}
+53
View File
@@ -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,
});
}
}
+110
View File
@@ -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)
}
+205
View File
@@ -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
}
}
+11
View File
@@ -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;
+270
View File
@@ -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
}
+127
View File
@@ -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)
}
}
}
}
-659
View File
@@ -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());
}
}
+248
View File
@@ -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
))
}
}
}
}
+15
View File
@@ -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;
+99
View File
@@ -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(&notes_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)),
}
}
}
+247
View File
@@ -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))
}
}
}
}
}
+104
View File
@@ -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)),
}
}
}
+189
View File
@@ -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)),
}
}
}
+89
View File
@@ -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)),
}
}
}
+206
View File
@@ -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 }))
}
}
}
}
+42
View File
@@ -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" }),
)
}
}
+226
View File
@@ -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)
}
+336
View File
@@ -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)),
}
}
}
+100
View File
@@ -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)),
}
}
}
+203
View File
@@ -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: &regex::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: &regex::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
}
+23
View File
@@ -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;
+165
View File
@@ -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)),
}
}
}
+48
View File
@@ -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))
}
}
+94
View File
@@ -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)),
}
}
}
+185
View File
@@ -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)),
}
}
}
+417
View File
@@ -0,0 +1,417 @@
// src/agent/tools/mod.rs
//
// 科研智能体工具集定义与实现。
// 每个工具遵循 AgentTool trait,向大模型声明 JSON Schema 参数定义,
// 并在 execute 中调用已有的服务层完成实际业务操作。
//
// 按功能域拆分为子模块:
// filesystem/ — 文件 I/Oread、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());
}
}
+181
View File
@@ -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"));
}
}
+222
View File
@@ -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,
}),
)
}
}
}
+132
View File
@@ -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,
)
}
}
}
+291
View File
@@ -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("团队管理器未初始化"),
}
}
}
+260
View File
@@ -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 blockedByID列表\
"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"todos": {
"type": "array",
"description": "任务列表,每项包含 id(唯一标识)、content(任务描述)、statuspending/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"]);
}
}
+171
View File
@@ -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());
}
}