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
+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);
}
}