feat: Agent 全栈升级——模块化重构、Hooks/Skills/Memory/SubAgent/Team 子系统、审计与任务持久化
架构重构: - Agent Runtime 由单文件拆为 runtime/ 目录 12 模块(熔断/流式执行/Token预算/文件缓存/权限等) - Agent Tools 由单文件拆为 tools/ 目录 20+ 模块(filesystem/astro/memory/skill/subagent/team 等) - 解析器体系重构(common.rs 836行变更),各解析器同步升级 - Download 服务重构(562行),反爬策略强化 - LLM 客户端重构(266行),流式调用优化 新子系统: - Hooks 生命周期系统(9种事件类型,PreToolUse/PostToolUse 支持输入输出拦截) - Skills 双层加载系统(system-reminder 轻量注入 + LoadSkillTool 按需加载,notify 文件监听热更新) - Memory 项目记忆管理(类型/提取/去重/衰减/保活/选择策略/护栏 7 模块) - SubAgent 上下文隔离子代理运行器(独立 ReAct 循环 + Hook 管道) - Team 多智能体团队协作(文件 inbox 通信、lead/teammate 协调) - TaskBoard DAG 任务依赖管理 - Trajectory 会话轨迹、Terminal 终止信号、Autonomous 自主模式、Background 异步通知 数据库: - agent_tasks 表(DAG 依赖模式,blocked_by JSON 数组) - agent_audit_log 表(工具调用审计:名称/状态/耗时/输出预览) - agent_identity 迁移(消息/审计/任务的 agent_name 归属,agent_team_members 团队注册表) API: - GET /chat/metrics 聚合指标端点 - GET /chat/sessions/:id/audit 会话审计查询 - GET /chat/questions + POST /chat/answer 人机交互问答 工程: - 新增依赖:serde_yaml、notify、glob、walkdir、lru - Skills 目录含 methodology/plotting/presentation 三个初始 SKILL.md - CLAUDE.md 完整项目架构文档
This commit is contained in:
@@ -0,0 +1,214 @@
|
||||
// src/agent/runtime/circuit_breaker.rs
|
||||
//
|
||||
// 熔断器 — 防止无限自动压缩循环。
|
||||
// 参考 Claude Code MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES 设计。
|
||||
//
|
||||
// 当自动压缩连续失败 MAX_CONSECUTIVE_FAILURES 次后,熔断器打开,
|
||||
// 停止后续压缩尝试,避免无限循环。
|
||||
|
||||
use std::time::Instant;
|
||||
use tracing::warn;
|
||||
|
||||
/// 最大连续失败次数,超出后熔断器打开
|
||||
const MAX_CONSECUTIVE_FAILURES: usize = 3;
|
||||
|
||||
/// 熔断器打开后,经过此时间自动进入 HalfOpen 状态尝试恢复
|
||||
const AUTO_RECOVERY_TIMEOUT_SECS: u64 = 300; // 5 分钟
|
||||
|
||||
/// 熔断器状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum CircuitState {
|
||||
/// 正常工作,允许压缩
|
||||
Closed,
|
||||
/// 熔断,拒绝后续压缩
|
||||
Open,
|
||||
/// 半开:允许一次试探性压缩以决定是否恢复
|
||||
HalfOpen,
|
||||
}
|
||||
|
||||
/// 压缩熔断器
|
||||
#[derive(Debug)]
|
||||
pub struct CompactionCircuitBreaker {
|
||||
/// 连续失败计数
|
||||
consecutive_failures: usize,
|
||||
/// 压缩总次数
|
||||
total_compactions: usize,
|
||||
/// 当前状态
|
||||
state: CircuitState,
|
||||
/// 熔断器打开的时间(用于自动恢复)
|
||||
opened_at: Option<Instant>,
|
||||
}
|
||||
|
||||
impl CompactionCircuitBreaker {
|
||||
/// 创建新的熔断器(初始状态 Closed)
|
||||
pub fn new() -> Self {
|
||||
CompactionCircuitBreaker {
|
||||
consecutive_failures: 0,
|
||||
total_compactions: 0,
|
||||
state: CircuitState::Closed,
|
||||
opened_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录一次成功的压缩(重置失败计数,关闭熔断器)
|
||||
pub fn record_success(&mut self) {
|
||||
self.consecutive_failures = 0;
|
||||
self.total_compactions += 1;
|
||||
self.state = CircuitState::Closed;
|
||||
self.opened_at = None;
|
||||
}
|
||||
|
||||
/// 记录一次失败的压缩(递增失败计数,可能触发熔断)
|
||||
pub fn record_failure(&mut self) {
|
||||
self.consecutive_failures += 1;
|
||||
self.total_compactions += 1;
|
||||
|
||||
if self.consecutive_failures >= MAX_CONSECUTIVE_FAILURES {
|
||||
let was_already_open = self.state == CircuitState::Open;
|
||||
self.state = CircuitState::Open;
|
||||
self.opened_at = Some(Instant::now());
|
||||
if !was_already_open {
|
||||
warn!(
|
||||
"[CircuitBreaker] 熔断器打开!连续 {} 次压缩失败,停止自动压缩。\
|
||||
{} 秒后将自动尝试恢复。",
|
||||
self.consecutive_failures, AUTO_RECOVERY_TIMEOUT_SECS
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 熔断器是否打开(应停止自动压缩)
|
||||
pub fn is_open(&self) -> bool {
|
||||
self.state == CircuitState::Open
|
||||
}
|
||||
|
||||
/// 是否可以尝试压缩。
|
||||
///
|
||||
/// 当熔断器打开超过 AUTO_RECOVERY_TIMEOUT_SECS 时,自动转为 HalfOpen 状态,
|
||||
/// 允许下一次压缩尝试以判断是否恢复。
|
||||
pub fn can_attempt(&mut self) -> bool {
|
||||
match self.state {
|
||||
CircuitState::Closed | CircuitState::HalfOpen => true,
|
||||
CircuitState::Open => {
|
||||
// 检查是否已超时,可自动进入 HalfOpen
|
||||
if let Some(opened) = self.opened_at {
|
||||
if opened.elapsed().as_secs() >= AUTO_RECOVERY_TIMEOUT_SECS {
|
||||
self.state = CircuitState::HalfOpen;
|
||||
warn!(
|
||||
"[CircuitBreaker] 熔断器超时,进入 HalfOpen 状态,\
|
||||
允许下一次压缩尝试"
|
||||
);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 重置熔断器到 Closed 状态
|
||||
pub fn reset(&mut self) {
|
||||
self.consecutive_failures = 0;
|
||||
self.state = CircuitState::Closed;
|
||||
self.opened_at = None;
|
||||
}
|
||||
|
||||
/// 获取连续失败次数
|
||||
pub fn consecutive_failures(&self) -> usize {
|
||||
self.consecutive_failures
|
||||
}
|
||||
|
||||
/// 获取压缩总次数
|
||||
pub fn total_compactions(&self) -> usize {
|
||||
self.total_compactions
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CompactionCircuitBreaker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_opens_after_max_failures() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
assert!(!breaker.is_open());
|
||||
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert!(!breaker.is_open()); // 2 failures, not yet open
|
||||
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open()); // 3 failures, now open
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_on_success() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert_eq!(breaker.consecutive_failures(), 2);
|
||||
|
||||
breaker.record_success();
|
||||
assert_eq!(breaker.consecutive_failures(), 0);
|
||||
assert!(!breaker.is_open());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_method() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open());
|
||||
|
||||
breaker.reset();
|
||||
assert!(!breaker.is_open());
|
||||
assert_eq!(breaker.consecutive_failures(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_can_attempt() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
assert!(breaker.can_attempt());
|
||||
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
// 刚打开,不应允许尝试
|
||||
assert!(!breaker.can_attempt());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_half_open_after_reset_or_success() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
// 触发熔断
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
assert!(breaker.is_open());
|
||||
|
||||
// success 直接重置到 Closed
|
||||
breaker.record_success();
|
||||
assert!(!breaker.is_open());
|
||||
assert!(breaker.can_attempt());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_record_failure_while_open_stays_open() {
|
||||
let mut breaker = CompactionCircuitBreaker::new();
|
||||
for _ in 0..3 {
|
||||
breaker.record_failure();
|
||||
}
|
||||
assert!(breaker.is_open());
|
||||
// 熔断器打开后再次失败,保持 Open
|
||||
breaker.record_failure();
|
||||
assert!(breaker.is_open());
|
||||
assert_eq!(breaker.consecutive_failures(), 4);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
// src/agent/runtime/context.rs
|
||||
//
|
||||
// 上下文构建:加载历史消息、注入系统提示词、添加用户消息、
|
||||
// 从数据库恢复持久化的任务状态。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::info;
|
||||
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
use super::session;
|
||||
|
||||
/// 构建初始 LLM 上下文:加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。
|
||||
pub async fn build_initial_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
system_prompt: &str,
|
||||
question: &str,
|
||||
_turn_index: i32,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let mut messages = session::load_history_for_llm(db, session_id).await?;
|
||||
|
||||
// 注入系统提示词(如果历史中没有)
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(system_prompt));
|
||||
}
|
||||
|
||||
// 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
|
||||
// 从数据库恢复持久化的任务状态
|
||||
if let Some(task_reminder) = restore_tasks_from_db(db, session_id).await {
|
||||
messages.push(ChatMessage::user(task_reminder));
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 从 agent_tasks 表恢复任务状态,返回格式化的提醒文本。
|
||||
///
|
||||
/// 如果表不存在或没有任务记录,返回 None。
|
||||
async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option<String> {
|
||||
let rows: Vec<(String, String, String, String, Option<String>)> = sqlx::query_as(
|
||||
"SELECT task_id, content, status, blocked_by, owner \
|
||||
FROM agent_tasks WHERE session_id = ? AND (owner = '' OR owner = 'lead') \
|
||||
ORDER BY created_at ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
if rows.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut lines: Vec<String> = Vec::new();
|
||||
for (task_id, content, status, blocked_by, owner) in &rows {
|
||||
let icon = match status.as_str() {
|
||||
"in_progress" => "🔄",
|
||||
"completed" => "✅",
|
||||
_ => "⏳",
|
||||
};
|
||||
|
||||
let blocked: Vec<String> = serde_json::from_str(blocked_by).unwrap_or_default();
|
||||
let mut line = format!("{} [{}] {}", icon, task_id, content);
|
||||
if !blocked.is_empty() {
|
||||
line.push_str(&format!(" (依赖: {})", blocked.join(", ")));
|
||||
}
|
||||
if let Some(o) = owner {
|
||||
if !o.is_empty() && o != "lead" {
|
||||
line.push_str(&format!(" (指派: {})", o));
|
||||
}
|
||||
}
|
||||
lines.push(line);
|
||||
}
|
||||
|
||||
info!("[Context] 从数据库恢复了 {} 个任务状态", rows.len());
|
||||
|
||||
Some(format!(
|
||||
"[当前任务状态]\n以下是上次会话中持久化的任务计划,请基于最新状态继续工作:\n\n{}\n\n\
|
||||
使用 todo_write 工具更新任务进度。",
|
||||
lines.join("\n")
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
// src/agent/runtime/error_recovery.rs
|
||||
//
|
||||
// 错误恢复阶梯。
|
||||
// 参考 Claude Code error recovery ladder 设计。
|
||||
//
|
||||
// 当 LLM 流返回可恢复的错误(如 prompt_too_long)时,
|
||||
// 按阶梯顺序尝试恢复:
|
||||
// 1. Aggressive Compact — 激进微压缩(保留更少的工具结果)
|
||||
// 2. Reactive Compact — 使用 LLM 摘要压缩对话历史
|
||||
// 3. Escalate Tokens — 临时提升 token 上限到 64k
|
||||
// 4. Multi-Turn — 注入 metacognitive 消息分步处理
|
||||
// 5. Surface — 放弃恢复,暴露错误给用户
|
||||
//
|
||||
// 每一步都有 `has_attempted` 守卫,防止无限循环。
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::token_budget::TokenBudget;
|
||||
|
||||
/// 错误类型分类
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ErrorKind {
|
||||
/// 上下文过长(prompt too long / 413)
|
||||
PromptTooLong,
|
||||
/// Token 耗尽
|
||||
TokenExhausted,
|
||||
/// 模型错误
|
||||
ModelError(String),
|
||||
/// 超时
|
||||
Timeout,
|
||||
/// 限流(HTTP 429)
|
||||
RateLimited,
|
||||
/// 服务过载(HTTP 529)
|
||||
Overloaded,
|
||||
}
|
||||
|
||||
/// 恢复步骤
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum RecoveryStep {
|
||||
/// 尝试更激进的 micro_compact
|
||||
AggressiveCompact,
|
||||
/// 使用 LLM 摘要压缩
|
||||
ReactiveCompact,
|
||||
/// 提升 token 上限
|
||||
EscalateTokens { new_hard_limit: usize },
|
||||
/// 分轮恢复(注入 meta 消息)
|
||||
MultiTurn,
|
||||
/// 放弃,暴露错误
|
||||
Surface,
|
||||
/// 指数退避重试(用于 429/529 瞬态错误)
|
||||
RetryWithBackoff { attempt: u32, delay_ms: u64 },
|
||||
}
|
||||
|
||||
/// 恢复尝试追踪
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RecoveryAttempts {
|
||||
pub aggressive_compact: bool,
|
||||
pub reactive_compact: bool,
|
||||
pub escalate_tokens: bool,
|
||||
pub multi_turn: bool,
|
||||
}
|
||||
|
||||
impl RecoveryAttempts {
|
||||
pub fn new() -> Self {
|
||||
RecoveryAttempts {
|
||||
aggressive_compact: false,
|
||||
reactive_compact: false,
|
||||
escalate_tokens: false,
|
||||
multi_turn: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 是否有未尝试的恢复步骤
|
||||
pub fn has_remaining(&self) -> bool {
|
||||
!self.aggressive_compact
|
||||
|| !self.reactive_compact
|
||||
|| !self.escalate_tokens
|
||||
|| !self.multi_turn
|
||||
}
|
||||
|
||||
/// 获取下一个应尝试的恢复步骤
|
||||
pub fn next_step(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
|
||||
// 429/529 使用退避重试,不消耗上下文恢复步骤
|
||||
if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) {
|
||||
return Some(RecoveryStep::RetryWithBackoff {
|
||||
attempt: 0,
|
||||
delay_ms: 500,
|
||||
});
|
||||
}
|
||||
|
||||
match error_kind {
|
||||
ErrorKind::PromptTooLong | ErrorKind::TokenExhausted => {
|
||||
if !self.aggressive_compact {
|
||||
self.aggressive_compact = true;
|
||||
return Some(RecoveryStep::AggressiveCompact);
|
||||
}
|
||||
if !self.reactive_compact {
|
||||
self.reactive_compact = true;
|
||||
return Some(RecoveryStep::ReactiveCompact);
|
||||
}
|
||||
if !self.escalate_tokens {
|
||||
self.escalate_tokens = true;
|
||||
return Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000,
|
||||
});
|
||||
}
|
||||
if !self.multi_turn {
|
||||
self.multi_turn = true;
|
||||
return Some(RecoveryStep::MultiTurn);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// 非上下文相关错误,直接暴露
|
||||
if !self.multi_turn {
|
||||
self.multi_turn = true;
|
||||
return Some(RecoveryStep::Surface);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RecoveryAttempts {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 错误恢复器
|
||||
pub struct ErrorRecovery {
|
||||
/// 恢复步骤追踪
|
||||
pub attempts: RecoveryAttempts,
|
||||
/// Token 预算(用于 escalate 步骤)
|
||||
pub token_budget: TokenBudget,
|
||||
}
|
||||
|
||||
impl ErrorRecovery {
|
||||
/// 创建新的错误恢复器
|
||||
pub fn new(token_budget: TokenBudget) -> Self {
|
||||
ErrorRecovery {
|
||||
attempts: RecoveryAttempts::new(),
|
||||
token_budget,
|
||||
}
|
||||
}
|
||||
|
||||
/// 尝试从错误中恢复。
|
||||
///
|
||||
/// 返回 `Some(RecoveryStep)` 表示找到了恢复步骤(调用方应执行该步骤后重试)。
|
||||
/// 返回 `None` 表示所有步骤已尝试完毕,应暴露错误给用户。
|
||||
pub fn try_recover(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
|
||||
let step = self.attempts.next_step(error_kind);
|
||||
match &step {
|
||||
Some(RecoveryStep::AggressiveCompact) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 1/4: AggressiveCompact");
|
||||
}
|
||||
Some(RecoveryStep::ReactiveCompact) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 2/4: ReactiveCompact");
|
||||
}
|
||||
Some(RecoveryStep::EscalateTokens { new_hard_limit }) => {
|
||||
info!(
|
||||
"[ErrorRecovery] 尝试步骤 3/4: EscalateTokens → {}",
|
||||
new_hard_limit
|
||||
);
|
||||
self.token_budget.escalate_hard_limit(*new_hard_limit);
|
||||
}
|
||||
Some(RecoveryStep::RetryWithBackoff { attempt, delay_ms }) => {
|
||||
info!(
|
||||
"[ErrorRecovery] 退避重试: attempt={}, delay={}ms",
|
||||
attempt, delay_ms
|
||||
);
|
||||
}
|
||||
Some(RecoveryStep::MultiTurn) => {
|
||||
info!("[ErrorRecovery] 尝试步骤 4/4: MultiTurn");
|
||||
}
|
||||
Some(RecoveryStep::Surface) => {
|
||||
warn!("[ErrorRecovery] 无法恢复,暴露错误");
|
||||
}
|
||||
None => {
|
||||
warn!("[ErrorRecovery] 所有恢复步骤已尝试完毕");
|
||||
}
|
||||
}
|
||||
step
|
||||
}
|
||||
|
||||
/// 生成 multi-turn 恢复消息(注入到对话中以继续处理)
|
||||
pub fn multi_turn_message() -> String {
|
||||
"由于 token 限制,当前回答被截断。请基于已收集的信息继续分析,\
|
||||
重点关注尚未完成的部分。你可以:\n\
|
||||
1. 总结已有发现\n\
|
||||
2. 使用 compress_context 手动压缩上下文\n\
|
||||
3. 分步完成剩余工作"
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// 检查是否需要恢复(错误是否可恢复)
|
||||
pub fn is_recoverable(error_kind: &ErrorKind) -> bool {
|
||||
matches!(
|
||||
error_kind,
|
||||
ErrorKind::PromptTooLong
|
||||
| ErrorKind::TokenExhausted
|
||||
| ErrorKind::RateLimited
|
||||
| ErrorKind::Overloaded
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算指数退避延迟(毫秒)。
|
||||
///
|
||||
/// 公式:min(500 * 2^attempt, 32000) + 25% 随机抖动
|
||||
/// 如果有 Retry-After header,优先使用。
|
||||
pub fn backoff_delay(attempt: u32, retry_after_secs: Option<u64>) -> u64 {
|
||||
if let Some(ra) = retry_after_secs {
|
||||
return ra * 1000;
|
||||
}
|
||||
let base = 500u64 * 2u64.pow(attempt.min(6)); // cap at 2^6 = 64 → 32000ms
|
||||
let base = base.min(32_000);
|
||||
// Simple deterministic jitter using attempt (avoid rand dependency)
|
||||
let jitter = (base / 4) * (attempt as u64 % 5) / 5;
|
||||
base + jitter
|
||||
}
|
||||
|
||||
/// 解析错误字符串中的 Retry-After header 值。
|
||||
///
|
||||
/// 期望格式: `retry_after=Some(N)` 出现在错误消息中。
|
||||
pub fn parse_retry_after(error_str: &str) -> Option<u64> {
|
||||
if let Some(pos) = error_str.find("retry_after=Some(") {
|
||||
let prefix_len = "retry_after=Some(".len(); // 19
|
||||
let rest = &error_str[pos + prefix_len..];
|
||||
if let Some(end) = rest.find(')') {
|
||||
return rest[..end].parse().ok();
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 从错误字符串分类错误类型。
|
||||
/// 解析 LLM API 返回的错误消息,映射到 ErrorKind。
|
||||
pub fn classify_error(error_str: &str) -> ErrorKind {
|
||||
let lower = error_str.to_lowercase();
|
||||
|
||||
// 先检测限流/过载(HTTP 状态码检查)
|
||||
if lower.contains("429")
|
||||
|| lower.contains("rate limit")
|
||||
|| lower.contains("rate_limit")
|
||||
|| lower.contains("too many requests")
|
||||
{
|
||||
return ErrorKind::RateLimited;
|
||||
}
|
||||
if lower.contains("529")
|
||||
|| lower.contains("overloaded")
|
||||
|| lower.contains("overload")
|
||||
|| lower.contains("service overloaded")
|
||||
{
|
||||
return ErrorKind::Overloaded;
|
||||
}
|
||||
|
||||
if lower.contains("prompt_too_long")
|
||||
|| lower.contains("prompt too long")
|
||||
|| lower.contains("context length")
|
||||
|| lower.contains("413")
|
||||
|| lower.contains("context_window_exceeded")
|
||||
|| lower.contains("input length")
|
||||
{
|
||||
return ErrorKind::PromptTooLong;
|
||||
}
|
||||
|
||||
if lower.contains("max_tokens")
|
||||
|| lower.contains("token limit")
|
||||
|| lower.contains("token_exhausted")
|
||||
|| lower.contains("maximum context length")
|
||||
|| lower.contains("reduce the length")
|
||||
{
|
||||
return ErrorKind::TokenExhausted;
|
||||
}
|
||||
|
||||
if lower.contains("timeout")
|
||||
|| lower.contains("timed out")
|
||||
|| lower.contains("deadline exceeded")
|
||||
|| lower.contains("408")
|
||||
|| lower.contains("504")
|
||||
{
|
||||
return ErrorKind::Timeout;
|
||||
}
|
||||
|
||||
// 默认归类为模型错误
|
||||
ErrorKind::ModelError(error_str.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_all_steps_sequence() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::AggressiveCompact)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::ReactiveCompact)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::PromptTooLong),
|
||||
Some(RecoveryStep::MultiTurn)
|
||||
);
|
||||
// 所有步骤已尝试
|
||||
assert_eq!(attempts.next_step(&ErrorKind::PromptTooLong), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_error_goes_straight_to_surface() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::ModelError("test".into())),
|
||||
Some(RecoveryStep::Surface)
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.next_step(&ErrorKind::ModelError("test".into())),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_remaining() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
assert!(attempts.has_remaining());
|
||||
|
||||
// 消耗所有步骤
|
||||
for _ in 0..4 {
|
||||
attempts.next_step(&ErrorKind::PromptTooLong);
|
||||
}
|
||||
assert!(!attempts.has_remaining());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_recoverable() {
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::PromptTooLong));
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::TokenExhausted));
|
||||
assert!(!ErrorRecovery::is_recoverable(&ErrorKind::ModelError(
|
||||
"test".into()
|
||||
)));
|
||||
assert!(!ErrorRecovery::is_recoverable(&ErrorKind::Timeout));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_rate_limited_429() {
|
||||
let kind = classify_error("HTTP 429: Too Many Requests");
|
||||
assert_eq!(kind, ErrorKind::RateLimited);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_overloaded_529() {
|
||||
let kind = classify_error("HTTP 529: Service Overloaded");
|
||||
assert_eq!(kind, ErrorKind::Overloaded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limited_is_recoverable() {
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited));
|
||||
assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_retry_after() {
|
||||
let err = "HTTP 429: retry_after=Some(30)";
|
||||
assert_eq!(parse_retry_after(err), Some(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_retry_after_none() {
|
||||
let err = "HTTP 500: Internal Server Error";
|
||||
assert_eq!(parse_retry_after(err), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_delay() {
|
||||
// Attempt 0: 500 + jitter
|
||||
let d0 = backoff_delay(0, None);
|
||||
assert!(d0 >= 500 && d0 <= 700);
|
||||
|
||||
// Attempt 3: 500*8=4000 + jitter
|
||||
let d3 = backoff_delay(3, None);
|
||||
assert!(d3 >= 4000 && d3 <= 5000);
|
||||
|
||||
// Capped at 32s
|
||||
let d10 = backoff_delay(10, None);
|
||||
assert!(d10 <= 40_000);
|
||||
|
||||
// Retry-After takes priority
|
||||
let d_ra = backoff_delay(0, Some(15));
|
||||
assert_eq!(d_ra, 15000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limited_goes_to_retry() {
|
||||
let mut attempts = RecoveryAttempts::new();
|
||||
let step = attempts.next_step(&ErrorKind::RateLimited);
|
||||
assert!(matches!(step, Some(RecoveryStep::RetryWithBackoff { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_escalation() {
|
||||
let budget = TokenBudget::new(32_000, 40_000);
|
||||
let mut recovery = ErrorRecovery::new(budget);
|
||||
assert_eq!(recovery.token_budget.hard_limit, 40_000);
|
||||
|
||||
recovery.attempts.next_step(&ErrorKind::PromptTooLong); // aggressive
|
||||
recovery.attempts.next_step(&ErrorKind::PromptTooLong); // reactive
|
||||
let step = recovery.try_recover(&ErrorKind::PromptTooLong); // escalate
|
||||
|
||||
assert_eq!(
|
||||
step,
|
||||
Some(RecoveryStep::EscalateTokens {
|
||||
new_hard_limit: 64_000
|
||||
})
|
||||
);
|
||||
assert_eq!(recovery.token_budget.hard_limit, 64_000);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
// src/agent/runtime/executor.rs
|
||||
//
|
||||
// 工具调用执行器:验证 → PreToolUse hooks → 并行执行 → 结果收集 → PostToolUse hooks。
|
||||
|
||||
use futures_util::stream::FuturesUnordered;
|
||||
use futures_util::StreamExt;
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{ChatMessage, ToolCall};
|
||||
|
||||
use super::file_cache::FileStateCache;
|
||||
use super::permission::PermissionChecker;
|
||||
use super::{AgentStreamEvent, DuplicateDetector};
|
||||
use crate::agent::hooks::{HookRegistry, PostToolUseContext, PreToolUseContext};
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result;
|
||||
use crate::agent::tools::{InterruptBehavior, ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 准备好的工具调用
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreparedCall {
|
||||
pub tool_call_id: String,
|
||||
pub tool_name: String,
|
||||
pub args: serde_json::Value,
|
||||
}
|
||||
|
||||
/// 单次工具执行后的消息 + 元数据
|
||||
pub struct ToolResultMessage {
|
||||
pub chat_message: ChatMessage,
|
||||
pub was_error: bool,
|
||||
}
|
||||
|
||||
/// 工具执行结果摘要
|
||||
pub struct ToolExecutionResult {
|
||||
/// 每条工具调用对应的 tool_result 消息(供调用方 push 到 messages)
|
||||
pub tool_messages: Vec<ToolResultMessage>,
|
||||
pub was_cancelled: bool,
|
||||
pub had_duplicate: bool,
|
||||
}
|
||||
|
||||
/// 验证工具调用:死循环检测 + 参数解析。
|
||||
///
|
||||
/// 返回 (prepared_calls, has_duplicate)。
|
||||
/// 死循环或参数无效时,错误消息直接注入到 messages。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn validate_and_prepare(
|
||||
tool_calls: &[ToolCall],
|
||||
duplicate_detector: &mut DuplicateDetector,
|
||||
duplicate_threshold: usize,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
) -> (Vec<PreparedCall>, bool) {
|
||||
let mut prepared_calls: Vec<PreparedCall> = Vec::new();
|
||||
let mut has_duplicate = false;
|
||||
|
||||
for tool_call in tool_calls {
|
||||
let tool_name = &tool_call.function.name;
|
||||
let tool_args_str = &tool_call.function.arguments;
|
||||
|
||||
// 死循环检测
|
||||
if duplicate_detector.record(tool_name, tool_args_str, duplicate_threshold) {
|
||||
warn!(
|
||||
"[Executor] 检测到死循环:{} 连续调用 {} 次",
|
||||
tool_name, duplicate_threshold
|
||||
);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用,已自动终止循环。", tool_name),
|
||||
});
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!(
|
||||
"错误:工具 {} 被连续重复调用 {} 次,参数完全相同。\
|
||||
请停止重复调用并直接给出目前收集到的答案。",
|
||||
tool_name, duplicate_threshold
|
||||
),
|
||||
);
|
||||
messages.push(error_msg);
|
||||
has_duplicate = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
// 解析参数
|
||||
let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
let error_output = format!("工具参数 JSON 解析失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output);
|
||||
save_tool_message_sync(db, session_id, turn_index, step, &tool_msg);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
prepared_calls.push(PreparedCall {
|
||||
tool_call_id: tool_call.id.clone(),
|
||||
tool_name: tool_name.clone(),
|
||||
args,
|
||||
});
|
||||
}
|
||||
|
||||
(prepared_calls, has_duplicate)
|
||||
}
|
||||
|
||||
/// 并行执行所有准备好的工具调用。
|
||||
///
|
||||
/// 流程:
|
||||
/// 1. 权限检查(deny 规则阻止不可执行工具)
|
||||
/// 2. 发送 ToolCall SSE 事件
|
||||
/// 3. 运行 PreToolUse hooks
|
||||
/// 4. 工具分区 + 并行执行(并发安全工具一批并行,不安全工具单独串行)
|
||||
/// 5. 收集结果、发送 ToolResult SSE、运行 PostToolUse hooks
|
||||
/// 6. 返回 ToolResultMessage 列表供调用方推入 messages
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn execute_parallel(
|
||||
prepared_calls: &[PreparedCall],
|
||||
tool_registry: &ToolRegistry,
|
||||
app_state: Arc<AppState>,
|
||||
hook_registry: &HookRegistry,
|
||||
_permission_checker: Option<&PermissionChecker>,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
agent_name: &str,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
tool_timeout_secs: u64,
|
||||
max_output_chars: usize,
|
||||
read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
) -> ToolExecutionResult {
|
||||
if prepared_calls.is_empty() {
|
||||
return ToolExecutionResult {
|
||||
tool_messages: Vec::new(),
|
||||
was_cancelled: false,
|
||||
had_duplicate: false,
|
||||
};
|
||||
}
|
||||
|
||||
let sid = session_id.to_string();
|
||||
|
||||
// Phase 1: 发送 ToolCall SSE 事件
|
||||
for prep in prepared_calls {
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
name: prep.tool_name.clone(),
|
||||
arguments: prep.args.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
|
||||
// Phase 2: PreToolUse hooks — 收集修改后的参数和附加上下文
|
||||
let exec_start = std::time::Instant::now();
|
||||
let mut mutated_args: Vec<serde_json::Value> = Vec::new();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
for prep in prepared_calls {
|
||||
let hook_ctx = PreToolUseContext {
|
||||
session_id: sid.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
tool_args: prep.args.clone(),
|
||||
step,
|
||||
};
|
||||
let result = hook_registry.run_pre_tool_use(&hook_ctx).await;
|
||||
if result.action.is_blocked() {
|
||||
let reason = result.action.block_reason().unwrap_or("unknown");
|
||||
warn!(
|
||||
"[Executor] PreToolUse hook 阻止了 {} 的执行: {}",
|
||||
prep.tool_name, reason
|
||||
);
|
||||
}
|
||||
// 使用 hook 可能修改后的参数
|
||||
mutated_args.push(result.final_args);
|
||||
if let Some(ctx) = result.additional_context {
|
||||
additional_contexts.push(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 3: 并行执行
|
||||
let cancelled = Arc::new(AtomicBool::new(false));
|
||||
let cancel_flag = cancelled.clone();
|
||||
let app_state_ref = app_state.clone();
|
||||
let sid_ref = sid.clone();
|
||||
|
||||
let cancel_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(locked) = app_state_ref.cancelled_runs.lock() {
|
||||
if locked.contains(&sid_ref) {
|
||||
cancel_flag.store(true, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs);
|
||||
|
||||
// Phase 3: 使用 FuturesUnordered 进行渐进式并行执行。
|
||||
// 每个工具完成后立即发送 SSE ToolResult 事件到前端(非阻塞),
|
||||
// 而后台继续等待其他工具完成。快工具的结果不会因慢工具而延迟。
|
||||
let mut exec_futs: FuturesUnordered<_> = prepared_calls
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, prep)| {
|
||||
let tool_name = prep.tool_name.clone();
|
||||
let args = mutated_args
|
||||
.get(i)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx = ToolContext::with_file_cache(app_state.clone(), read_file_state.clone());
|
||||
let cancelled = cancelled.clone();
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
Box::pin(async move {
|
||||
let output = match tool_opt {
|
||||
Some(tool) => {
|
||||
let interrupt_behavior = tool.interrupt_behavior();
|
||||
let is_blocking = interrupt_behavior == InterruptBehavior::Block;
|
||||
|
||||
let tool_fut = tool.execute(args, &tool_ctx);
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if !is_blocking && cancelled.load(Ordering::SeqCst) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
res = tokio::time::timeout(timeout_dur, tool_fut) => {
|
||||
match res {
|
||||
Ok(output) => output,
|
||||
Err(_) => ToolOutput::error(format!(
|
||||
"工具 {} 执行超时({}秒)",
|
||||
tool_name,
|
||||
timeout_dur.as_secs()
|
||||
)),
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
ToolOutput::error("执行已被用户取消")
|
||||
}
|
||||
}
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
let was_cancelled = cancelled.load(Ordering::SeqCst);
|
||||
(
|
||||
prep.tool_call_id.clone(),
|
||||
prep.tool_name.clone(),
|
||||
prep.args.clone(),
|
||||
output,
|
||||
was_cancelled,
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut tool_messages: Vec<ToolResultMessage> = Vec::new();
|
||||
let mut was_cancelled = false;
|
||||
|
||||
// 渐进式处理结果:每个工具一完成就立即处理(SSE 事件 + PostToolUse hook + 持久化)
|
||||
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
|
||||
exec_futs.next().await
|
||||
{
|
||||
if cancelled_flag {
|
||||
was_cancelled = true;
|
||||
}
|
||||
|
||||
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
|
||||
|
||||
// SSE 事件 — 立即推送到前端
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
// 输出处理:小结果直接传递,大结果持久化到磁盘并返回 stub
|
||||
let tool_results_dir = app_state.config.library_dir.join("tool-results");
|
||||
let (processed_content, _persisted_path) = maybe_persist_tool_result(
|
||||
&output.content,
|
||||
&tool_call_id,
|
||||
max_output_chars,
|
||||
&tool_results_dir,
|
||||
);
|
||||
|
||||
// PostToolUse hook
|
||||
let post_ctx = PostToolUseContext {
|
||||
session_id: sid.clone(),
|
||||
agent_name: agent_name.to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
tool_args,
|
||||
output_content: processed_content.clone(),
|
||||
is_error: output.is_error,
|
||||
step,
|
||||
elapsed_ms,
|
||||
};
|
||||
let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
|
||||
let final_content = post_result.final_content;
|
||||
|
||||
let chat_message = ChatMessage::tool_result(&tool_call_id, &final_content);
|
||||
|
||||
// 持久化到数据库(fire-and-forget)
|
||||
save_tool_message_sync(db, &sid, turn_index, step, &chat_message);
|
||||
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message,
|
||||
was_error: output.is_error,
|
||||
});
|
||||
}
|
||||
|
||||
cancel_handle.abort();
|
||||
|
||||
ToolExecutionResult {
|
||||
tool_messages,
|
||||
was_cancelled,
|
||||
had_duplicate: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 同步保存 tool 角色消息到数据库。
|
||||
fn save_tool_message_sync(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step_index: usize,
|
||||
msg: &ChatMessage,
|
||||
) {
|
||||
let db_clone = db.clone();
|
||||
let session_id = session_id.to_string();
|
||||
let content = msg.content.as_deref().unwrap_or("").to_string();
|
||||
let tool_call_id = msg.tool_call_id.clone();
|
||||
// fire-and-forget: tool 消息保存失败不影响主流程
|
||||
tokio::spawn(async move {
|
||||
let token_count = content.len() as i32 / 4;
|
||||
if let Err(e) = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, token_count, agent_name) \
|
||||
VALUES (?, ?, ?, 'tool', ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step_index as i32)
|
||||
.bind(&content)
|
||||
.bind(&tool_call_id)
|
||||
.bind(token_count)
|
||||
.bind("lead")
|
||||
.execute(&db_clone)
|
||||
.await
|
||||
{
|
||||
warn!("[Executor] 保存 tool 消息失败(非致命): {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,571 @@
|
||||
// src/agent/runtime/file_cache.rs
|
||||
//
|
||||
// 文件状态缓存 — 参考 Claude Code FileStateCache 设计。
|
||||
//
|
||||
// 在 Read 工具调用前检查缓存:
|
||||
// 1. 路径已缓存 → 读取磁盘 mtime → mtime 相同 + offset/limit 一致 → 返回 stub
|
||||
// 2. mtime 不同或新文件 → 正常读取 → 写入缓存
|
||||
//
|
||||
// 压缩时:
|
||||
// - 压缩前:快照缓存到普通对象
|
||||
// - 压缩后:清空缓存,将最近 N 个文件作为上下文注入
|
||||
//
|
||||
// 缓存上限:100 个条目,25MB 内容总大小(LRU 自动淘汰)。
|
||||
|
||||
use lru::LruCache;
|
||||
use std::num::NonZeroUsize;
|
||||
use std::path::Path;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// 缓存条目最大数量
|
||||
pub const MAX_ENTRIES: usize = 100;
|
||||
/// 缓存内容总大小上限(25MB)
|
||||
pub const MAX_CACHE_SIZE_BYTES: usize = 25 * 1024 * 1024;
|
||||
|
||||
/// 文件不变时的占位消息(参考 Claude Code FILE_UNCHANGED_STUB)
|
||||
pub const FILE_UNCHANGED_STUB: &str =
|
||||
"File unchanged since last read. The content from the earlier read_file tool_result \
|
||||
in this conversation is still current — refer to that instead of re-reading.";
|
||||
|
||||
/// 压缩后恢复的最大文件数
|
||||
pub const POST_COMPACT_MAX_FILES_TO_RESTORE: usize = 5;
|
||||
/// 压缩后恢复的每文件最大 token 数(~字符数)
|
||||
pub const POST_COMPACT_MAX_CHARS_PER_FILE: usize = 4_000;
|
||||
|
||||
/// 单个文件的缓存状态
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FileState {
|
||||
/// 上次读取的文件内容
|
||||
pub content: String,
|
||||
/// 文件修改时间(Unix 时间戳,秒级)
|
||||
pub timestamp: i64,
|
||||
/// 读取起始行(1-based)
|
||||
pub offset: usize,
|
||||
/// 行数限制
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// 文件状态快照(用于压缩前后传递,纯数据,不含 LRU 结构)
|
||||
pub type FileStateSnapshot = Vec<(String, FileState)>;
|
||||
|
||||
/// 文件状态缓存。
|
||||
///
|
||||
/// 包装 `LruCache<String, FileState>` + 内容总大小追踪。
|
||||
/// 通过 `Arc<Mutex<FileStateCache>>` 在工具调用间共享。
|
||||
pub struct FileStateCache {
|
||||
cache: LruCache<String, FileState>,
|
||||
/// 当前缓存中所有内容的字节数总和(近似,使用 content.len())
|
||||
current_size_bytes: usize,
|
||||
/// 最大字节数
|
||||
max_size_bytes: usize,
|
||||
}
|
||||
|
||||
impl FileStateCache {
|
||||
/// 创建新的缓存实例
|
||||
pub fn new() -> Self {
|
||||
let max_entries = NonZeroUsize::new(MAX_ENTRIES).unwrap();
|
||||
FileStateCache {
|
||||
cache: LruCache::new(max_entries),
|
||||
current_size_bytes: 0,
|
||||
max_size_bytes: MAX_CACHE_SIZE_BYTES,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带自定义参数的新缓存
|
||||
pub fn with_limits(max_entries: usize, max_size_bytes: usize) -> Self {
|
||||
let max_entries =
|
||||
NonZeroUsize::new(max_entries).unwrap_or(NonZeroUsize::new(MAX_ENTRIES).unwrap());
|
||||
FileStateCache {
|
||||
cache: LruCache::new(max_entries),
|
||||
current_size_bytes: 0,
|
||||
max_size_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
/// 规范化路径 key(确保一致性)
|
||||
fn normalize_key(path: &str) -> String {
|
||||
// 去除尾随斜杠,规范化重复斜杠
|
||||
let p = Path::new(path);
|
||||
// 尝试 canonicalize(跟随符号链接),失败则用简单的字符串规范化
|
||||
match p.canonicalize() {
|
||||
Ok(canon) => canon.to_string_lossy().to_string(),
|
||||
Err(_) => {
|
||||
// 简单规范化:折叠重复的 /
|
||||
let mut result = String::with_capacity(path.len());
|
||||
let mut prev_slash = false;
|
||||
for ch in path.chars() {
|
||||
if ch == '/' || ch == '\\' {
|
||||
if !prev_slash {
|
||||
result.push('/');
|
||||
prev_slash = true;
|
||||
}
|
||||
} else {
|
||||
result.push(ch);
|
||||
prev_slash = false;
|
||||
}
|
||||
}
|
||||
// 去除尾随 /
|
||||
if result.ends_with('/') && result.len() > 1 {
|
||||
result.pop();
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取缓存的条目,返回 Some(&FileState) 若存在
|
||||
pub fn get(&mut self, path: &str) -> Option<&FileState> {
|
||||
let key = Self::normalize_key(path);
|
||||
self.cache.get(&key)
|
||||
}
|
||||
|
||||
/// 写入缓存条目,自动处理容量限制
|
||||
pub fn set(&mut self, path: &str, state: FileState) {
|
||||
let key = Self::normalize_key(path);
|
||||
let content_len = state.content.len();
|
||||
|
||||
// 如果 key 已存在,先减去旧内容的 size
|
||||
if let Some(old) = self.cache.get(&key) {
|
||||
self.current_size_bytes = self.current_size_bytes.saturating_sub(old.content.len());
|
||||
}
|
||||
|
||||
// 驱逐旧条目直到有足够空间
|
||||
while self.current_size_bytes + content_len > self.max_size_bytes && !self.cache.is_empty() {
|
||||
if let Some((_, evicted)) = self.cache.pop_lru() {
|
||||
self.current_size_bytes = self
|
||||
.current_size_bytes
|
||||
.saturating_sub(evicted.content.len());
|
||||
}
|
||||
}
|
||||
|
||||
// 如果单个文件超过上限,仍存储但记录警告
|
||||
if content_len > self.max_size_bytes {
|
||||
warn!(
|
||||
"[FileCache] 单个文件内容 ({} bytes) 超过缓存上限 ({} bytes)",
|
||||
content_len, self.max_size_bytes
|
||||
);
|
||||
}
|
||||
|
||||
self.current_size_bytes += content_len;
|
||||
self.cache.push(key, state);
|
||||
|
||||
info!(
|
||||
"[FileCache] 缓存写入: path={}, size={} bytes, cache_entries={}, cache_size={}",
|
||||
path,
|
||||
content_len,
|
||||
self.cache.len(),
|
||||
self.current_size_bytes
|
||||
);
|
||||
}
|
||||
|
||||
/// 检查 key 是否存在
|
||||
pub fn contains(&mut self, path: &str) -> bool {
|
||||
let key = Self::normalize_key(path);
|
||||
self.cache.contains(&key)
|
||||
}
|
||||
|
||||
/// 删除缓存条目
|
||||
pub fn remove(&mut self, path: &str) -> bool {
|
||||
let key = Self::normalize_key(path);
|
||||
if let Some(removed) = self.cache.pop(&key) {
|
||||
self.current_size_bytes = self
|
||||
.current_size_bytes
|
||||
.saturating_sub(removed.content.len());
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 清空缓存
|
||||
pub fn clear(&mut self) {
|
||||
self.cache.clear();
|
||||
self.current_size_bytes = 0;
|
||||
info!("[FileCache] 缓存已清空");
|
||||
}
|
||||
|
||||
/// 缓存条目数
|
||||
pub fn len(&self) -> usize {
|
||||
self.cache.len()
|
||||
}
|
||||
|
||||
/// 缓存是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.cache.len() == 0
|
||||
}
|
||||
|
||||
/// 当前缓存内容总大小(近似字节数)
|
||||
pub fn current_size_bytes(&self) -> usize {
|
||||
self.current_size_bytes
|
||||
}
|
||||
|
||||
// ── 压缩集成 ──
|
||||
|
||||
/// 生成快照(纯数据,不含 LRU 结构)。
|
||||
/// 在压缩前调用,用于压缩后恢复文件上下文。
|
||||
pub fn to_snapshot(&mut self) -> FileStateSnapshot {
|
||||
// 按 timestamp 降序排序(最近读的在前)
|
||||
let mut entries: Vec<(String, FileState)> = self
|
||||
.cache
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.collect();
|
||||
entries.sort_by(|a, b| b.1.timestamp.cmp(&a.1.timestamp));
|
||||
entries
|
||||
}
|
||||
|
||||
/// 从快照恢复指定数量的最近文件到缓存。
|
||||
/// 在压缩后调用。
|
||||
pub fn restore_from_snapshot(&mut self, snapshot: &FileStateSnapshot, max_files: usize) {
|
||||
for (path, state) in snapshot.iter().take(max_files) {
|
||||
// 不恢复过大的文件(已有提示说可能过时)
|
||||
if state.content.len() > POST_COMPACT_MAX_CHARS_PER_FILE {
|
||||
continue;
|
||||
}
|
||||
self.set(path, state.clone());
|
||||
}
|
||||
info!(
|
||||
"[FileCache] 从快照恢复了 {} 个文件 (快照大小: {})",
|
||||
self.len().min(max_files),
|
||||
snapshot.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// 从快照生成上 下文注入文本(用于压缩后注入到对话中)。
|
||||
/// 返回格式化的 markdown 块列表。
|
||||
pub fn build_restore_context(snapshot: &FileStateSnapshot, max_files: usize) -> Vec<String> {
|
||||
if snapshot.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut contexts: Vec<String> = Vec::new();
|
||||
let mut used_chars = 0usize;
|
||||
let total_budget = POST_COMPACT_MAX_CHARS_PER_FILE * max_files;
|
||||
|
||||
for (path, state) in snapshot.iter().take(max_files) {
|
||||
let preview: String = state
|
||||
.content
|
||||
.chars()
|
||||
.take(POST_COMPACT_MAX_CHARS_PER_FILE)
|
||||
.collect();
|
||||
let truncated = if state.content.len() > preview.len() {
|
||||
format!(
|
||||
"{}…\n[内容已截断: {} 字符 → {} 字符]",
|
||||
preview,
|
||||
state.content.len(),
|
||||
preview.len()
|
||||
)
|
||||
} else {
|
||||
preview
|
||||
};
|
||||
|
||||
if used_chars + truncated.len() > total_budget {
|
||||
break;
|
||||
}
|
||||
|
||||
let block = format!(
|
||||
"[压缩后恢复: {}]\n上次读取时间戳: {}\n内容:\n```\n{}\n```",
|
||||
path, state.timestamp, truncated
|
||||
);
|
||||
used_chars += block.len();
|
||||
contexts.push(block);
|
||||
}
|
||||
|
||||
if !contexts.is_empty() {
|
||||
info!(
|
||||
"[FileCache] 生成压缩恢复上下文: {} 文件, {} chars",
|
||||
contexts.len(),
|
||||
used_chars
|
||||
);
|
||||
}
|
||||
|
||||
contexts
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FileStateCache {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取文件的当前 mtime(Unix 时间戳,秒)。
|
||||
/// 失败时返回 None。
|
||||
pub fn get_file_mtime(path: &str) -> Option<i64> {
|
||||
match std::fs::metadata(path) {
|
||||
Ok(meta) => match meta.modified() {
|
||||
Ok(time) => match time.duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(d) => Some(d.as_secs() as i64),
|
||||
Err(_) => None,
|
||||
},
|
||||
Err(_) => None,
|
||||
},
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_new_cache_is_empty() {
|
||||
let cache = FileStateCache::new();
|
||||
assert!(cache.is_empty());
|
||||
assert_eq!(cache.len(), 0);
|
||||
assert_eq!(cache.current_size_bytes(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_and_get() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "hello world".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(!cache.is_empty());
|
||||
|
||||
let entry = cache.get("/tmp/test.txt");
|
||||
assert!(entry.is_some());
|
||||
assert_eq!(entry.unwrap().content, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_path_normalization() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp//test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 规范化后的路径应该能命中
|
||||
assert!(cache.get("/tmp/test.txt").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains() {
|
||||
let mut cache = FileStateCache::new();
|
||||
assert!(!cache.contains("/tmp/test.txt"));
|
||||
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(cache.contains("/tmp/test.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/test.txt",
|
||||
FileState {
|
||||
content: "test".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
assert!(cache.remove("/tmp/test.txt"));
|
||||
assert!(cache.is_empty());
|
||||
assert!(!cache.remove("/tmp/test.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "a".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/b.txt",
|
||||
FileState {
|
||||
content: "b".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.clear();
|
||||
assert!(cache.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snapshot_sorted_by_timestamp_desc() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/old.txt",
|
||||
FileState {
|
||||
content: "old".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/new.txt",
|
||||
FileState {
|
||||
content: "new".to_string(),
|
||||
timestamp: 3000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/mid.txt",
|
||||
FileState {
|
||||
content: "mid".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
assert_eq!(snapshot.len(), 3);
|
||||
// 按 timestamp 降序
|
||||
assert_eq!(snapshot[0].1.timestamp, 3000);
|
||||
assert_eq!(snapshot[1].1.timestamp, 2000);
|
||||
assert_eq!(snapshot[2].1.timestamp, 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_snapshot() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "a".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/b.txt",
|
||||
FileState {
|
||||
content: "b".to_string(),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
cache.clear();
|
||||
|
||||
cache.restore_from_snapshot(&snapshot, 1);
|
||||
assert_eq!(cache.len(), 1);
|
||||
// 应该恢复 timestamp 最高的
|
||||
assert!(cache.get("/tmp/b.txt").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_restore_context() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set(
|
||||
"/tmp/a.txt",
|
||||
FileState {
|
||||
content: "file a content here".to_string(),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
|
||||
let snapshot = cache.to_snapshot();
|
||||
let contexts = FileStateCache::build_restore_context(&snapshot, 2);
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert!(contexts[0].contains("/tmp/a.txt"));
|
||||
assert!(contexts[0].contains("file a content here"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_snapshot_builds_no_context() {
|
||||
let contexts = FileStateCache::build_restore_context(&Vec::new(), 5);
|
||||
assert!(contexts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction_on_size() {
|
||||
// 创建一个小容量缓存(仅 200 bytes)
|
||||
let mut cache = FileStateCache::with_limits(100, 200);
|
||||
|
||||
// 写入 3 个 100 字节内容 → 应触发 LRU 淘汰
|
||||
cache.set(
|
||||
"/tmp/1.txt",
|
||||
FileState {
|
||||
content: "x".repeat(100),
|
||||
timestamp: 1000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
cache.set(
|
||||
"/tmp/2.txt",
|
||||
FileState {
|
||||
content: "y".repeat(100),
|
||||
timestamp: 2000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 此时应该有 2 个条目(200 bytes)
|
||||
assert_eq!(cache.len(), 2);
|
||||
|
||||
// 写入第 3 个 → 应淘汰最旧的(1.txt)
|
||||
cache.set(
|
||||
"/tmp/3.txt",
|
||||
FileState {
|
||||
content: "z".repeat(100),
|
||||
timestamp: 3000,
|
||||
offset: 1,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
// 旧条目被淘汰
|
||||
assert!(!cache.contains("/tmp/1.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_file_mtime() {
|
||||
// 创建临时文件
|
||||
let tmp = std::env::temp_dir().join("test_mtime.txt");
|
||||
std::fs::write(&tmp, "test").unwrap();
|
||||
|
||||
let mtime = get_file_mtime(&tmp.to_string_lossy());
|
||||
assert!(mtime.is_some());
|
||||
assert!(mtime.unwrap() > 0);
|
||||
|
||||
// 不存在的文件
|
||||
let mtime = get_file_mtime("/tmp/nonexistent_12345_xxx.txt");
|
||||
assert!(mtime.is_none());
|
||||
|
||||
std::fs::remove_file(&tmp).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_file_unchanged_stub_is_static() {
|
||||
assert!(FILE_UNCHANGED_STUB.contains("File unchanged"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// src/agent/runtime/finalize.rs
|
||||
//
|
||||
// 会话收尾:更新元信息、生成标题、保存指标、发送 Done 事件。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::info;
|
||||
|
||||
use super::{AgentMetrics, AgentStreamEvent};
|
||||
use crate::agent::hooks::{HookRegistry, SessionStopContext};
|
||||
use crate::agent::memory::extraction::{run_extraction, ExtractionConfig};
|
||||
use crate::agent::terminal::TurnTerminal;
|
||||
use crate::agent::trajectory::TrajectoryExporter;
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 完成会话回合:更新 session 元信息、触发 OnSessionStop hook、发送 Done、
|
||||
/// 导出 trajectory 数据。
|
||||
///
|
||||
/// `terminal` 参数允许传入实际的终止原因(取消/错误/超限等),
|
||||
/// 传入 `None` 时默认使用 `Completed`。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn finalize_turn(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
metrics: &AgentMetrics,
|
||||
question: &str,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
hook_registry: &HookRegistry,
|
||||
terminal: Option<TurnTerminal>,
|
||||
// Trajectory 导出所需参数
|
||||
library_dir: Option<&PathBuf>,
|
||||
llm_model: Option<&str>,
|
||||
system_prompt: Option<&str>,
|
||||
// 自动记忆提取所需
|
||||
app_state: Option<Arc<AppState>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let new_turn_count: i32 = sqlx::query_scalar(
|
||||
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
let metrics_json = serde_json::to_value(metrics).unwrap_or_default();
|
||||
info!(
|
||||
"[Finalize] 会话 {} 指标: steps={}, tools={:?}, compressions={}, duplicates={}",
|
||||
session_id,
|
||||
metrics.total_steps,
|
||||
metrics.tool_calls,
|
||||
metrics.compression_count,
|
||||
metrics.duplicate_detections
|
||||
);
|
||||
|
||||
// 首轮自动生成标题
|
||||
if new_turn_count <= 1 {
|
||||
let title = generate_title(question);
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET title = ?, turn_count = ?, metadata = ?, \
|
||||
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
|
||||
)
|
||||
.bind(&title)
|
||||
.bind(new_turn_count)
|
||||
.bind(&metrics_json)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
} else {
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET turn_count = ?, metadata = ?, \
|
||||
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
|
||||
)
|
||||
.bind(new_turn_count)
|
||||
.bind(&metrics_json)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
// OnSessionStop hook — 唯一触发点,使用传入的 terminal 或默认 Completed
|
||||
let default_terminal = TurnTerminal::Completed {
|
||||
session_id: session_id.to_string(),
|
||||
total_steps: metrics.total_steps,
|
||||
};
|
||||
let actual_terminal = terminal.unwrap_or(default_terminal);
|
||||
hook_registry
|
||||
.run_on_session_stop(&SessionStopContext {
|
||||
session_id: session_id.to_string(),
|
||||
terminal: &actual_terminal,
|
||||
total_steps: metrics.total_steps,
|
||||
})
|
||||
.await;
|
||||
|
||||
// ── Trajectory 导出(非阻塞,失败不影响主流程) ──
|
||||
if let (Some(lib_dir), Some(model), Some(sys_prompt)) = (library_dir, llm_model, system_prompt)
|
||||
{
|
||||
if let Err(e) = TrajectoryExporter::export(
|
||||
db,
|
||||
session_id,
|
||||
lib_dir,
|
||||
model,
|
||||
sys_prompt,
|
||||
metrics,
|
||||
Some(&actual_terminal),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("[Finalize] Trajectory 导出失败(非致命): {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// ── 自动记忆提取(fire-and-forget,不阻塞会话关闭) ──
|
||||
if let Some(app_state) = app_state {
|
||||
let extraction_config = ExtractionConfig::from_env();
|
||||
let session_id_owned = session_id.to_string();
|
||||
let mem_mgr = app_state.memory_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
run_extraction(app_state, session_id_owned, mem_mgr, extraction_config).await;
|
||||
});
|
||||
}
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 根据用户首条问题生成会话标题(截取前 50 字符)
|
||||
fn generate_title(question: &str) -> String {
|
||||
let chars: String = question.chars().take(50).collect();
|
||||
if question.len() > 50 {
|
||||
format!("{}...", chars)
|
||||
} else {
|
||||
chars
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,189 @@
|
||||
// src/agent/runtime/partitioner.rs
|
||||
//
|
||||
// 工具调用并发分区器。
|
||||
// 将准备好的工具调用按并发安全性分组成批处理。
|
||||
// 参考 Claude Code partitionToolCalls() 设计。
|
||||
//
|
||||
// 分区规则:
|
||||
// 1. 连续的并发安全工具放在同一个并行批次
|
||||
// 2. 非并发安全的工具单独一个批次(串行执行)
|
||||
// 3. 每个并行批次最多 max_concurrency 个工具
|
||||
|
||||
use tracing::debug;
|
||||
|
||||
use super::executor::PreparedCall;
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
|
||||
/// 一批工具调用
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolBatch {
|
||||
/// 该批次是否可以并行执行
|
||||
pub is_parallel: bool,
|
||||
/// 批次中的工具调用(按原始顺序)
|
||||
pub calls: Vec<PreparedCall>,
|
||||
}
|
||||
|
||||
/// 工具调用分区器
|
||||
pub struct ToolPartitioner {
|
||||
/// 最大并行度(并行批次中最多执行的工具数)
|
||||
max_concurrency: usize,
|
||||
}
|
||||
|
||||
impl ToolPartitioner {
|
||||
/// 创建分区器
|
||||
///
|
||||
/// `max_concurrency` 为 0 时使用默认值 10。
|
||||
pub fn new(max_concurrency: usize) -> Self {
|
||||
let concurrency = if max_concurrency == 0 {
|
||||
10
|
||||
} else {
|
||||
max_concurrency
|
||||
};
|
||||
ToolPartitioner {
|
||||
max_concurrency: concurrency,
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 prepared_calls 分区为顺序批处理。
|
||||
///
|
||||
/// 返回的批次列表按顺序依次执行:
|
||||
/// - `is_parallel = true` 的批次内工具可并发执行
|
||||
/// - `is_parallel = false` 的批次内只有一个工具,需串行执行
|
||||
pub fn partition(
|
||||
&self,
|
||||
prepared_calls: &[PreparedCall],
|
||||
tool_registry: &ToolRegistry,
|
||||
) -> Vec<ToolBatch> {
|
||||
let mut batches: Vec<ToolBatch> = Vec::new();
|
||||
|
||||
for prep in prepared_calls {
|
||||
let is_safe = self.is_concurrent_safe(prep, tool_registry);
|
||||
|
||||
// 尝试追加到上一个并行批次
|
||||
if is_safe {
|
||||
if let Some(last) = batches.last_mut() {
|
||||
if last.is_parallel && last.calls.len() < self.max_concurrency {
|
||||
last.calls.push(prep.clone());
|
||||
continue;
|
||||
}
|
||||
}
|
||||
// 新建并行批次
|
||||
batches.push(ToolBatch {
|
||||
is_parallel: true,
|
||||
calls: vec![prep.clone()],
|
||||
});
|
||||
} else {
|
||||
// 非并发安全,单独一个串行批次
|
||||
// 如果前一个也是串行,可以合并(但为了简单,保持每个非安全工具独立批次)
|
||||
batches.push(ToolBatch {
|
||||
is_parallel: false,
|
||||
calls: vec![prep.clone()],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
debug!(
|
||||
"[Partitioner] 分区完成: {} calls → {} batches",
|
||||
prepared_calls.len(),
|
||||
batches.len()
|
||||
);
|
||||
batches
|
||||
}
|
||||
|
||||
/// 判断单个工具调用是否可并发安全执行
|
||||
fn is_concurrent_safe(&self, prep: &PreparedCall, tool_registry: &ToolRegistry) -> bool {
|
||||
tool_registry
|
||||
.get(&prep.tool_name)
|
||||
.map(|t| t.is_concurrency_safe(&prep.args))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use super::*;
|
||||
use crate::agent::skills::SkillRegistry;
|
||||
|
||||
fn make_registry() -> ToolRegistry {
|
||||
ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
|
||||
"./skills",
|
||||
)))))
|
||||
}
|
||||
|
||||
fn make_prep(name: &str) -> PreparedCall {
|
||||
PreparedCall {
|
||||
tool_call_id: format!("call_{}", name),
|
||||
tool_name: name.to_string(),
|
||||
args: serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_concurrent_in_single_batch() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("rag_search"),
|
||||
make_prep("get_paper_metadata"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// search_papers, rag_search, get_paper_metadata 都是并发安全的
|
||||
// 它们应该在一个并行批次中
|
||||
assert_eq!(batches.len(), 1);
|
||||
assert!(batches[0].is_parallel);
|
||||
assert_eq!(batches[0].calls.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_download_splits_batch() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("download_paper"),
|
||||
make_prep("rag_search"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// search_papers (安全) → download_paper (不安全) → rag_search (安全)
|
||||
// 应分为 3 个批次
|
||||
assert_eq!(batches.len(), 3);
|
||||
assert!(batches[0].is_parallel); // search_papers
|
||||
assert_eq!(batches[0].calls.len(), 1);
|
||||
assert!(!batches[1].is_parallel); // download_paper (串行)
|
||||
assert_eq!(batches[1].calls.len(), 1);
|
||||
assert!(batches[2].is_parallel); // rag_search
|
||||
assert_eq!(batches[2].calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_max_concurrency_limit() {
|
||||
let registry = make_registry();
|
||||
let partitioner = ToolPartitioner::new(2);
|
||||
|
||||
let calls = vec![
|
||||
make_prep("search_papers"),
|
||||
make_prep("rag_search"),
|
||||
make_prep("get_paper_metadata"),
|
||||
make_prep("load_skill"),
|
||||
];
|
||||
|
||||
let batches = partitioner.partition(&calls, ®istry);
|
||||
|
||||
// 4 个并发安全调用,max_concurrency=2 → 应分为 2 个并行批次
|
||||
assert_eq!(batches.len(), 2);
|
||||
assert!(batches[0].is_parallel);
|
||||
assert_eq!(batches[0].calls.len(), 2);
|
||||
assert!(batches[1].is_parallel);
|
||||
assert_eq!(batches[1].calls.len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// src/agent/runtime/permission.rs
|
||||
//
|
||||
// 权限检查管道 — 优先级排序的规则链。
|
||||
// 参考 Claude Code PermissionChecker 设计。
|
||||
//
|
||||
// 规则优先级(从高到低):
|
||||
// 1. Deny — 不可覆盖的拒绝
|
||||
// 2. Allow — 允许
|
||||
// 3. Ask — 需要用户确认
|
||||
//
|
||||
// 支持通配符 "*" 匹配所有工具。
|
||||
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::PermissionRule;
|
||||
|
||||
/// 权限检查结果
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionResult {
|
||||
/// 被拒绝(不可覆盖)
|
||||
Denied { reason: String },
|
||||
/// 允许
|
||||
Allowed,
|
||||
/// 需要用户确认
|
||||
AskUser { message: String },
|
||||
}
|
||||
|
||||
impl PermissionResult {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, PermissionResult::Allowed)
|
||||
}
|
||||
|
||||
pub fn is_denied(&self) -> bool {
|
||||
matches!(self, PermissionResult::Denied { .. })
|
||||
}
|
||||
}
|
||||
|
||||
/// 权限检查器 — 维护有序规则列表并逐条匹配
|
||||
pub struct PermissionChecker {
|
||||
rules: Vec<PermissionRule>,
|
||||
}
|
||||
|
||||
impl PermissionChecker {
|
||||
/// 创建空的检查器(默认允许所有)
|
||||
pub fn new() -> Self {
|
||||
PermissionChecker { rules: Vec::new() }
|
||||
}
|
||||
|
||||
/// 添加规则。先添加的优先级更高。
|
||||
pub fn add_rule(&mut self, rule: PermissionRule) {
|
||||
self.rules.push(rule);
|
||||
}
|
||||
|
||||
/// 检查指定工具是否可以执行。
|
||||
///
|
||||
/// 遍历规则列表,返回第一个匹配的决策。
|
||||
/// 无匹配规则时默认 Allow。
|
||||
pub fn check(&self, tool_name: &str) -> PermissionResult {
|
||||
for rule in &self.rules {
|
||||
match rule {
|
||||
PermissionRule::Deny {
|
||||
tool_name: name,
|
||||
reason,
|
||||
} if Self::matches(name, tool_name) => {
|
||||
info!("[Permission] 拒绝工具 {}: {}", tool_name, reason);
|
||||
return PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
}
|
||||
PermissionRule::Allow { tool_name: name } if Self::matches(name, tool_name) => {
|
||||
return PermissionResult::Allowed;
|
||||
}
|
||||
PermissionRule::Ask {
|
||||
tool_name: name,
|
||||
message,
|
||||
} if Self::matches(name, tool_name) => {
|
||||
return PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// 默认允许
|
||||
PermissionResult::Allowed
|
||||
}
|
||||
|
||||
/// 检查是否有明确拒绝该工具的规则
|
||||
pub fn is_denied(&self, tool_name: &str) -> bool {
|
||||
self.check(tool_name).is_denied()
|
||||
}
|
||||
|
||||
/// 规则名称匹配:支持精确匹配和通配符 "*"
|
||||
fn matches(pattern: &str, tool_name: &str) -> bool {
|
||||
pattern == "*" || pattern == tool_name
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PermissionChecker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::tools::PermissionRule;
|
||||
|
||||
#[test]
|
||||
fn test_empty_checker_allows_all() {
|
||||
let checker = PermissionChecker::new();
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
assert_eq!(checker.check("download_paper"), PermissionResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_wins_over_allow() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "download_paper".into(),
|
||||
reason: "blocked".into(),
|
||||
});
|
||||
checker.add_rule(PermissionRule::Allow {
|
||||
tool_name: "download_paper".into(),
|
||||
});
|
||||
|
||||
let result = checker.check("download_paper");
|
||||
assert!(result.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_deny_blocks_all() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "*".into(),
|
||||
reason: "all blocked".into(),
|
||||
});
|
||||
|
||||
assert!(checker.check("search_papers").is_denied());
|
||||
assert!(checker.check("download_paper").is_denied());
|
||||
assert!(checker.is_denied("rag_search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_returns_ask_user() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Ask {
|
||||
tool_name: "delete_paper".into(),
|
||||
message: "Are you sure?".into(),
|
||||
});
|
||||
|
||||
let result = checker.check("delete_paper");
|
||||
assert_eq!(
|
||||
result,
|
||||
PermissionResult::AskUser {
|
||||
message: "Are you sure?".into()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_match_defaults_to_allow() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "download_paper".into(),
|
||||
reason: "blocked".into(),
|
||||
});
|
||||
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rule_ordering_first_match_wins() {
|
||||
let mut checker = PermissionChecker::new();
|
||||
// 先添加 Allow,后添加 Deny — Allow 先匹配
|
||||
checker.add_rule(PermissionRule::Allow {
|
||||
tool_name: "search_papers".into(),
|
||||
});
|
||||
checker.add_rule(PermissionRule::Deny {
|
||||
tool_name: "search_papers".into(),
|
||||
reason: "should not match".into(),
|
||||
});
|
||||
|
||||
assert_eq!(checker.check("search_papers"), PermissionResult::Allowed);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
// src/agent/runtime/session.rs
|
||||
//
|
||||
// 会话生命周期管理:创建/恢复/验证 Agent 会话。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::clients::llm::LlmClient;
|
||||
|
||||
/// 会话信息摘要
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionInfo {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
}
|
||||
|
||||
/// 创建新会话或恢复已有会话。
|
||||
///
|
||||
/// 返回会话信息。如果指定的 session_id 不存在则返回错误。
|
||||
pub async fn create_or_resume_session(
|
||||
db: &SqlitePool,
|
||||
session_id: Option<String>,
|
||||
llm: &LlmClient,
|
||||
) -> anyhow::Result<SessionInfo> {
|
||||
match session_id {
|
||||
Some(id) => {
|
||||
// 验证会话存在且未被软删除
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
|
||||
}
|
||||
|
||||
// 计算当前轮次号
|
||||
let turn_index: i32 = sqlx::query_scalar(
|
||||
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?",
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: id,
|
||||
turn_index,
|
||||
})
|
||||
}
|
||||
None => {
|
||||
let new_id = uuid::Uuid::new_v4().to_string();
|
||||
sqlx::query("INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)")
|
||||
.bind(&new_id)
|
||||
.bind("")
|
||||
.bind(llm.model())
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: new_id,
|
||||
turn_index: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 加载会话的历史消息(供 LLM 上下文使用)。
|
||||
///
|
||||
/// `agent_name` 参数用于消息隔离:
|
||||
/// - `"lead"` — 只加载 Lead Agent 自己的消息(默认)
|
||||
/// - `"*"` — 加载所有 agent 的消息(调试/审计用)
|
||||
pub async fn load_history_for_llm(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
load_history_for_agent(db, session_id, "lead").await
|
||||
}
|
||||
|
||||
/// 加载指定 agent 的历史消息。
|
||||
pub async fn load_history_for_agent(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
agent_name: &str,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
let rows: Vec<(
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
)> = if agent_name == "*" {
|
||||
sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? AND agent_name = ? ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(agent_name)
|
||||
.fetch_all(db)
|
||||
.await?
|
||||
};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (role_str, content, tool_calls_json, tool_call_id, thought) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"system" => MessageRole::System,
|
||||
"user" => MessageRole::User,
|
||||
"assistant" => MessageRole::Assistant,
|
||||
"tool" => MessageRole::Tool,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let tool_calls: Option<Vec<crate::clients::llm::ToolCall>> =
|
||||
tool_calls_json.and_then(|json_str| serde_json::from_str(&json_str).ok());
|
||||
|
||||
messages.push(ChatMessage {
|
||||
role,
|
||||
content: if content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(content)
|
||||
},
|
||||
tool_call_id,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content: thought,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
// src/agent/runtime/streaming.rs
|
||||
//
|
||||
// LLM 流式响应处理:消费 chat_stream 返回的 StreamEvent 通道,
|
||||
// 支持并发取消检测,累积推理内容、文本增量和工具调用。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::error;
|
||||
|
||||
use crate::clients::llm::{
|
||||
ChatMessage, LlmClient, StreamEvent, TokenUsage, ToolCall, ToolDefinition,
|
||||
};
|
||||
|
||||
use super::AgentStreamEvent;
|
||||
|
||||
/// 流式处理的结果
|
||||
#[derive(Debug)]
|
||||
pub struct StreamOutput {
|
||||
pub content: String,
|
||||
pub reasoning: Option<String>,
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
pub is_tool_call_step: bool,
|
||||
pub status: StreamStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum StreamStatus {
|
||||
Success,
|
||||
Error(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// 处理 LLM 流式响应。
|
||||
///
|
||||
/// 使用 tokio::select! 在流式读取和取消信号之间竞速。
|
||||
/// 实时发送 Thought/TextDelta SSE 事件给前端。
|
||||
pub async fn process_llm_stream(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
tool_defs: &[ToolDefinition],
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
step: usize,
|
||||
session_id: &str,
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
) -> StreamOutput {
|
||||
// 1. 发起 LLM 流式调用
|
||||
let mut stream_rx = match llm.chat_stream(messages, tool_defs).await {
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[Streaming] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
return StreamOutput {
|
||||
content: String::new(),
|
||||
reasoning: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
is_tool_call_step: false,
|
||||
status: StreamStatus::Error(e.to_string()),
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
let mut accumulated_content = String::new();
|
||||
let mut accumulated_reasoning = String::new();
|
||||
let mut accumulated_tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut usage: Option<TokenUsage> = None;
|
||||
let mut is_tool_call_step = false;
|
||||
|
||||
// 2. 取消监视 future
|
||||
let sid = session_id.to_string();
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 3. tokio::select! 竞速:流式事件 vs 取消信号
|
||||
let mut cancel_pinned = Box::pin(cancel_fut);
|
||||
let mut error_msg = None;
|
||||
let mut was_cancelled = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
event_opt = stream_rx.recv() => {
|
||||
match event_opt {
|
||||
Some(event) => {
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
is_tool_call_step = true;
|
||||
accumulated_tool_calls = Some(tool_calls);
|
||||
}
|
||||
StreamEvent::ToolCallDelta { .. } => {}
|
||||
StreamEvent::Usage(u) => {
|
||||
usage = Some(u);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
break;
|
||||
}
|
||||
StreamEvent::Error(e) => {
|
||||
error_msg = Some(e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = &mut cancel_pinned => {
|
||||
was_cancelled = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 构建结果
|
||||
let status = if was_cancelled {
|
||||
StreamStatus::Cancelled
|
||||
} else if let Some(e) = error_msg {
|
||||
StreamStatus::Error(e)
|
||||
} else {
|
||||
StreamStatus::Success
|
||||
};
|
||||
|
||||
let reasoning = if accumulated_reasoning.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_reasoning)
|
||||
};
|
||||
|
||||
StreamOutput {
|
||||
content: accumulated_content,
|
||||
reasoning,
|
||||
tool_calls: accumulated_tool_calls,
|
||||
usage,
|
||||
is_tool_call_step,
|
||||
status,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
// src/agent/runtime/streaming_executor.rs
|
||||
//
|
||||
// 流式工具执行器。
|
||||
// 参考 Claude Code StreamingToolExecutor 设计。
|
||||
//
|
||||
// 当 LLM 流式输出 tool_use 块时,立即开始执行并发安全的工具。
|
||||
// 非并发安全的工具排队等待。结果按流中顺序 yield。
|
||||
//
|
||||
// 功能:
|
||||
// 1. 流式执行 — tool_use 到达时立即调度
|
||||
// 2. Sibling Abort — 副效应工具报错时中止兄弟并行执行
|
||||
// 3. Progress 流式 — 长时间操作可发送进度更新
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{broadcast, mpsc, oneshot};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::partitioner::ToolPartitioner;
|
||||
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 流式工具执行状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TrackedToolStatus {
|
||||
/// 工具调用已从 LLM 流中接收到
|
||||
Queued,
|
||||
/// 正在执行中
|
||||
Executing,
|
||||
/// 执行完成,等待 yield
|
||||
Completed,
|
||||
/// 结果已 yield 给调用方
|
||||
Yielded,
|
||||
}
|
||||
|
||||
/// 跟踪中的工具执行
|
||||
#[derive(Debug)]
|
||||
struct TrackedTool {
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
status: TrackedToolStatus,
|
||||
/// 执行完成后的输出
|
||||
output: Option<ToolOutput>,
|
||||
/// 取消通道(Sibling Abort 使用)
|
||||
#[allow(dead_code)]
|
||||
cancel_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
/// Sibling Abort 原因
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AbortReason {
|
||||
/// 兄弟工具出错触发的级联取消
|
||||
SiblingError { description: String },
|
||||
/// 用户主动中断
|
||||
UserInterrupted,
|
||||
}
|
||||
|
||||
/// 流式工具执行器
|
||||
pub struct StreamingToolExecutor {
|
||||
/// 所有跟踪中的工具
|
||||
tracked: Vec<TrackedTool>,
|
||||
/// 工具注册表
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
/// 并发分区器(保留用于未来并发策略优化)
|
||||
#[allow(dead_code)]
|
||||
partitioner: ToolPartitioner,
|
||||
/// 工具上下文
|
||||
tool_context: ToolContext,
|
||||
/// Sibling Abort 广播通道 (tx)
|
||||
abort_tx: broadcast::Sender<AbortReason>,
|
||||
/// Sibling Abort 广播通道 (rx)
|
||||
abort_rx: broadcast::Receiver<AbortReason>,
|
||||
/// 当前是否已发生错误(触发 sibling abort)
|
||||
has_errored: bool,
|
||||
/// 出错工具的描述
|
||||
errored_tool_desc: String,
|
||||
/// 下一个 stream_index
|
||||
next_index: usize,
|
||||
/// 最大工具输出字符数
|
||||
max_output_chars: usize,
|
||||
}
|
||||
|
||||
impl StreamingToolExecutor {
|
||||
/// 创建新的流式执行器
|
||||
pub fn new(
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
tool_context: ToolContext,
|
||||
max_concurrency: usize,
|
||||
max_output_chars: usize,
|
||||
) -> Self {
|
||||
let (abort_tx, abort_rx) = broadcast::channel(16);
|
||||
StreamingToolExecutor {
|
||||
tracked: Vec::new(),
|
||||
tool_registry,
|
||||
partitioner: ToolPartitioner::new(max_concurrency),
|
||||
tool_context,
|
||||
abort_tx,
|
||||
abort_rx,
|
||||
has_errored: false,
|
||||
errored_tool_desc: String::new(),
|
||||
next_index: 0,
|
||||
max_output_chars,
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取 abort 广播发送端(供外部注入取消信号)
|
||||
pub fn abort_sender(&self) -> broadcast::Sender<AbortReason> {
|
||||
self.abort_tx.clone()
|
||||
}
|
||||
|
||||
/// 当 LLM 流产生一个新的 tool_use 时调用。
|
||||
///
|
||||
/// 返回 true 表示该工具已立即开始执行(并发安全),false 表示排队。
|
||||
pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool {
|
||||
let _index = self.next_index;
|
||||
self.next_index += 1;
|
||||
|
||||
let is_concurrency_safe = self
|
||||
.tool_registry
|
||||
.get(&name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let (cancel_tx, _cancel_rx) = oneshot::channel();
|
||||
|
||||
let tool = TrackedTool {
|
||||
tool_call_id: call_id.clone(),
|
||||
tool_name: name.clone(),
|
||||
args: args.clone(),
|
||||
status: TrackedToolStatus::Queued,
|
||||
output: None,
|
||||
cancel_tx: Some(cancel_tx),
|
||||
};
|
||||
|
||||
self.tracked.push(tool);
|
||||
|
||||
if is_concurrency_safe {
|
||||
info!("[StreamingExecutor] 立即调度并发安全工具: {}", name);
|
||||
self.try_execute_pending();
|
||||
true
|
||||
} else {
|
||||
info!("[StreamingExecutor] 排队非并发安全工具: {}", name);
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM 流结束后调用,执行所有剩余排队工具。
|
||||
pub async fn flush(&mut self) {
|
||||
info!(
|
||||
"[StreamingExecutor] flush: {} tracked, {} queued",
|
||||
self.tracked.len(),
|
||||
self.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Queued)
|
||||
.count()
|
||||
);
|
||||
|
||||
// 将剩余排队的工具分批执行
|
||||
let queued: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in queued {
|
||||
self.execute_one(idx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 按流顺序获取下一个完成的结果(非阻塞)。
|
||||
pub fn next_result(&mut self) -> Option<(String, ToolOutput)> {
|
||||
for tool in &mut self.tracked {
|
||||
if tool.status == TrackedToolStatus::Completed {
|
||||
tool.status = TrackedToolStatus::Yielded;
|
||||
let output = tool
|
||||
.output
|
||||
.take()
|
||||
.unwrap_or_else(|| ToolOutput::error("工具执行异常:无输出"));
|
||||
return Some((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 是否有未 yield 的结果
|
||||
pub fn has_pending_results(&self) -> bool {
|
||||
self.tracked
|
||||
.iter()
|
||||
.any(|t| t.status == TrackedToolStatus::Completed)
|
||||
}
|
||||
|
||||
/// 是否有未完成的工具
|
||||
pub fn has_unfinished(&self) -> bool {
|
||||
self.tracked.iter().any(|t| {
|
||||
t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing
|
||||
})
|
||||
}
|
||||
|
||||
/// 获取所有已完成的结果(包括已 yield 和未 yield 的)
|
||||
pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> {
|
||||
let mut results = Vec::new();
|
||||
for tool in &mut self.tracked {
|
||||
if let Some(output) = tool.output.take() {
|
||||
results.push((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
// ── 内部方法 ──
|
||||
|
||||
/// 尝试执行可执行的排队工具
|
||||
fn try_execute_pending(&mut self) {
|
||||
// 简单策略:如果有正在执行的且它不是并发的,则不启动新的
|
||||
let has_executing = self
|
||||
.tracked
|
||||
.iter()
|
||||
.any(|t| t.status == TrackedToolStatus::Executing);
|
||||
|
||||
if !has_executing {
|
||||
// 启动所有排队的并发安全工具
|
||||
let indices: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in indices {
|
||||
// 在同步上下文中只能标记状态,实际执行在 async 上下文中
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行单个工具(内部辅助)
|
||||
async fn execute_one(&mut self, idx: usize) {
|
||||
if idx >= self.tracked.len() {
|
||||
return;
|
||||
}
|
||||
|
||||
// 检查 sibling abort
|
||||
if self.has_errored {
|
||||
if let Ok(reason) = self.abort_rx.try_recv() {
|
||||
let msg = match reason {
|
||||
AbortReason::SiblingError { ref description } => {
|
||||
format!("取消:并行工具 {} 出错,已级联取消", description)
|
||||
}
|
||||
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
|
||||
};
|
||||
self.tracked[idx].output = Some(ToolOutput::error(msg));
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
let tool_name = self.tracked[idx].tool_name.clone();
|
||||
|
||||
let output = match self.tool_registry.get(&tool_name) {
|
||||
Some(tool) => {
|
||||
let (progress_tx, _progress_rx) = mpsc::unbounded_channel();
|
||||
|
||||
let tool_args = self.tracked[idx].args.clone();
|
||||
let tool_fut =
|
||||
tool.execute_with_progress(tool_args, &self.tool_context, Some(&progress_tx));
|
||||
|
||||
tool_fut.await
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
// 检查是否需要触发 sibling abort
|
||||
if output.is_error {
|
||||
let causes_abort = self
|
||||
.tool_registry
|
||||
.get(&tool_name)
|
||||
.map(|t| t.causes_sibling_abort())
|
||||
.unwrap_or(false);
|
||||
|
||||
if causes_abort {
|
||||
warn!(
|
||||
"[StreamingExecutor] 工具 {} 出错,触发 sibling abort",
|
||||
tool_name
|
||||
);
|
||||
self.has_errored = true;
|
||||
self.errored_tool_desc = tool_name.clone();
|
||||
let _ = self.abort_tx.send(AbortReason::SiblingError {
|
||||
description: tool_name.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 截断输出
|
||||
let truncated = if output.content.len() > self.max_output_chars {
|
||||
let t: String = output.content.chars().take(self.max_output_chars).collect();
|
||||
ToolOutput {
|
||||
content: format!(
|
||||
"{}...\n[输出已截断,原始长度: {} 字符]",
|
||||
t,
|
||||
output.content.len()
|
||||
),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata,
|
||||
}
|
||||
} else {
|
||||
output
|
||||
};
|
||||
|
||||
self.tracked[idx].output = Some(truncated);
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
// src/agent/runtime/system_prompt.rs
|
||||
//
|
||||
// 模块化系统提示词组装 — 参考 Claude Code s10 System Prompt 设计。
|
||||
//
|
||||
// 将硬编码的提示词拆分为独立 section,运行时按需拼接。
|
||||
// 静态 section 在前以最大化 Anthropic prompt cache 命中率。
|
||||
|
||||
/// 系统提示词组装器
|
||||
pub struct SystemPrompt {
|
||||
sections: Vec<(&'static str, String)>,
|
||||
}
|
||||
|
||||
impl SystemPrompt {
|
||||
pub fn new() -> Self {
|
||||
SystemPrompt {
|
||||
sections: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 添加一个 section(先添加的排在前面)
|
||||
pub fn add_section(&mut self, name: &'static str, content: String) {
|
||||
self.sections.push((name, content));
|
||||
}
|
||||
|
||||
/// 组装最终的系统提示词(section 间用双换行分隔)
|
||||
pub fn assemble(&self) -> String {
|
||||
self.sections
|
||||
.iter()
|
||||
.map(|(_, content)| content.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
}
|
||||
|
||||
/// 是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.sections.is_empty()
|
||||
}
|
||||
|
||||
/// 获取 section 数量
|
||||
pub fn section_count(&self) -> usize {
|
||||
self.sections.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SystemPrompt {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 静态身份 section(始终加载,最大化 prompt cache 命中率)
|
||||
pub const IDENTITY_SECTION: &str = "\
|
||||
你是一位专业的天体物理学研究助手,具备丰富的天文学知识。";
|
||||
|
||||
/// 静态核心原则 section
|
||||
pub const PRINCIPLES_SECTION: &str = "\
|
||||
核心原则:
|
||||
1. 主动使用工具搜索最新文献,不要仅凭训练数据回答。
|
||||
2. 优先使用本地资源(get_paper_content / rag_search),必要时再检索新文献。
|
||||
3. 收集到足够信息后立即给出最终答案,避免无意义的重复工具调用。
|
||||
4. 回答时引用具体文献来源,使用 ADS bibcode 标注。
|
||||
5. 对于数学公式,使用标准 LaTeX 格式。
|
||||
6. 用中文回答,保持科学术语的准确性(可附带英文原文)。
|
||||
7. 对于复杂任务(如文献综述),调用 load_skill 获取方法论指引,再用 todo_write 制定计划。
|
||||
8. 如果某个工具调用失败,不要用相同参数重试,尝试换一种方式或工具。
|
||||
9. 任务状态会在每轮开始时从数据库恢复,请基于最新状态继续工作。";
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_assemble_empty() {
|
||||
let sp = SystemPrompt::new();
|
||||
assert_eq!(sp.assemble(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_assemble_multiple_sections() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
sp.add_section("a", "Section A".to_string());
|
||||
sp.add_section("b", "Section B".to_string());
|
||||
let result = sp.assemble();
|
||||
assert_eq!(result, "Section A\n\nSection B");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_static_sections_first() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
sp.add_section("identity", "I am".to_string());
|
||||
sp.add_section("dynamic", "Tools: ...".to_string());
|
||||
let result = sp.assemble();
|
||||
assert!(result.starts_with("I am"));
|
||||
assert!(result.contains("Tools: ..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_section_count() {
|
||||
let mut sp = SystemPrompt::new();
|
||||
assert_eq!(sp.section_count(), 0);
|
||||
sp.add_section("a", "A".to_string());
|
||||
assert_eq!(sp.section_count(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
// src/agent/runtime/token_budget.rs
|
||||
//
|
||||
// Token 预算管理。
|
||||
// 参考 Claude Code TokenBudget 设计。
|
||||
// 软限制:接近上限时注入 nudging 消息提醒模型。
|
||||
// 硬限制:达到上限时触发强制压缩或终止。
|
||||
|
||||
/// Token 预算管理器。
|
||||
///
|
||||
/// 参考 Claude Code TokenBudget 设计,增加:
|
||||
/// - 多级渐进式 nudge(near_soft / over_soft / over_hard)
|
||||
/// - Diminishing returns 检测(防止模型在死循环中消耗预算)
|
||||
/// - Continuation 计数追踪
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TokenBudget {
|
||||
/// 软限制(触发 nudging 提醒)
|
||||
pub soft_limit: usize,
|
||||
/// 硬限制(触发强制动作)
|
||||
pub hard_limit: usize,
|
||||
/// 已消耗输入 tokens
|
||||
pub input_tokens_spent: usize,
|
||||
/// 已消耗输出 tokens
|
||||
pub output_tokens_spent: usize,
|
||||
/// 延续次数(每个 ReAct 步骤递增)
|
||||
pub continuation_count: usize,
|
||||
/// 上次检查时的总消耗(用于 diminishing returns 检测)
|
||||
last_total_spent: usize,
|
||||
/// 连续无进展次数
|
||||
consecutive_no_progress: usize,
|
||||
/// 是否已触发 diminishing returns
|
||||
pub diminishing_returns: bool,
|
||||
}
|
||||
|
||||
impl TokenBudget {
|
||||
/// 创建预算管理器
|
||||
pub fn new(soft_limit: usize, hard_limit: usize) -> Self {
|
||||
TokenBudget {
|
||||
soft_limit,
|
||||
hard_limit,
|
||||
input_tokens_spent: 0,
|
||||
output_tokens_spent: 0,
|
||||
continuation_count: 0,
|
||||
last_total_spent: 0,
|
||||
consecutive_no_progress: 0,
|
||||
diminishing_returns: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录输入 token 消耗
|
||||
pub fn spend_input(&mut self, tokens: usize) {
|
||||
self.input_tokens_spent += tokens;
|
||||
}
|
||||
|
||||
/// 记录输出 token 消耗
|
||||
pub fn spend_output(&mut self, tokens: usize) {
|
||||
self.output_tokens_spent += tokens;
|
||||
}
|
||||
|
||||
/// 总消耗
|
||||
pub fn total_spent(&self) -> usize {
|
||||
self.input_tokens_spent + self.output_tokens_spent
|
||||
}
|
||||
|
||||
/// 记录一次延续(每个 ReAct 步骤调用一次)。
|
||||
/// 同时检查 diminishing returns。
|
||||
pub fn record_continuation(&mut self) -> bool {
|
||||
self.continuation_count += 1;
|
||||
self.check_diminishing_returns_inner()
|
||||
}
|
||||
|
||||
/// 检测 diminishing returns — 模型在同一问题上打转而不产生实质进展。
|
||||
///
|
||||
/// 触发条件:3 次以上延续,且连续 2 次检查的 token 增量 < 500。
|
||||
/// 返回 true 表示已检测到无进展循环。
|
||||
pub fn check_diminishing_returns(&mut self) -> bool {
|
||||
self.check_diminishing_returns_inner()
|
||||
}
|
||||
|
||||
fn check_diminishing_returns_inner(&mut self) -> bool {
|
||||
if self.diminishing_returns {
|
||||
return true; // 已触发过,保持状态
|
||||
}
|
||||
if self.continuation_count < 3 {
|
||||
return false;
|
||||
}
|
||||
let delta = self.total_spent().saturating_sub(self.last_total_spent);
|
||||
self.last_total_spent = self.total_spent();
|
||||
if delta < 500 {
|
||||
self.consecutive_no_progress += 1;
|
||||
if self.consecutive_no_progress >= 2 {
|
||||
self.diminishing_returns = true;
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.consecutive_no_progress = 0;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 是否接近软限制(超过 80%)
|
||||
pub fn near_soft_limit(&self) -> bool {
|
||||
if self.soft_limit == 0 {
|
||||
return false;
|
||||
}
|
||||
self.total_spent() >= self.soft_limit * 8 / 10
|
||||
}
|
||||
|
||||
/// 是否超过软限制
|
||||
pub fn over_soft_limit(&self) -> bool {
|
||||
self.total_spent() >= self.soft_limit
|
||||
}
|
||||
|
||||
/// 是否超过硬限制
|
||||
pub fn over_hard_limit(&self) -> bool {
|
||||
self.total_spent() >= self.hard_limit
|
||||
}
|
||||
|
||||
/// 已使用预算的百分比(相对于软限制)
|
||||
pub fn usage_pct(&self) -> u32 {
|
||||
if self.soft_limit == 0 {
|
||||
return 0;
|
||||
}
|
||||
(self.total_spent() * 100 / self.soft_limit) as u32
|
||||
}
|
||||
|
||||
/// 剩余可用 tokens(硬限制 - 已消耗)
|
||||
pub fn remaining(&self) -> usize {
|
||||
self.hard_limit.saturating_sub(self.total_spent())
|
||||
}
|
||||
|
||||
/// 生成渐进式 nudging 提醒消息。
|
||||
///
|
||||
/// 三级:
|
||||
/// - near_soft (80-99%): 温和提醒
|
||||
/// - over_soft (100-硬): 明确警告
|
||||
/// - over_hard (>硬限制): 强制完成
|
||||
/// - diminishing_returns: 要求最终答案
|
||||
pub fn nudge_message(&self) -> Option<String> {
|
||||
if self.diminishing_returns {
|
||||
return Some(
|
||||
"⚠️ 已检测到重复操作模式 — 后续步骤未产生新信息。\
|
||||
请基于已收集的全部信息直接给出最终答案,不要再调用工具。"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
if self.over_hard_limit() {
|
||||
Some(format!(
|
||||
"🔴 Token 预算已耗尽({}/{} tokens, {}%)。\
|
||||
请立即总结当前发现并给出最终答案,不要再调用任何工具。",
|
||||
self.total_spent(),
|
||||
self.hard_limit,
|
||||
self.usage_pct()
|
||||
))
|
||||
} else if self.over_soft_limit() {
|
||||
let pct = self.usage_pct();
|
||||
Some(format!(
|
||||
"🟡 Token 预算警告:已使用 {}/{} tokens ({}%)。\
|
||||
请尽快总结关键发现并给出最终答案。如非必要,不要再调用工具。",
|
||||
self.total_spent(),
|
||||
self.soft_limit,
|
||||
pct
|
||||
))
|
||||
} else if self.near_soft_limit() {
|
||||
let pct = self.usage_pct();
|
||||
Some(format!(
|
||||
"💡 Token 预算提示:已使用 {}/{} tokens ({}%)。\
|
||||
请注意控制后续步骤的深度,优先处理最重要的发现。",
|
||||
self.total_spent(),
|
||||
self.soft_limit,
|
||||
pct
|
||||
))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 将硬限制提升到指定值(用于 error recovery 中的 escalate 步骤)
|
||||
pub fn escalate_hard_limit(&mut self, new_limit: usize) {
|
||||
self.hard_limit = new_limit;
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TokenBudget {
|
||||
fn default() -> Self {
|
||||
Self::new(32_000, 40_000)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_spend_tracking() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(500);
|
||||
budget.spend_output(300);
|
||||
assert_eq!(budget.total_spent(), 800);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_near_soft_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
assert!(!budget.near_soft_limit());
|
||||
budget.spend_input(850); // 85% > 80%
|
||||
assert!(budget.near_soft_limit());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_over_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(2100);
|
||||
assert!(budget.over_hard_limit());
|
||||
assert_eq!(budget.remaining(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nudge_message_at_soft_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(1000); // exactly at soft limit
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("Token 预算警告"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nudge_message_at_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(2000); // at hard limit
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("已耗尽"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_nudge_when_under_limit() {
|
||||
let budget = TokenBudget::new(1000, 2000);
|
||||
assert!(budget.nudge_message().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_escalate_hard_limit() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.escalate_hard_limit(64000);
|
||||
assert_eq!(budget.hard_limit, 64000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_budget() {
|
||||
let budget = TokenBudget::default();
|
||||
assert_eq!(budget.soft_limit, 32_000);
|
||||
assert_eq!(budget.hard_limit, 40_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_near_soft_limit_nudge() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(850); // 85% — near soft
|
||||
let msg = budget.nudge_message();
|
||||
assert!(msg.is_some());
|
||||
assert!(msg.unwrap().contains("Token 预算提示"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_not_triggered_early() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
// < 3 continuations — should not trigger
|
||||
budget.record_continuation();
|
||||
assert!(!budget.diminishing_returns);
|
||||
budget.record_continuation();
|
||||
assert!(!budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_triggers_after_stagnation() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
// First 3 continuations establish baseline (all with 0 spending)
|
||||
// After the 3rd, consecutive_no_progress becomes 1 (delta=0 < 500)
|
||||
for _ in 0..3 {
|
||||
budget.record_continuation();
|
||||
}
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
// 4th continuation with tiny spending — 2nd consecutive <500
|
||||
budget.spend_input(100);
|
||||
budget.record_continuation();
|
||||
// consecutive_no_progress is now 2 → triggered
|
||||
assert!(budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_resets_on_progress() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
// Establish baseline with spending BEFORE the 3-continuation threshold
|
||||
budget.spend_input(2000);
|
||||
for _ in 0..3 {
|
||||
budget.record_continuation();
|
||||
}
|
||||
// delta = total - last_total. After the first check, last_total is set.
|
||||
// total_spent=2000, last_total=2000 after first continuation ≥ 3
|
||||
|
||||
budget.spend_input(100); // total=2100
|
||||
budget.record_continuation(); // delta=100 < 500, cons=1
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
budget.spend_input(600); // total=2700
|
||||
budget.record_continuation(); // delta=600 >= 500, cons resets to 0
|
||||
assert!(!budget.diminishing_returns);
|
||||
|
||||
budget.spend_input(100); // total=2800
|
||||
budget.record_continuation(); // delta=100 < 500, cons=1
|
||||
assert!(!budget.diminishing_returns);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diminishing_returns_nudge_message() {
|
||||
let mut budget = TokenBudget::new(1000, 5000);
|
||||
budget.diminishing_returns = true;
|
||||
let msg = budget.nudge_message().unwrap();
|
||||
assert!(msg.contains("重复操作"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_pct() {
|
||||
let mut budget = TokenBudget::new(1000, 2000);
|
||||
budget.spend_input(500);
|
||||
assert_eq!(budget.usage_pct(), 50);
|
||||
budget.spend_output(300);
|
||||
assert_eq!(budget.usage_pct(), 80);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user