feat: 科研分析层全栈落地——光谱/时域/运动学分析工具链 + JWST/X 射线数据源 + 定时文献同步
数据分析层（新增 services/{spectrum,timeseries,analysis}）:
- 光谱参数提取 parameters.rs：LAMOST/SDSS/APOGEE/DESI FITS header 跨源归一化读取
Teff/logg/[Fe/H]/RV 及 ASPCAP 20+ 元素丰度，rayon 并发批量提取
- 谱线测量 lines.rs：内置真空/空气波长谱线表，窗口内极值搜索 + 梯形法积分 EW + FWHM，支持自定义谱线
- 交叉相关测速 cross_correlate.rs：对数波长重采样对齐，内置 Pickles 模板按光谱型插值，
CCF 峰值位置提取 RV 及不确定度
- 周期搜索 periodicity.rs：Lomb-Scargle 周期图（含 FAP 误报概率）+ BLS 凌星检测 + 相位折叠
- 变星分类 classification.rs：振幅/偏度/峰度/过零率/eta 等统计特征 + 规则分类（RR Lyrae/Cepheid/食双星/AGN 等）
- SED 拟合 sed.rs：多波段测光黑体模型拟合，输出 T_eff/半径/消光 A_V/光度及不确定度
- 运动学 kinematics.rs：视差+自行+RV → 银河系 UVW 空间速度，含移动星群成员概率（Banyan Σ 简化版）
- 化学丰度 chemistry.rs：[α/Fe] vs [Fe/H] 计算，厚盘/薄盘/晕星族判别
- 观测规划 observability.rs：目标升落时间/airmass/月相影响/曝光时间估算
- 赫罗图 hr_diagram.rs：Gaia TAP CMD 查询，新增 GET /api/analysis/hr-diagram 端点
数据获取层:
- JWST：clients/mast/jwst.rs 封装 MAST Portal 锥形检索 + JwstSpectrumFetcher（NIRSpec/MIRI 光谱）
- X 射线：clients/heasarc 封装 HEASARC TAP（ADQL）+ XMM-Newton/Chandra 光谱 fetcher
- 图像 cutout：SDSS SkyServer/STScI DSS/Pan-STARRS 三源 cutout + 发现图（Finding Chart）生成
- Source 枚举新增 Jwst/Xmm/Chandra 并注册 ObservationRegistry，前端 SOURCE_THEME 与筛选器同步三源
Agent 工具集（24→35）:
- 新增 9 个分析工具：get_spectrum_parameters / measure_spectral_lines / measure_radial_velocity /
find_period / classify_variable_star / fit_sed / analyze_kinematics / analyze_abundance_pattern / plan_observation
- batch_process：批量样本"查询→下载→分析→报告"流水线，并发控制防数据源速率限制
- literature_monitor：按 ADS 查询式/时间窗/最低引用数检查最新文献
定时文献同步:
- sync_queries 表新增 is_scheduled 列（migration 20260713）
- 新增 POST /sync/queries/:id/schedule 端点
- 服务启动时拉起每小时调度器，对 is_scheduled=1 的检索配置静默执行 ADS（entdate 增量）/arXiv 增量同步
- search_history 工具收敛至 services/search::search_agent_history，消除 FTS 查询逻辑重复
其他:
- plotting skill 由占位填充为完整科研绘图规范：光谱/光变/折叠曲线/CMD/SED/[α/Fe]/周期图/Mollweide/发现图 9 类 matplotlib 模板
- 删除死代码 streaming_executor.rs（929 行，仅剩 mod 声明引用，无调用方）
- 新增 docs/roadmap-research-features.md 科研功能路线图及实现状态
This commit is contained in:
@@ -70,8 +70,12 @@ const DEFAULT_EXCLUDES: &[&str] = &[
|
||||
"Thumbs.db",
|
||||
];
|
||||
|
||||
/// 每个 turn 最多快照一次的工具
|
||||
const CHECKPOINT_TRIGGER_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"];
|
||||
/// 每个 turn 最多快照一次的工具。
|
||||
///
|
||||
/// 历史上执行器按此名单触发快照；现在由各工具通过
|
||||
/// `AgentTool::causes_file_changes()` 声明，本名单仅作为
|
||||
/// 注册表查询不到工具时的保守回退。
|
||||
const CHECKPOINT_FALLBACK_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"];
|
||||
|
||||
/// Checkpoint 元数据
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -465,9 +469,9 @@ impl CheckpointManager {
|
||||
))
|
||||
}
|
||||
|
||||
/// 检查指定工具是否需要触发 checkpoint。
|
||||
/// 检查指定工具是否需要触发 checkpoint（注册表查询不到时的保守回退）。
|
||||
pub fn should_checkpoint(tool_name: &str) -> bool {
|
||||
CHECKPOINT_TRIGGER_TOOLS.contains(&tool_name)
|
||||
CHECKPOINT_FALLBACK_TOOLS.contains(&tool_name)
|
||||
}
|
||||
|
||||
/// 获取 repo 路径
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
//
|
||||
// 上下文构建：加载历史消息、注入系统提示词、添加用户消息、
|
||||
// 从数据库恢复持久化的任务状态。
|
||||
//
|
||||
// 上下文快照回放：若上一 turn 结束时保存过压缩后的折叠上下文
|
||||
// （context_snapshot 事件），则加载"快照消息 + id > base 的增量消息"，
|
||||
// 避免从原始消息重建后再次触发 LLM 摘要压缩（重复付费且信息有损）。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::info;
|
||||
@@ -9,22 +13,51 @@ use tracing::info;
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
use super::session;
|
||||
use super::session_events;
|
||||
|
||||
/// 构建初始 LLM 上下文：加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。
|
||||
/// 构建初始 LLM 上下文。
|
||||
///
|
||||
/// 加载顺序：上下文快照（若有）+ 增量历史 → 插入系统提示词 →
|
||||
/// 动态上下文快照（内容变更时才追加并持久化）→ 用户消息 → 任务恢复。
|
||||
pub async fn build_initial_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
system_prompt: &str,
|
||||
question: &str,
|
||||
_turn_index: i32,
|
||||
turn_index: i32,
|
||||
dynamic_context: Option<(String, u64)>,
|
||||
image_context: Option<&str>,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let mut messages = session::load_history_for_llm(db, session_id).await?;
|
||||
let mut messages = load_folded_history(db, session_id).await?;
|
||||
|
||||
// 注入系统提示词（如果历史中没有）
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(system_prompt));
|
||||
}
|
||||
|
||||
// 动态上下文快照（durable user-role 消息）：哈希与上次持久化的一致时跳过。
|
||||
// 这是 KV-cache 纪律的另一半——易变内容不进 system prompt（否则任一
|
||||
// 变化都使整条前缀缓存失效），而以追加式快照进入历史，字节一旦写入
|
||||
// 就永不变更。
|
||||
if let Some((text, hash)) = dynamic_context {
|
||||
let last_hash = last_dynamic_context_hash(db, session_id).await;
|
||||
if last_hash.as_deref() != Some(hash.to_string().as_str()) {
|
||||
let reminder =
|
||||
ChatMessage::user(format!("<system-reminder>\n{}\n</system-reminder>", text));
|
||||
// 先持久化再注入（模型可见 ⟺ 已日志化）
|
||||
persist_dynamic_context(db, session_id, turn_index, &reminder, hash).await;
|
||||
messages.push(reminder);
|
||||
}
|
||||
}
|
||||
|
||||
// 图片上下文：以 system-reminder 形式注入在用户问题之前
|
||||
if let Some(img_ctx) = image_context {
|
||||
messages.push(ChatMessage::user(format!(
|
||||
"<system-reminder>\n{}\n</system-reminder>",
|
||||
img_ctx
|
||||
)));
|
||||
}
|
||||
|
||||
// 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
|
||||
@@ -36,6 +69,29 @@ pub async fn build_initial_context(
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 加载折叠后的历史消息：优先回放上下文快照 + 增量；无快照时全量加载。
|
||||
async fn load_folded_history(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
if let Some((snapshot_messages, base_id)) =
|
||||
session_events::load_context_snapshot(db, session_id).await
|
||||
{
|
||||
let incremental = session_events::load_messages_after(db, session_id, base_id).await?;
|
||||
info!(
|
||||
"[Context] 回放上下文快照: {} 条快照消息 + {} 条增量消息 (base_id={})",
|
||||
snapshot_messages.len(),
|
||||
incremental.len(),
|
||||
base_id
|
||||
);
|
||||
let mut messages = snapshot_messages;
|
||||
messages.extend(incremental);
|
||||
Ok(messages)
|
||||
} else {
|
||||
session::load_history_for_llm(db, session_id).await
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 agent_tasks 表恢复任务状态，返回格式化的提醒文本。
|
||||
///
|
||||
/// 如果表不存在或没有任务记录，返回 None。
|
||||
@@ -82,3 +138,44 @@ async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option<Stri
|
||||
lines.join("\n")
|
||||
))
|
||||
}
|
||||
|
||||
/// 查询会话中最近一次动态上下文快照的哈希（无则 None）
|
||||
async fn last_dynamic_context_hash(db: &SqlitePool, session_id: &str) -> Option<String> {
|
||||
let metadata: Option<String> = sqlx::query_scalar(
|
||||
"SELECT metadata FROM agent_messages WHERE session_id = ? AND role = 'user' AND active = 1 AND metadata LIKE '%\"dynamic_context\"%' ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
metadata
|
||||
.and_then(|m| serde_json::from_str::<serde_json::Value>(&m).ok())
|
||||
.and_then(|v| v.get("context_hash")?.as_str().map(|s| s.to_string()))
|
||||
}
|
||||
|
||||
/// 持久化动态上下文快照（metadata 携带哈希供下轮比较）
|
||||
async fn persist_dynamic_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
msg: &ChatMessage,
|
||||
hash: u64,
|
||||
) {
|
||||
use crate::agent::engine::{DbMessageSink, MessageSink as _};
|
||||
let sink = DbMessageSink {
|
||||
db: db.clone(),
|
||||
session_id: session_id.to_string(),
|
||||
agent_name: "lead".to_string(),
|
||||
fixed_metadata: None,
|
||||
};
|
||||
sink.save(
|
||||
turn_index,
|
||||
0,
|
||||
msg,
|
||||
None,
|
||||
Some(serde_json::json!({ "dynamic_context": true, "context_hash": hash.to_string() })),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -1,27 +1,95 @@
|
||||
// src/agent/runtime/duplicate_detector.rs
|
||||
//
|
||||
// 同质调用检测器：检测连续重复的工具调用，防止死循环。
|
||||
// 同质调用检测器：检测重复的工具调用，防止死循环。
|
||||
//
|
||||
// 历史实现是单槽（只记住上一次调用），A/B 交替死循环
|
||||
// （call A → call B → call A → call B ...）检测不到。
|
||||
// 现改为"连续计数 + 滑动窗口计数"双通道：
|
||||
// - 连续通道：连续相同调用 ≥ threshold（原语义，立即触发）
|
||||
// - 窗口通道：最近 threshold*2 次调用中同一调用出现 ≥ threshold+2 次
|
||||
// （捕获 A/B 交替；门槛略高以放过合理的重复只读调用）
|
||||
|
||||
use std::collections::VecDeque;
|
||||
|
||||
/// 同质调用检测器
|
||||
#[derive(Debug, Default)]
|
||||
pub struct DuplicateDetector {
|
||||
last_call: Option<(String, String)>, // (tool_name, arguments)
|
||||
consecutive_count: usize,
|
||||
/// 最近调用的滑动窗口（name, arguments）
|
||||
recent: VecDeque<(String, String)>,
|
||||
}
|
||||
|
||||
impl DuplicateDetector {
|
||||
/// 记录一次调用，返回是否检测到死循环
|
||||
pub fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool {
|
||||
let key = (tool_name.to_string(), arguments.to_string());
|
||||
if self.last_call.as_ref() == Some(&key) {
|
||||
self.consecutive_count += 1;
|
||||
if self.consecutive_count >= threshold {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.last_call = Some(key);
|
||||
self.consecutive_count = 1;
|
||||
let window = (threshold * 2).max(4);
|
||||
|
||||
self.recent.push_back(key.clone());
|
||||
while self.recent.len() > window {
|
||||
self.recent.pop_front();
|
||||
}
|
||||
|
||||
// 连续通道：尾部连续出现次数
|
||||
let consecutive = self.recent.iter().rev().take_while(|k| *k == &key).count();
|
||||
if threshold >= 2 && consecutive >= threshold {
|
||||
return true;
|
||||
}
|
||||
|
||||
// 窗口通道：同一调用在窗口内出现次数（捕获 A/B 交替）。
|
||||
// 门槛与 threshold 相同：纯交替下窗口内单键最多出现 window/2 =
|
||||
// threshold 次，恰好可达；合理的间隔重复（≤ threshold-1 次）不受影响。
|
||||
let occurrences = self.recent.iter().filter(|k| **k == key).count();
|
||||
threshold >= 2 && occurrences >= threshold
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_consecutive_duplicates_detected() {
|
||||
let mut det = DuplicateDetector::default();
|
||||
assert!(!det.record("search_papers", "q=1", 3));
|
||||
assert!(!det.record("search_papers", "q=1", 3));
|
||||
assert!(det.record("search_papers", "q=1", 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alternating_loop_detected() {
|
||||
// A/B 交替：单槽实现检测不到，窗口通道应捕获
|
||||
let mut det = DuplicateDetector::default();
|
||||
let mut tripped = false;
|
||||
for i in 0..6 {
|
||||
let a = det.record("tool_a", "{}", 3);
|
||||
let b = det.record("tool_b", "{}", 3);
|
||||
if a || b {
|
||||
// 交替 5 次出现（threshold+2 = 5）后触发
|
||||
assert!(i >= 1, "不应过早触发 (round {})", i);
|
||||
tripped = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(tripped, "A/B 交替死循环应被检测到");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spaced_repeats_not_detected() {
|
||||
// 合理的重复只读调用（间隔其他调用）不应触发
|
||||
let mut det = DuplicateDetector::default();
|
||||
det.record("read_file", "a.rs", 3);
|
||||
det.record("grep_files", "pat", 3);
|
||||
det.record("read_file", "b.rs", 3);
|
||||
det.record("glob_files", "*.rs", 3);
|
||||
let tripped = det.record("read_file", "a.rs", 3);
|
||||
assert!(!tripped);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_args_not_duplicates() {
|
||||
let mut det = DuplicateDetector::default();
|
||||
for q in ["q=1", "q=2", "q=3", "q=4", "q=5", "q=6", "q=7", "q=8"] {
|
||||
assert!(!det.record("search_papers", q, 3));
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
// src/agent/runtime/executor/helpers.rs
|
||||
//
|
||||
// 执行器辅助类型与函数：PreparedCall, ToolResultMessage, ToolExecutionResult,
|
||||
// execute_single_tool, process_single_result, save_tool_message_sync。
|
||||
// execute_single_tool, process_single_result。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::clients::llm::ChatMessage;
|
||||
@@ -15,6 +13,7 @@ use super::AgentStreamEvent;
|
||||
use crate::agent::hooks::{
|
||||
event_label, HookRegistry, PostToolUseContext, PostToolUseFailureContext,
|
||||
};
|
||||
use crate::agent::tools::ToolOutput;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreparedCall {
|
||||
@@ -42,8 +41,6 @@ pub struct ToolExecutionResult {
|
||||
}
|
||||
|
||||
/// 执行单个工具调用（含超时和取消检测）。
|
||||
///
|
||||
/// 从原 `execute_parallel` 的闭包提取，供分区后的批次执行复用。
|
||||
pub(super) async fn execute_single_tool(
|
||||
tool_opt: Option<&dyn crate::agent::tools::AgentTool>,
|
||||
args: serde_json::Value,
|
||||
@@ -51,10 +48,10 @@ pub(super) async fn execute_single_tool(
|
||||
cancelled: &Arc<AtomicBool>,
|
||||
timeout_dur: std::time::Duration,
|
||||
tool_name: &str,
|
||||
) -> crate::agent::tools::ToolOutput {
|
||||
) -> ToolOutput {
|
||||
let tool = match tool_opt {
|
||||
Some(t) => t,
|
||||
None => return crate::agent::tools::ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
None => return ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
let interrupt_behavior = tool.interrupt_behavior();
|
||||
@@ -76,7 +73,7 @@ pub(super) async fn execute_single_tool(
|
||||
res = tokio::time::timeout(timeout_dur, tool_fut) => {
|
||||
match res {
|
||||
Ok(output) => output,
|
||||
Err(_) => crate::agent::tools::ToolOutput::error(format!(
|
||||
Err(_) => ToolOutput::error(format!(
|
||||
"工具 {} 执行超时（{}秒）",
|
||||
tool_name,
|
||||
timeout_dur.as_secs()
|
||||
@@ -84,50 +81,50 @@ pub(super) async fn execute_single_tool(
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
crate::agent::tools::ToolOutput::error("执行已被用户取消")
|
||||
ToolOutput::error("执行已被用户取消")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 处理单个工具执行结果（SSE 事件、PostToolUse hooks、持久化）。
|
||||
///
|
||||
/// 从原 `execute_parallel` 的结果处理循环提取。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn process_single_result(
|
||||
tool_call_id: &str,
|
||||
tool_name: &str,
|
||||
tool_args: &serde_json::Value,
|
||||
output: &crate::agent::tools::ToolOutput,
|
||||
output: &ToolOutput,
|
||||
cancelled_flag: bool,
|
||||
exec_start: std::time::Instant,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
hook_registry: &HookRegistry,
|
||||
library_dir: &std::path::Path,
|
||||
sid: &str,
|
||||
agent_name: &str,
|
||||
step: usize,
|
||||
max_output_chars: usize,
|
||||
ctx: &super::ExecutorContext<'_>,
|
||||
tool_messages: &mut Vec<ToolResultMessage>,
|
||||
additional_contexts: &mut Vec<String>,
|
||||
db: &SqlitePool,
|
||||
turn_index: i32,
|
||||
tool_registry: &crate::agent::tools::ToolRegistry,
|
||||
) {
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result;
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result_for;
|
||||
|
||||
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
|
||||
let step = ctx.step;
|
||||
let tap = &ctx.tap;
|
||||
|
||||
let (is_internal, display_name) = if let Some(tool) = tool_registry.get(tool_name) {
|
||||
let empty_hooks;
|
||||
let hook_registry: &HookRegistry = match ctx.hook_registry {
|
||||
Some(h) => h,
|
||||
None => {
|
||||
empty_hooks = HookRegistry::new();
|
||||
&empty_hooks
|
||||
}
|
||||
};
|
||||
|
||||
let (is_internal, display_name) = if let Some(tool) = ctx.tool_registry.get(tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, tool_name.to_string())
|
||||
};
|
||||
|
||||
// SSE 事件 — 立即推送到前端
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
name: tool_name.to_string(),
|
||||
display_name,
|
||||
name: tap.display(tool_name),
|
||||
display_name: tap.display(&display_name),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
@@ -137,23 +134,30 @@ pub(super) async fn process_single_result(
|
||||
|
||||
// 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub
|
||||
// 但对于已从磁盘读取内容的工具（如 read_file），跳过持久化以防止级联
|
||||
let tool_results_dir = library_dir.join(".agent").join("tool-results");
|
||||
let tool_results_dir = ctx
|
||||
.app_state
|
||||
.config
|
||||
.storage
|
||||
.library_dir
|
||||
.join(".agent")
|
||||
.join("tool-results");
|
||||
let (processed_content, _persisted_path) = if output.skip_persist {
|
||||
(output.content.clone(), None)
|
||||
} else {
|
||||
maybe_persist_tool_result(
|
||||
maybe_persist_tool_result_for(
|
||||
&output.content,
|
||||
tool_call_id,
|
||||
max_output_chars,
|
||||
ctx.max_output_chars,
|
||||
&tool_results_dir,
|
||||
tool_name,
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
// PostToolUse hook
|
||||
let post_ctx = PostToolUseContext {
|
||||
session_id: sid.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
session_id: ctx.session_id.to_string(),
|
||||
agent_name: ctx.agent_name.to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
tool_args: tool_args.clone(),
|
||||
output_content: processed_content.clone(),
|
||||
@@ -164,9 +168,14 @@ pub(super) async fn process_single_result(
|
||||
let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
|
||||
let final_content = post_result.final_content;
|
||||
|
||||
// 非可信内容包裹（间接 prompt 注入防御）
|
||||
let llm_content =
|
||||
crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content);
|
||||
// 非可信内容包裹（间接 prompt 注入防御）。
|
||||
// 是否包裹由工具通过 AgentTool::untrusted_output 声明；
|
||||
// 未知工具回退到名字启发式（mcp__*/web_* 前缀）。
|
||||
let llm_content = if ctx.tool_registry.untrusted_output(tool_name) {
|
||||
crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content)
|
||||
} else {
|
||||
final_content
|
||||
};
|
||||
|
||||
// 收集 PostToolUse hook 注入的上下文
|
||||
if !post_result.tagged_contexts.is_empty() {
|
||||
@@ -179,8 +188,8 @@ pub(super) async fn process_single_result(
|
||||
));
|
||||
}
|
||||
} else {
|
||||
for ctx in &post_result.additional_contexts {
|
||||
additional_contexts.push(ctx.clone());
|
||||
for c in &post_result.additional_contexts {
|
||||
additional_contexts.push(c.clone());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -200,8 +209,8 @@ pub(super) async fn process_single_result(
|
||||
// PostToolUseFailure hook
|
||||
if output.is_error {
|
||||
let failure_ctx = PostToolUseFailureContext {
|
||||
session_id: sid.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
session_id: ctx.session_id.to_string(),
|
||||
agent_name: ctx.agent_name.to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
tool_args: tool_args.clone(),
|
||||
error_message: output.content.clone(),
|
||||
@@ -217,51 +226,25 @@ pub(super) async fn process_single_result(
|
||||
// 发送给 LLM 使用包裹后的内容（安全防御）
|
||||
let chat_message = ChatMessage::tool_result(tool_call_id, &llm_content);
|
||||
|
||||
// 持久化到数据库（fire-and-forget）
|
||||
save_tool_message_sync(db, sid, turn_index, step, &chat_message);
|
||||
// 持久化到数据库（fire-and-forget；无 sink 时跳过）。
|
||||
// canonical value 一并写入 metadata（重放/审计消费结构化结果）。
|
||||
if let Some(sink) = &ctx.sink {
|
||||
let sink = sink.clone();
|
||||
let msg = chat_message.clone();
|
||||
let turn_index = ctx.turn_index;
|
||||
let step_i = ctx.step;
|
||||
let extra = output
|
||||
.value
|
||||
.clone()
|
||||
.map(|v| serde_json::json!({ "canonical_value": v }));
|
||||
tokio::spawn(async move {
|
||||
sink.save(turn_index, step_i as i32, &msg, None, extra)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message,
|
||||
was_error: output.is_error,
|
||||
});
|
||||
}
|
||||
|
||||
/// 同步保存 tool 角色消息到数据库。
|
||||
pub(super) 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.text().unwrap_or("").to_string();
|
||||
let tool_call_id = msg.tool_call_id.clone();
|
||||
// 提前序列化，避免闭包内的生命周期问题
|
||||
let metadata_str =
|
||||
serde_json::to_string(&serde_json::json!({ "role": "tool" })).unwrap_or_default();
|
||||
let raw_json = serde_json::to_string(&msg).unwrap_or_default();
|
||||
// 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, metadata, raw_json, 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(&metadata_str)
|
||||
.bind(&raw_json)
|
||||
.bind("lead")
|
||||
.execute(&db_clone)
|
||||
.await
|
||||
{
|
||||
warn!("[Executor] 保存 tool 消息失败（非致命）: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
+256
-280
@@ -1,18 +1,21 @@
|
||||
// src/agent/runtime/executor/mod.rs
|
||||
//
|
||||
// 工具调用验证与并行执行器。
|
||||
//
|
||||
// 执行参数通过 ExecutorContext 组合传递（历史实现有约 20 个位置参数）。
|
||||
// 行为注入点：
|
||||
// - ask_policy: AskUser 权限请求是交互等待（主代理）还是 fail-closed 自动拒绝（子代理/队友）
|
||||
// - tap: SSE 事件出口（None = 静默；prefix = "[sub]"）
|
||||
// - sink: 消息持久化出口（None = 不落库）
|
||||
|
||||
mod helpers;
|
||||
|
||||
// Re-export 公共类型
|
||||
pub use helpers::{PreparedCall, ToolExecutionResult, ToolResultMessage};
|
||||
|
||||
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, oneshot};
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::api::{AppState, PendingPermission};
|
||||
@@ -26,26 +29,23 @@ use super::partitioner::ToolPartitioner;
|
||||
use super::permission::{PermissionChecker, PermissionResult};
|
||||
use super::permission_explainer::explain_permission;
|
||||
use super::{AgentStreamEvent, DuplicateDetector};
|
||||
use crate::agent::engine::{AskPolicy, EventTap, MessageSink};
|
||||
use crate::agent::hooks::{event_label, HookRegistry, PreToolUseContext};
|
||||
use crate::agent::tools::{ToolContext, ToolRegistry};
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
|
||||
use helpers::{execute_single_tool, process_single_result, save_tool_message_sync};
|
||||
use helpers::{execute_single_tool, process_single_result};
|
||||
|
||||
/// 验证工具调用：死循环检测 + 参数解析。
|
||||
///
|
||||
/// 返回 (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>,
|
||||
tool_registry: &ToolRegistry,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
tap: &EventTap,
|
||||
step: usize,
|
||||
) -> (Vec<PreparedCall>, bool) {
|
||||
let mut prepared_calls: Vec<PreparedCall> = Vec::new();
|
||||
@@ -68,7 +68,7 @@ pub fn validate_and_prepare(
|
||||
"[Executor] 检测到死循环：{} 连续调用 {} 次",
|
||||
tool_name, duplicate_threshold
|
||||
);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
tap.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name),
|
||||
});
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
@@ -94,10 +94,10 @@ pub fn validate_and_prepare(
|
||||
} else {
|
||||
(false, tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
name: tool_name.clone(),
|
||||
display_name,
|
||||
display_name: tap.display(&display_name),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
@@ -105,7 +105,6 @@ pub fn validate_and_prepare(
|
||||
is_internal,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&call_id, &error_output);
|
||||
save_tool_message_sync(db, session_id, turn_index, step, &tool_msg);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
@@ -121,36 +120,48 @@ pub fn validate_and_prepare(
|
||||
(prepared_calls, has_duplicate)
|
||||
}
|
||||
|
||||
/// 执行器上下文 — 组合全部行为注入点（替代约 20 个位置参数）。
|
||||
pub struct ExecutorContext<'a> {
|
||||
pub tool_registry: &'a ToolRegistry,
|
||||
pub app_state: Arc<AppState>,
|
||||
/// hooks 缺省时使用空注册表语义（teammate 等无人值守场景）
|
||||
pub hook_registry: Option<&'a HookRegistry>,
|
||||
pub permission_checker: Option<&'a PermissionChecker>,
|
||||
pub session_checker: Option<&'a PermissionChecker>,
|
||||
pub denial_tracker: Option<&'a std::sync::Mutex<DenialTracker>>,
|
||||
pub checkpoint_manager: Option<&'a CheckpointManager>,
|
||||
pub db: &'a sqlx::SqlitePool,
|
||||
pub session_id: &'a str,
|
||||
pub agent_name: &'a str,
|
||||
pub turn_index: i32,
|
||||
pub step: usize,
|
||||
pub tool_timeout_secs: u64,
|
||||
pub max_output_chars: usize,
|
||||
pub read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
pub enable_thinking: bool,
|
||||
pub additional_allowed_dirs: Vec<String>,
|
||||
/// AskUser 权限请求处理策略
|
||||
pub ask_policy: AskPolicy,
|
||||
/// 取消观察：会话取消表（交互式）或外部原子标志
|
||||
pub cancel: Arc<AtomicBool>,
|
||||
pub interactive_cancel: bool,
|
||||
pub cancel_session_id: Option<String>,
|
||||
pub tap: EventTap,
|
||||
pub sink: Option<Arc<dyn MessageSink>>,
|
||||
}
|
||||
|
||||
/// 并行执行所有准备好的工具调用。
|
||||
///
|
||||
/// 流程：
|
||||
/// 1. 权限检查（deny 规则阻止不可执行工具）
|
||||
/// 2. 发送 ToolCall SSE 事件
|
||||
/// 3. 运行 PreToolUse hooks
|
||||
/// 1. Hardline 预检查（不可绕过的参数级拒绝）
|
||||
/// 2. 权限检查（deny 规则拦截；ask 按策略交互等待或 fail-closed 拒绝）
|
||||
/// 3. 发送 ToolCall SSE 事件 + 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>,
|
||||
session_permission_checker: Option<&PermissionChecker>,
|
||||
denial_tracker: Option<&std::sync::Mutex<DenialTracker>>,
|
||||
checkpoint_manager: Option<&std::sync::Arc<CheckpointManager>>,
|
||||
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>>,
|
||||
enable_thinking: bool,
|
||||
additional_allowed_dirs: Vec<String>,
|
||||
ctx: &ExecutorContext<'_>,
|
||||
) -> ToolExecutionResult {
|
||||
if prepared_calls.is_empty() {
|
||||
return ToolExecutionResult {
|
||||
@@ -162,7 +173,72 @@ pub async fn execute_parallel(
|
||||
};
|
||||
}
|
||||
|
||||
let sid = session_id.to_string();
|
||||
let tool_registry = ctx.tool_registry;
|
||||
let tap = &ctx.tap;
|
||||
let step = ctx.step;
|
||||
let sid = ctx.session_id.to_string();
|
||||
let _turn_index = ctx.turn_index;
|
||||
let empty_hooks;
|
||||
let hook_registry: &HookRegistry = match ctx.hook_registry {
|
||||
Some(h) => h,
|
||||
None => {
|
||||
empty_hooks = HookRegistry::new();
|
||||
&empty_hooks
|
||||
}
|
||||
};
|
||||
|
||||
/// 工具消息的统一落库出口（无 sink 时跳过）
|
||||
fn persist_tool_message(
|
||||
sink: &Option<Arc<dyn MessageSink>>,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
msg: &ChatMessage,
|
||||
) {
|
||||
if let Some(sink) = sink {
|
||||
let sink = sink.clone();
|
||||
let msg = msg.clone();
|
||||
tokio::spawn(async move {
|
||||
sink.save(turn_index, step as i32, &msg, None, None).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// 拒绝结果的统一出口：SSE 事件 + 消息落库 + 拒绝追踪
|
||||
async fn record_denial(
|
||||
ctx: &ExecutorContext<'_>,
|
||||
tap: &EventTap,
|
||||
tool_call_id: &str,
|
||||
tool_name: &str,
|
||||
err_output: &str,
|
||||
metadata: serde_json::Value,
|
||||
tool_messages: &mut Vec<ToolResultMessage>,
|
||||
) {
|
||||
let (is_internal, display_name) = match ctx.tool_registry.get(tool_name) {
|
||||
Some(tool) => (tool.is_internal(), tool.display_name().to_string()),
|
||||
None => (false, tool_name.to_string()),
|
||||
};
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
name: tool_name.to_string(),
|
||||
display_name: tap.display(&display_name),
|
||||
output: err_output.to_string(),
|
||||
is_error: true,
|
||||
metadata,
|
||||
step: ctx.step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(tool_call_id, err_output);
|
||||
persist_tool_message(&ctx.sink, ctx.turn_index, ctx.step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 1: 发送 ToolCall SSE 事件
|
||||
for prep in prepared_calls {
|
||||
@@ -171,10 +247,10 @@ pub async fn execute_parallel(
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
tap.send(AgentStreamEvent::ToolCall {
|
||||
id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
name: tap.display(&prep.tool_name),
|
||||
display_name: tap.display(&display_name),
|
||||
arguments: prep.args.clone(),
|
||||
step,
|
||||
is_internal,
|
||||
@@ -203,14 +279,12 @@ pub async fn execute_parallel(
|
||||
prep.tool_name, reason
|
||||
);
|
||||
}
|
||||
// 收集所有阻塞错误详情（含多个 hook 同时 block 的情况）
|
||||
for be in &result.blocking_errors {
|
||||
hook_blocking_errors.push(format!(
|
||||
"[{}] 阻止 {}: {}",
|
||||
be.hook_name, prep.tool_name, be.reason
|
||||
));
|
||||
}
|
||||
// 收集 hook 的权限请求（保留完整信息用于 AskUser prompt）
|
||||
if let Some((permission, tool_name)) = result.permission_info() {
|
||||
info!(
|
||||
"[Executor] Hook 请求了工具 {} 的权限确认: {}",
|
||||
@@ -220,9 +294,7 @@ pub async fn execute_parallel(
|
||||
} else {
|
||||
hook_permission_info.push(None);
|
||||
}
|
||||
// 使用 hook 可能修改后的参数
|
||||
mutated_args.push(result.final_args);
|
||||
// 收集所有 hook 注入的上下文（优先使用带来源标记的 tagged_contexts）
|
||||
if !result.tagged_contexts.is_empty() {
|
||||
for tc in &result.tagged_contexts {
|
||||
additional_contexts.push(format!(
|
||||
@@ -233,39 +305,25 @@ pub async fn execute_parallel(
|
||||
));
|
||||
}
|
||||
} else {
|
||||
for ctx in &result.additional_contexts {
|
||||
additional_contexts.push(ctx.clone());
|
||||
for c in &result.additional_contexts {
|
||||
additional_contexts.push(c.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2.5: 权限检查 — PermissionChecker 规则引擎拦截被拒绝的工具。
|
||||
// Phase 2.5: Hardline 预检查 + 权限检查。
|
||||
// 被拒绝的工具直接注入错误 result，不进入执行队列。
|
||||
let mut tool_messages: Vec<ToolResultMessage> = Vec::new();
|
||||
let mut denied_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
|
||||
|
||||
// ── Hardline 预检查（在任何模式下都不可绕过）──
|
||||
// 在 PermissionChecker 之前执行，确保 hardline 规则始终生效。
|
||||
// 检查逻辑由工具自身通过 AgentTool::hardline_check 声明（按参数路由，
|
||||
// 不再按工具名字符串匹配）。
|
||||
for (i, prep) in prepared_calls.iter().enumerate() {
|
||||
let hardline_result = match prep.tool_name.as_str() {
|
||||
"run_bash" => {
|
||||
if let Some(cmd) = prep.args.get("command").and_then(|v| v.as_str()) {
|
||||
hardline::check_command(cmd)
|
||||
} else {
|
||||
hardline::HardlineResult::allowed()
|
||||
}
|
||||
}
|
||||
"file_write" | "file_edit" => {
|
||||
if let Some(path) = prep.args.get("file_path").and_then(|v| v.as_str()) {
|
||||
hardline::check_dangerous_path(path)
|
||||
} else if let Some(path) = prep.args.get("path").and_then(|v| v.as_str()) {
|
||||
hardline::check_dangerous_path(path)
|
||||
} else {
|
||||
hardline::HardlineResult::allowed()
|
||||
}
|
||||
}
|
||||
_ => hardline::HardlineResult::allowed(),
|
||||
};
|
||||
let hardline_result = tool_registry
|
||||
.get(&prep.tool_name)
|
||||
.and_then(|tool| tool.hardline_check(&prep.args))
|
||||
.unwrap_or_else(hardline::HardlineResult::allowed);
|
||||
|
||||
if hardline_result.blocked {
|
||||
warn!(
|
||||
@@ -275,48 +333,30 @@ pub async fn execute_parallel(
|
||||
hardline_result.reason
|
||||
);
|
||||
let err_output = hardline_result.reason.clone();
|
||||
let (is_internal, display_name) = if let Some(tool) = tool_registry.get(&prep.tool_name)
|
||||
{
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({
|
||||
"hardline_blocked": true,
|
||||
"hardline_category": hardline_result.category,
|
||||
}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output);
|
||||
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
// 记录拒绝追踪
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(checker) = permission_checker {
|
||||
if let Some(checker) = ctx.permission_checker {
|
||||
for (i, prep) in prepared_calls.iter().enumerate() {
|
||||
// ── 权限决策合并（单调不变量：Deny 粘滞，只能收紧不能放松）──
|
||||
let mut perm_result = checker.check(&prep.tool_name, Some(&prep.args));
|
||||
perm_result = checker.apply_mode(perm_result, &prep.tool_name);
|
||||
|
||||
// Hook PermissionRequired — 若 Checker 返回 Allowed，升级为 Ask
|
||||
// 使用 hook 提供的具体权限描述替换泛型消息
|
||||
if let Some(Some((ref perm_desc, _))) = hook_permission_info.get(i) {
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser {
|
||||
@@ -330,14 +370,12 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
|
||||
// 工具级 check_permissions() — 在 PermissionChecker 结果基础上叠加
|
||||
// PermissionChecker Deny/Ask 优先，工具级规则在 Allow 时可升级为 Ask
|
||||
// 工具级 check_permissions() — 在已有结果基础上收紧
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
let tool_rules = tool.check_permissions(&prep.args);
|
||||
for tool_rule in &tool_rules {
|
||||
match tool_rule {
|
||||
crate::agent::tools::PermissionRule::Deny { reason, .. } => {
|
||||
// 工具级 Deny 仅在 PermissionChecker 未 Deny 时生效
|
||||
if !perm_result.is_denied() {
|
||||
perm_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
@@ -345,7 +383,6 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
crate::agent::tools::PermissionRule::Ask { message, .. } => {
|
||||
// 工具级 Ask：若 PermissionChecker 返回 Allowed，升级为 Ask
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
@@ -357,67 +394,57 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
|
||||
// 会话级权限检查（API 动态添加的规则，优先级高于环境变量规则）
|
||||
if let Some(session_checker) = session_permission_checker {
|
||||
let session_result = session_checker.check(&prep.tool_name, Some(&prep.args));
|
||||
// 会话规则结果覆盖或升级
|
||||
match session_result {
|
||||
PermissionResult::Denied { reason } => {
|
||||
// 会话 Deny 强制覆盖
|
||||
perm_result = PermissionResult::Denied { reason };
|
||||
}
|
||||
PermissionResult::AskUser { message } => {
|
||||
// 会话 Ask 在 Allow 时升级
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser { message };
|
||||
}
|
||||
}
|
||||
PermissionResult::Allowed => {
|
||||
// 会话 Allow 仅覆盖 Allowed，保持 Deny/AskUser 不变
|
||||
// 避免覆盖工具级 check_permissions() 升级的 AskUser
|
||||
}
|
||||
}
|
||||
// 会话级权限检查（API 动态添加的规则，只能收紧）
|
||||
if let Some(session_checker) = ctx.session_checker {
|
||||
perm_result = super::permission::tighten(
|
||||
perm_result,
|
||||
session_checker.check(&prep.tool_name, Some(&prep.args)),
|
||||
);
|
||||
}
|
||||
|
||||
match perm_result {
|
||||
PermissionResult::Denied { reason } => {
|
||||
warn!(
|
||||
"[Executor] PermissionChecker 拒绝了工具 {}: {}",
|
||||
"[Executor] 权限检查拒绝了工具 {}: {}",
|
||||
prep.tool_name, reason
|
||||
);
|
||||
let err_output =
|
||||
format!("工具 {} 被权限规则拒绝执行: {}", prep.tool_name, reason);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output);
|
||||
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
// 记录拒绝追踪
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({}),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
PermissionResult::AskUser { message } => {
|
||||
// 无人值守上下文（子代理/队友）无用户可问：fail-closed 自动拒绝
|
||||
if ctx.ask_policy == AskPolicy::AutoDeny {
|
||||
let reason =
|
||||
match crate::agent::engine::auto_deny_ask(&prep.tool_name, &message) {
|
||||
PermissionResult::Denied { reason } => reason,
|
||||
_ => "需要用户确认但在无人值守上下文中不可用".to_string(),
|
||||
};
|
||||
let err_output = format!("工具 {} 被拒绝执行: {}", prep.tool_name, reason);
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({ "auto_denied_ask": true }),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
continue;
|
||||
}
|
||||
|
||||
info!(
|
||||
"[Executor] PermissionChecker 请求用户确认工具 {}: {}",
|
||||
prep.tool_name, message
|
||||
@@ -427,8 +454,7 @@ pub async fn execute_parallel(
|
||||
let permission_exp = explain_permission(&prep.tool_name, &prep.args);
|
||||
let explanation_json = serde_json::to_value(&permission_exp).ok();
|
||||
|
||||
// 发送权限请求 SSE 事件
|
||||
let _ = tx.send(AgentStreamEvent::PermissionRequest {
|
||||
tap.send(AgentStreamEvent::PermissionRequest {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
message: message.clone(),
|
||||
@@ -439,17 +465,14 @@ pub async fn execute_parallel(
|
||||
// 创建 oneshot 通道等待用户响应
|
||||
let (resp_tx, resp_rx) = oneshot::channel();
|
||||
let perm_id = uuid::Uuid::new_v4().to_string();
|
||||
let tc_id = prep.tool_call_id.clone();
|
||||
let t_name = prep.tool_name.clone();
|
||||
|
||||
// 存储待处理的权限请求
|
||||
{
|
||||
let mut perms = app_state.session.pending_permissions.lock().await;
|
||||
let mut perms = ctx.app_state.session.pending_permissions.lock().await;
|
||||
perms.insert(
|
||||
perm_id.clone(),
|
||||
PendingPermission {
|
||||
tool_call_id: tc_id.clone(),
|
||||
tool_name: t_name.clone(),
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
message: message.clone(),
|
||||
arguments: prep.args.clone(),
|
||||
response_tx: resp_tx,
|
||||
@@ -458,12 +481,11 @@ pub async fn execute_parallel(
|
||||
);
|
||||
}
|
||||
|
||||
// 等待用户响应（120 秒超时）
|
||||
// 等待用户响应（120 秒超时 = fail-closed 拒绝）
|
||||
let timeout_dur = std::time::Duration::from_secs(120);
|
||||
let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await;
|
||||
|
||||
// 清理待处理的权限请求
|
||||
app_state
|
||||
ctx.app_state
|
||||
.session
|
||||
.pending_permissions
|
||||
.lock()
|
||||
@@ -473,99 +495,61 @@ pub async fn execute_parallel(
|
||||
match perm_result {
|
||||
Ok(Ok(response)) if response.allowed => {
|
||||
info!("[Executor] 用户允许了工具 {} 的执行", prep.tool_name);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: true,
|
||||
});
|
||||
// 用户允许 → 重置连续拒绝计数
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_success();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Ok(_response)) => {
|
||||
// 用户拒绝
|
||||
info!("[Executor] 用户拒绝了工具 {}", prep.tool_name);
|
||||
let err_output = format!("用户拒绝了工具 {} 的执行", prep.tool_name);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output);
|
||||
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
// 用户拒绝 → 记录拒绝追踪
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
denied_indices.insert(i);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: false,
|
||||
});
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({}),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
_ => {
|
||||
// 超时或通道关闭
|
||||
// 超时或通道关闭 → fail-closed
|
||||
warn!("[Executor] 权限请求超时或取消: {}", prep.tool_name);
|
||||
let err_output =
|
||||
format!("权限请求超时 (120s): {} 未获得用户确认", prep.tool_name);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output);
|
||||
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
// 超时 → 记录拒绝追踪
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
denied_indices.insert(i);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: false,
|
||||
});
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({ "timeout": true }),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
PermissionResult::Allowed => {
|
||||
// 工具被允许 → 重置连续拒绝计数
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_success();
|
||||
}
|
||||
@@ -575,38 +559,46 @@ pub async fn execute_parallel(
|
||||
}
|
||||
} // if let Some(checker)
|
||||
|
||||
// Phase 3: 分区并行执行（参考 Claude Code partitionToolCalls + runTools）。
|
||||
//
|
||||
// 改进：原实现将所有非拒绝工具放入单个 FuturesUnordered 无差别并发，
|
||||
// 可能导致非并发安全工具（如 run_bash）错误地并行执行。
|
||||
// 新实现使用 ToolPartitioner 将工具按并发安全性分批：
|
||||
// Phase 3: 分区并行执行。
|
||||
// ToolPartitioner 将工具按并发安全性分批：
|
||||
// - 连续的并发安全工具放入同一个并行批次（FuturesUnordered）
|
||||
// - 非并发安全工具独占一个串行批次（逐次执行）
|
||||
// 批次内工具执行完成后立即推送 SSE 事件，不等待整个批次完成。
|
||||
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 app_state_ref = ctx.app_state.clone();
|
||||
let _sid_ref = sid.clone();
|
||||
let interactive = ctx.interactive_cancel;
|
||||
let cancel_sid = ctx.cancel_session_id.clone();
|
||||
let external_flag = ctx.cancel.clone();
|
||||
|
||||
let cancel_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if app_state_ref.session.cancelled_runs.contains_key(&sid_ref) {
|
||||
let tripped = if interactive {
|
||||
cancel_sid
|
||||
.as_ref()
|
||||
.is_some_and(|s| app_state_ref.session.cancelled_runs.contains_key(s))
|
||||
} else {
|
||||
external_flag.load(Ordering::SeqCst)
|
||||
};
|
||||
if tripped {
|
||||
cancel_flag.store(true, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs);
|
||||
let timeout_dur = std::time::Duration::from_secs(ctx.tool_timeout_secs);
|
||||
|
||||
// ── Checkpoint 预触发：对文件变更类工具在执行前创建快照 ──
|
||||
if let Some(ckpt) = checkpoint_manager {
|
||||
// 由工具通过 AgentTool::causes_file_changes 声明
|
||||
if let Some(ckpt) = ctx.checkpoint_manager {
|
||||
let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
|
||||
for prep in prepared_calls
|
||||
.iter()
|
||||
.filter(|p| CheckpointManager::should_checkpoint(&p.tool_name))
|
||||
{
|
||||
for prep in prepared_calls.iter().filter(|p| {
|
||||
tool_registry
|
||||
.get(&p.tool_name)
|
||||
.is_some_and(|t| t.causes_file_changes())
|
||||
}) {
|
||||
ckpt.ensure_checkpoint(&cwd, &format!("pre-{}", prep.tool_name));
|
||||
}
|
||||
}
|
||||
@@ -624,7 +616,6 @@ pub async fn execute_parallel(
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
let batches = partitioner.partition(&non_denied_calls, tool_registry);
|
||||
|
||||
// 预设非拒绝工具中哪些原索引属于已拒绝列表（不会有，但安全起见）
|
||||
let original_index_of: std::collections::HashMap<String, usize> = non_denied
|
||||
.iter()
|
||||
.map(|(orig_idx, prep)| (prep.tool_call_id.clone(), *orig_idx))
|
||||
@@ -641,7 +632,6 @@ pub async fn execute_parallel(
|
||||
let mut was_cancelled = false;
|
||||
|
||||
// ── Phase 3c: 逐批次执行 ──
|
||||
// 批次之间串行；并行批次内工具并发执行；串行批次内工具逐个执行。
|
||||
for batch in &batches {
|
||||
if was_cancelled {
|
||||
break;
|
||||
@@ -649,7 +639,7 @@ pub async fn execute_parallel(
|
||||
|
||||
if batch.is_parallel {
|
||||
// ── 并行批次：FuturesUnordered 并发执行 ──
|
||||
let mut exec_futs: FuturesUnordered<_> = batch
|
||||
let mut exec_futs: futures_util::stream::FuturesUnordered<_> = batch
|
||||
.calls
|
||||
.iter()
|
||||
.map(|prep| {
|
||||
@@ -662,14 +652,16 @@ pub async fn execute_parallel(
|
||||
.get(orig_idx)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx =
|
||||
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
|
||||
.with_sse_tx(tx.clone())
|
||||
.with_session_id(session_id.to_string())
|
||||
.with_thinking(enable_thinking)
|
||||
.with_additional_dirs(additional_allowed_dirs.clone())
|
||||
.with_tool_call_id(prep.tool_call_id.clone())
|
||||
.with_max_output_chars(max_output_chars);
|
||||
let tool_ctx = crate::agent::tools::ToolContext::with_file_cache(
|
||||
ctx.app_state.clone(),
|
||||
ctx.read_file_state.clone(),
|
||||
)
|
||||
.with_sse_tx_opt(ctx.tap.sender())
|
||||
.with_session_id(sid.clone())
|
||||
.with_thinking(ctx.enable_thinking)
|
||||
.with_additional_dirs(ctx.additional_allowed_dirs.clone())
|
||||
.with_tool_call_id(prep.tool_call_id.clone())
|
||||
.with_max_output_chars(ctx.max_output_chars);
|
||||
let cancelled = cancelled.clone();
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
@@ -696,6 +688,7 @@ pub async fn execute_parallel(
|
||||
.collect();
|
||||
|
||||
// 渐进式处理：每个工具一完成就处理
|
||||
use futures_util::StreamExt;
|
||||
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
|
||||
exec_futs.next().await
|
||||
{
|
||||
@@ -709,18 +702,9 @@ pub async fn execute_parallel(
|
||||
&output,
|
||||
cancelled_flag,
|
||||
exec_start,
|
||||
tx,
|
||||
hook_registry,
|
||||
&app_state.config.storage.library_dir,
|
||||
&sid,
|
||||
agent_name,
|
||||
step,
|
||||
max_output_chars,
|
||||
ctx,
|
||||
&mut tool_messages,
|
||||
&mut additional_contexts,
|
||||
db,
|
||||
turn_index,
|
||||
tool_registry,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -736,12 +720,13 @@ pub async fn execute_parallel(
|
||||
.get(orig_idx)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx =
|
||||
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
|
||||
.with_sse_tx(tx.clone())
|
||||
.with_session_id(session_id.to_string())
|
||||
.with_thinking(enable_thinking)
|
||||
.with_additional_dirs(additional_allowed_dirs.clone());
|
||||
let tool_ctx = crate::agent::tools::ToolContext::with_file_cache(
|
||||
ctx.app_state.clone(),
|
||||
ctx.read_file_state.clone(),
|
||||
)
|
||||
.with_session_id(sid.clone())
|
||||
.with_thinking(ctx.enable_thinking)
|
||||
.with_additional_dirs(ctx.additional_allowed_dirs.clone());
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
let output = execute_single_tool(
|
||||
@@ -765,18 +750,9 @@ pub async fn execute_parallel(
|
||||
&output,
|
||||
cancelled_flag,
|
||||
exec_start,
|
||||
tx,
|
||||
hook_registry,
|
||||
&app_state.config.storage.library_dir,
|
||||
&sid,
|
||||
agent_name,
|
||||
step,
|
||||
max_output_chars,
|
||||
ctx,
|
||||
&mut tool_messages,
|
||||
&mut additional_contexts,
|
||||
db,
|
||||
turn_index,
|
||||
tool_registry,
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
+224
-980
File diff suppressed because it is too large
Load Diff
+86
-163
@@ -508,109 +508,33 @@ impl PermissionChecker {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Permission Precedence Resolver ──
|
||||
// ── Monotonic Combination（权限单调性不变量）──
|
||||
|
||||
/// 多源权限决策的最终裁决。遵循正式的优先级规则表：
|
||||
/// 权限决策合并 — 整个权限体系的唯一不变量：
|
||||
/// **只能收紧，不能放松**（deny > ask > allow，最严格胜出）。
|
||||
///
|
||||
/// | Priority | Source | Overridable By |
|
||||
/// |----------|----------------------------------|----------------|
|
||||
/// | P0 | PermissionChecker::Deny | Nothing |
|
||||
/// | P1 | Tool-level PermissionRule::Deny | Nothing |
|
||||
/// | P2 | Session-level Checker::Deny | Nothing |
|
||||
/// | P3 | Hook PreToolUseAction::Block | P0-P2 |
|
||||
/// | P4 | Hook PermissionRequired | P0-P3 |
|
||||
/// | P5-P7 | Checker::Ask / Tool::Ask / Allow | Normal |
|
||||
/// 多源决策（env 规则、工具级声明、会话级规则、hook 请求）逐层调用本函数
|
||||
/// 合并，任何一层都无法放行另一层已收紧的决策——这是 guard 单调性的
|
||||
/// 权限等价物（参考 deepseek-harness "guards only tighten" 约束）。
|
||||
///
|
||||
/// `conflict_log` 记录被覆盖的决策，便于审计和调试。
|
||||
pub fn resolve_permission_precedence(
|
||||
checker_result: PermissionResult,
|
||||
tool_rules: &[crate::agent::tools::PermissionRule],
|
||||
hook_permission: Option<&(String, String)>, // (permission_desc, tool_name)
|
||||
hook_blocked: bool,
|
||||
session_result: Option<PermissionResult>,
|
||||
) -> (PermissionResult, Vec<String>) {
|
||||
let mut final_result = checker_result;
|
||||
let mut conflict_log: Vec<String> = Vec::new();
|
||||
|
||||
// ── P1: Tool-level Deny ──
|
||||
for rule in tool_rules {
|
||||
if let crate::agent::tools::PermissionRule::Deny { reason, .. } = rule {
|
||||
if !final_result.is_denied() {
|
||||
conflict_log.push(format!("Tool-level Deny overrides checker: {reason}"));
|
||||
final_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
} else {
|
||||
conflict_log.push(format!(
|
||||
"Tool-level Deny '{reason}' ignored: already Denied"
|
||||
));
|
||||
}
|
||||
break; // only handle first Deny
|
||||
}
|
||||
/// 具体语义：
|
||||
/// - `base` 已 Denied → 结果必为 Denied（粘滞，不可覆盖）
|
||||
/// - `base` 已 AskUser → 候选 Allowed 不生效；候选 Denied 升级为 Denied
|
||||
/// - `base` Allowed → 采用候选（候选可为 Allowed/AskUser/Denied）
|
||||
pub fn tighten(base: PermissionResult, candidate: PermissionResult) -> PermissionResult {
|
||||
match base {
|
||||
PermissionResult::Denied { .. } => base, // Deny 粘滞
|
||||
PermissionResult::AskUser { .. } => match candidate {
|
||||
// Ask 不能被放松为 Allow；可升级为 Deny
|
||||
PermissionResult::Denied { reason } => PermissionResult::Denied { reason },
|
||||
_ => base,
|
||||
},
|
||||
PermissionResult::Allowed => candidate,
|
||||
}
|
||||
|
||||
// ── P2: Session-level Deny ──
|
||||
if let Some(PermissionResult::Denied { reason }) = &session_result {
|
||||
conflict_log.push(format!("Session-level Deny overrides current: {reason}"));
|
||||
final_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
}
|
||||
|
||||
// ── P3: Hook Block ──
|
||||
if hook_blocked {
|
||||
conflict_log.push("Hook Block prevents execution".to_string());
|
||||
// Block is already handled in the executor via denied_indices;
|
||||
// here we record it for the conflict log.
|
||||
}
|
||||
|
||||
// ── P4: Hook PermissionRequired ──
|
||||
if let Some((perm_desc, _tool_name)) = hook_permission {
|
||||
if final_result.is_allowed() {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired upgrades Allowed → Ask: {perm_desc}"
|
||||
));
|
||||
} else if !final_result.is_denied() {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired coexists with current state: {perm_desc}"
|
||||
));
|
||||
} else {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired '{perm_desc}' ignored: already Denied"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// ── P5: Session-level Ask ──
|
||||
if let Some(PermissionResult::AskUser { message }) = &session_result {
|
||||
if final_result.is_allowed() {
|
||||
final_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
conflict_log.push("Session-level Ask overrides Allow".to_string());
|
||||
} else {
|
||||
conflict_log.push("Session-level Ask ignored: not Allowed".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
// ── P6: Tool-level Ask ──
|
||||
for rule in tool_rules {
|
||||
if let crate::agent::tools::PermissionRule::Ask { message, .. } = rule {
|
||||
if final_result.is_allowed() {
|
||||
conflict_log.push(format!("Tool-level Ask upgrades Allow: {message}"));
|
||||
final_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
} else {
|
||||
conflict_log.push("Tool-level Ask ignored: not Allowed".to_string());
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
(final_result, conflict_log)
|
||||
}
|
||||
|
||||
// ── Permission Precedence Resolver ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -965,85 +889,84 @@ mod tests {
|
||||
assert!(matches!(result, PermissionResult::AskUser { .. }));
|
||||
}
|
||||
|
||||
// ── resolve_permission_precedence tests ──
|
||||
// ── tighten（单调合并不变量）tests ──
|
||||
|
||||
#[test]
|
||||
fn test_precedence_checker_deny_wins_over_all() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
fn test_tighten_deny_is_sticky() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::Denied {
|
||||
reason: "policy".into(),
|
||||
};
|
||||
// 任何候选都无法放松 Deny
|
||||
assert!(tighten(base.clone(), PermissionResult::Allowed).is_denied());
|
||||
assert!(tighten(
|
||||
base,
|
||||
PermissionResult::AskUser {
|
||||
message: "m".into()
|
||||
}
|
||||
)
|
||||
.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tighten_ask_cannot_be_relaxed() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::AskUser {
|
||||
message: "need confirm".into(),
|
||||
};
|
||||
// Ask + Allow → Ask（不可放松）
|
||||
let r = tighten(base.clone(), PermissionResult::Allowed);
|
||||
assert!(matches!(r, PermissionResult::AskUser { .. }));
|
||||
// Ask + Deny → Deny（可升级）
|
||||
let r = tighten(
|
||||
base,
|
||||
PermissionResult::Denied {
|
||||
reason: "blocked by policy".into(),
|
||||
reason: "no".into(),
|
||||
},
|
||||
&[],
|
||||
Some(&("need confirmation".to_string(), "test_tool".to_string())),
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(
|
||||
!log.is_empty(),
|
||||
"conflict log should record the interaction"
|
||||
);
|
||||
assert!(r.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_tool_deny_overrides_allow() {
|
||||
use crate::agent::tools::{PermissionRule, PermissionRuleSource};
|
||||
let tool_rules = vec![PermissionRule::Deny {
|
||||
tool_name: "test_tool".into(),
|
||||
reason: "tool self-protection".into(),
|
||||
source: PermissionRuleSource::Env,
|
||||
}];
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&tool_rules,
|
||||
None,
|
||||
false,
|
||||
None,
|
||||
fn test_tighten_allowed_adopts_candidate() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::Allowed;
|
||||
assert_eq!(
|
||||
tighten(base.clone(), PermissionResult::Allowed),
|
||||
PermissionResult::Allowed
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(!log.is_empty());
|
||||
assert!(matches!(
|
||||
tighten(
|
||||
base.clone(),
|
||||
PermissionResult::AskUser {
|
||||
message: "m".into()
|
||||
}
|
||||
),
|
||||
PermissionResult::AskUser { .. }
|
||||
));
|
||||
assert!(tighten(base, PermissionResult::Denied { reason: "d".into() }).is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_hook_block_recorded() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&[],
|
||||
None,
|
||||
true, // hook blocked
|
||||
None,
|
||||
);
|
||||
// Hook Block doesn't directly return Denied — it's logged for executor handling
|
||||
assert!(result.is_allowed());
|
||||
assert!(log.iter().any(|l| l.contains("Block")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_session_deny_overrides() {
|
||||
let (result, _log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&[],
|
||||
None,
|
||||
false,
|
||||
Some(PermissionResult::Denied {
|
||||
reason: "session deny".into(),
|
||||
}),
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_hook_permission_ignored_when_denied() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Denied {
|
||||
reason: "policy deny".into(),
|
||||
fn test_tighten_layers_only_tighten() {
|
||||
// 模拟三层合并：env Allowed + 工具级 Ask + 会话级 Allow → 最终 Ask
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let merged = PermissionResult::Allowed;
|
||||
let merged = tighten(
|
||||
merged,
|
||||
PermissionResult::AskUser {
|
||||
message: "tool ask".into(),
|
||||
},
|
||||
&[],
|
||||
Some(&("need confirm".to_string(), "test_tool".to_string())),
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(log.iter().any(|l| l.contains("ignored")));
|
||||
let merged = tighten(merged, PermissionResult::Allowed);
|
||||
assert!(matches!(merged, PermissionResult::AskUser { .. }));
|
||||
|
||||
// 反向：env Deny + 工具级 Allow + 会话级 Allow → 最终 Deny
|
||||
let merged = PermissionResult::Denied {
|
||||
reason: "env".into(),
|
||||
};
|
||||
let merged = tighten(merged.clone(), PermissionResult::Allowed);
|
||||
let merged = tighten(merged, PermissionResult::Allowed);
|
||||
assert!(merged.is_denied());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,6 +97,45 @@ pub async fn create_or_resume_session(
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用调用方预分配的会话 ID 创建（或恢复）会话。
|
||||
///
|
||||
/// 与 `create_or_resume_session` 的区别：会话不存在时以该 ID **创建**而非报错。
|
||||
/// 供 API 层使用——请求进入时预生成会话 ID，使取消标记、运行时缓存等
|
||||
/// 按 ID 索引的机制对新会话的首个请求同样生效（历史上新会话首请求
|
||||
/// 超时只能 abort 任务，无法写入取消标记）。
|
||||
pub async fn create_or_resume_session_preallocated(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
llm: &LlmClient,
|
||||
mode: &str,
|
||||
) -> anyhow::Result<SessionInfo> {
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if exists {
|
||||
return create_or_resume_session(db, Some(session_id.to_string()), llm, mode).await;
|
||||
}
|
||||
|
||||
sqlx::query("INSERT INTO agent_sessions (session_id, title, model, mode) VALUES (?, ?, ?, ?)")
|
||||
.bind(session_id)
|
||||
.bind("")
|
||||
.bind(llm.model().await)
|
||||
.bind(mode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: session_id.to_string(),
|
||||
turn_index: 0,
|
||||
mode: mode.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 从数据库加载会话的运行模式。
|
||||
///
|
||||
/// 返回 None 表示会话不存在或已删除。
|
||||
@@ -288,6 +327,9 @@ pub async fn rewind_to_message(
|
||||
|
||||
let preview: String = target_content.chars().take(120).collect();
|
||||
|
||||
// 回退改变了历史 → 上下文快照失效（其 base 高水位与消息集合不再对应）
|
||||
super::session_events::remove_context_snapshots(db, session_id).await;
|
||||
|
||||
info!(
|
||||
"[Session] 回退完成: session={}, rewound={}, to_id={}, new_turn={}",
|
||||
session_id, count, target_message_id, new_turn_index
|
||||
@@ -483,6 +525,9 @@ pub async fn retry_last_turn(
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 硬删除改变了历史 → 上下文快照失效
|
||||
super::session_events::remove_context_snapshots(db, session_id).await;
|
||||
|
||||
info!(
|
||||
"[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}, has_image={}",
|
||||
session_id,
|
||||
@@ -595,6 +640,7 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
|
||||
|
||||
tx.commit().await?;
|
||||
|
||||
// 分叉是新会话：不复制事件（快照引用旧会话的消息 id，直接作废）
|
||||
info!(
|
||||
"[Session] 分叉完成: parent={}, branch={}, copied={} messages, forked_at={}",
|
||||
session_id, branch_id, copied, forked_at
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
// src/agent/runtime/session_events.rs
|
||||
//
|
||||
// 会话事件日志 — turn/compaction 生命周期与上下文快照的事件溯源层。
|
||||
//
|
||||
// 不变量："模型可见 ⟺ 已日志化" 的生命周期侧面：
|
||||
// - 每个 turn 以 turn_start 开始、turn_end 关闭；崩溃留下可检测的
|
||||
// 开口（dangling start），恢复时合成 interrupted 关闭而不截断
|
||||
// 已持久化的内容（参考 deepseek-harness 的崩溃恢复策略）。
|
||||
// - 压缩以 compaction_start/compaction_end 构成日志化锁。
|
||||
// - context_snapshot 记录压缩后的折叠上下文 + 消息高水位
|
||||
// （base_message_id），下一 turn 回放快照 + 增量消息，避免重复压缩。
|
||||
|
||||
use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// 事件类型常量
|
||||
pub mod event_types {
|
||||
pub const TURN_START: &str = "turn_start";
|
||||
pub const TURN_END: &str = "turn_end";
|
||||
pub const COMPACTION_START: &str = "compaction_start";
|
||||
pub const COMPACTION_END: &str = "compaction_end";
|
||||
/// 恢复时合成的中断标记（检测到开口的 turn_start/compaction_start）
|
||||
pub const TURN_INTERRUPTED: &str = "turn_interrupted";
|
||||
pub const COMPACTION_INTERRUPTED: &str = "compaction_interrupted";
|
||||
/// 压缩后的折叠上下文快照（payload 含 messages + base_message_id）
|
||||
pub const CONTEXT_SNAPSHOT: &str = "context_snapshot";
|
||||
}
|
||||
|
||||
/// 追加一条会话事件
|
||||
pub async fn append_event(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
event_type: &str,
|
||||
payload: Value,
|
||||
) {
|
||||
let payload_str = serde_json::to_string(&payload).unwrap_or_default();
|
||||
if let Err(e) = sqlx::query(
|
||||
"INSERT INTO agent_events (session_id, turn_index, event_type, payload) \
|
||||
VALUES (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(event_type)
|
||||
.bind(&payload_str)
|
||||
.execute(db)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 写入 {} 事件失败（非致命）: {}",
|
||||
event_type, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 最新一条指定类型的事件（无则 None）
|
||||
pub async fn latest_event(db: &SqlitePool, session_id: &str, event_type: &str) -> Option<Value> {
|
||||
let payload: Option<String> = sqlx::query_scalar(
|
||||
"SELECT payload FROM agent_events \
|
||||
WHERE session_id = ? AND event_type = ? ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(event_type)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
payload.and_then(|p| serde_json::from_str(&p).ok())
|
||||
}
|
||||
|
||||
/// 检查是否存在"开口"事件：最新的 `open_type` 之后没有新的 `close_type`。
|
||||
/// 用于崩溃恢复——开口的 turn_start/compaction_start 表示上次执行被中断。
|
||||
pub async fn has_open_event(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
open_type: &str,
|
||||
close_type: &str,
|
||||
) -> bool {
|
||||
let open_id: Option<i64> = sqlx::query_scalar(
|
||||
"SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(open_type)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
|
||||
let Some(open_id) = open_id else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let close_id: Option<i64> = sqlx::query_scalar(
|
||||
"SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(close_type)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
|
||||
close_id.is_none_or(|c| c < open_id)
|
||||
}
|
||||
|
||||
/// 恢复会话时的中断协调：检测开口的 turn/compaction，合成 interrupted 关闭。
|
||||
///
|
||||
/// 这是对崩溃/进程被杀的容错：不截断任何已持久化内容，
|
||||
/// 只补写"上次被中断"的事实标记。
|
||||
pub async fn reconcile_interrupted(db: &SqlitePool, session_id: &str) {
|
||||
if has_open_event(
|
||||
db,
|
||||
session_id,
|
||||
event_types::TURN_START,
|
||||
event_types::TURN_END,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 会话 {} 检测到未关闭的 turn_start（上次执行被中断），合成 interrupted 关闭",
|
||||
session_id
|
||||
);
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::TURN_END,
|
||||
serde_json::json!({ "reason": "interrupted", "synthesized": true }),
|
||||
)
|
||||
.await;
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::TURN_INTERRUPTED,
|
||||
serde_json::json!({ "note": "检测到崩溃/中断残留，已合成关闭" }),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
if has_open_event(
|
||||
db,
|
||||
session_id,
|
||||
event_types::COMPACTION_START,
|
||||
event_types::COMPACTION_END,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 会话 {} 检测到未关闭的 compaction_start，合成 interrupted 关闭",
|
||||
session_id
|
||||
);
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::COMPACTION_END,
|
||||
serde_json::json!({ "interrupted": true, "synthesized": true }),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 删除会话的上下文快照（rewind/retry/branch 后快照失效）
|
||||
pub async fn remove_context_snapshots(db: &SqlitePool, session_id: &str) {
|
||||
if let Err(e) = sqlx::query("DELETE FROM agent_events WHERE session_id = ? AND event_type = ?")
|
||||
.bind(session_id)
|
||||
.bind(event_types::CONTEXT_SNAPSHOT)
|
||||
.execute(db)
|
||||
.await
|
||||
{
|
||||
warn!("[SessionEvents] 清理上下文快照失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
/// 保存上下文快照：折叠后的消息 + 消息高水位。
|
||||
///
|
||||
/// `base_message_id` 是快照时 agent_messages 的最大 id；恢复时加载
|
||||
/// 快照消息 + id > base 的增量消息，避免重复压缩。
|
||||
pub async fn save_context_snapshot(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
messages: &[crate::clients::llm::ChatMessage],
|
||||
base_message_id: i64,
|
||||
) {
|
||||
// 排除 system 消息（每 turn 重建，且包含缓存的动态 section）
|
||||
let body: Vec<&crate::clients::llm::ChatMessage> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role != crate::clients::llm::MessageRole::System)
|
||||
.collect();
|
||||
|
||||
let serialized: Vec<Value> = body
|
||||
.iter()
|
||||
.map(|m| serde_json::to_value(m).unwrap_or(Value::Null))
|
||||
.collect();
|
||||
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
turn_index,
|
||||
event_types::CONTEXT_SNAPSHOT,
|
||||
serde_json::json!({
|
||||
"base_message_id": base_message_id,
|
||||
"message_count": serialized.len(),
|
||||
"messages": serialized,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
info!(
|
||||
"[SessionEvents] 保存上下文快照: session={} base_msg_id={} messages={}",
|
||||
session_id,
|
||||
base_message_id,
|
||||
serialized.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// 加载上下文快照（若存在）：返回 (快照消息, base_message_id)。
|
||||
pub async fn load_context_snapshot(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> Option<(Vec<crate::clients::llm::ChatMessage>, i64)> {
|
||||
let payload = latest_event(db, session_id, event_types::CONTEXT_SNAPSHOT).await?;
|
||||
|
||||
let base_message_id = payload.get("base_message_id")?.as_i64()?;
|
||||
let messages_json = payload.get("messages")?.as_array()?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for m in messages_json {
|
||||
match serde_json::from_value::<crate::clients::llm::ChatMessage>(m.clone()) {
|
||||
Ok(msg) => messages.push(msg),
|
||||
Err(e) => {
|
||||
warn!("[SessionEvents] 快照消息反序列化失败，放弃快照: {}", e);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some((messages, base_message_id))
|
||||
}
|
||||
|
||||
/// 获取会话消息的当前最大 id（快照高水位）
|
||||
pub async fn max_message_id(db: &SqlitePool, session_id: &str) -> i64 {
|
||||
sqlx::query_scalar("SELECT COALESCE(MAX(id), 0) FROM agent_messages WHERE session_id = ?")
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// 加载 id 大于 `after_id` 的增量历史消息（快照回放后拼接）
|
||||
pub async fn load_messages_after(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
after_id: i64,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
let rows: Vec<(
|
||||
i64,
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
)> = sqlx::query_as(
|
||||
"SELECT id, role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? AND id > ? AND active = 1 AND agent_name = 'lead' \
|
||||
ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(after_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (_id, 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(|j| serde_json::from_str(&j).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)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
async fn setup_db() -> SqlitePool {
|
||||
let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE agent_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
event_type TEXT NOT NULL,
|
||||
payload TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE agent_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
step_index INTEGER NOT NULL DEFAULT 0,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thought TEXT,
|
||||
tool_calls TEXT,
|
||||
tool_call_id TEXT,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
metadata TEXT,
|
||||
raw_json TEXT,
|
||||
agent_name TEXT NOT NULL DEFAULT 'lead',
|
||||
active INTEGER NOT NULL DEFAULT 1,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_event_detection() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s1";
|
||||
|
||||
// 无事件 → 无开口
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
|
||||
// turn_start 后无 turn_end → 开口
|
||||
append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await;
|
||||
assert!(has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
|
||||
// 关闭后 → 无开口
|
||||
append_event(
|
||||
&db,
|
||||
sid,
|
||||
0,
|
||||
"turn_end",
|
||||
serde_json::json!({ "reason": "completed" }),
|
||||
)
|
||||
.await;
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reconcile_interrupted_synthesizes_close() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s2";
|
||||
|
||||
append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await;
|
||||
append_event(&db, sid, 0, "compaction_start", serde_json::json!({})).await;
|
||||
|
||||
reconcile_interrupted(&db, sid).await;
|
||||
|
||||
// 开口已闭合
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
assert!(!has_open_event(&db, sid, "compaction_start", "compaction_end").await);
|
||||
// 合成了中断标记
|
||||
assert!(latest_event(&db, sid, "turn_interrupted").await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_context_snapshot_roundtrip() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s3";
|
||||
|
||||
use crate::clients::llm::ChatMessage;
|
||||
sqlx::query("INSERT INTO agent_messages (session_id, role, content, agent_name) VALUES (?, 'user', 'm1', 'lead')")
|
||||
.bind(sid)
|
||||
.execute(&db)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let base = max_message_id(&db, sid).await;
|
||||
assert_eq!(base, 1);
|
||||
|
||||
let messages = vec![
|
||||
ChatMessage::system("sys"),
|
||||
ChatMessage::user("[历史对话摘要]\n关于黑洞"),
|
||||
ChatMessage::assistant("结论"),
|
||||
];
|
||||
save_context_snapshot(&db, sid, 0, &messages, base).await;
|
||||
|
||||
let (loaded, loaded_base) = load_context_snapshot(&db, sid).await.unwrap();
|
||||
assert_eq!(loaded_base, base);
|
||||
// system 消息被排除
|
||||
assert_eq!(loaded.len(), 2);
|
||||
assert!(loaded[0].text().unwrap().contains("历史对话摘要"));
|
||||
|
||||
// 删除快照
|
||||
remove_context_snapshots(&db, sid).await;
|
||||
assert!(load_context_snapshot(&db, sid).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_messages_after() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s4";
|
||||
|
||||
for i in 1..=3 {
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_messages (id, session_id, role, content, agent_name) \
|
||||
VALUES (?, ?, 'user', ?, 'lead')",
|
||||
)
|
||||
.bind(i)
|
||||
.bind(sid)
|
||||
.bind(format!("msg-{}", i))
|
||||
.execute(&db)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let incremental = load_messages_after(&db, sid, 1).await.unwrap();
|
||||
assert_eq!(incremental.len(), 2);
|
||||
assert_eq!(incremental[0].text(), Some("msg-2"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// src/agent/runtime/session_registry.rs
|
||||
//
|
||||
// 会话级运行时注册表 — AgentRuntime 按会话缓存复用。
|
||||
//
|
||||
// 历史问题：chat_agent 每个 HTTP 请求 new 一个 AgentRuntime，导致注释中
|
||||
// 声称"跨 turn 共享"的状态（后台任务队列、压缩折叠日志、文件缓存、
|
||||
// 拒绝追踪器、压缩熔断器）实际随请求销毁：
|
||||
// - 后台任务结果跨请求丢失（队列无人 drain）
|
||||
// - 压缩后的上下文不回写，下一 turn 从原始消息重建再重新压缩
|
||||
// - prompt cache / file cache 每请求冷启动
|
||||
//
|
||||
// 本注册表将 runtime 生命周期与 session 对齐（参考 deepseek-harness 的
|
||||
// "runtime 状态属于会话而非请求"原则）：
|
||||
// - 同一会话的请求复用同一 runtime 实例
|
||||
// - 每个会话持有一个 turn 互斥锁，防止同会话并发 turn 互相破坏
|
||||
// - 空闲超过 TTL 的条目由后台清扫回收
|
||||
// - 会话删除时显式移除
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use tracing::info;
|
||||
|
||||
use super::AgentRuntime;
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 单个会话的注册表条目
|
||||
pub struct SessionRuntimeEntry {
|
||||
pub runtime: Arc<AgentRuntime>,
|
||||
/// 同会话 turn 串行化锁（防止并发 turn 破坏 turn_index/消息顺序）
|
||||
pub turn_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
pub last_used: Instant,
|
||||
/// 创建时的模式 ID（用于诊断"模式漂移"）
|
||||
pub mode_id: String,
|
||||
}
|
||||
|
||||
/// 会话运行时注册表
|
||||
pub struct SessionRuntimeRegistry {
|
||||
entries: DashMap<String, SessionRuntimeEntry>,
|
||||
idle_ttl: Duration,
|
||||
}
|
||||
|
||||
impl Default for SessionRuntimeRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl SessionRuntimeRegistry {
|
||||
pub fn new() -> Self {
|
||||
SessionRuntimeRegistry {
|
||||
entries: DashMap::new(),
|
||||
// 2 小时无活动的会话 runtime 允许被回收；
|
||||
// 后台任务等通过 Arc 自然延长实际生命周期
|
||||
idle_ttl: Duration::from_secs(2 * 60 * 60),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取（或创建）指定会话的 runtime。
|
||||
///
|
||||
/// `mode_id` 只在首次创建时生效——后续请求从 DB 恢复会话模式，
|
||||
/// 不会因请求参数不同而静默切换（模式由会话决定，一次创建后保持稳定）。
|
||||
pub fn get_or_create(
|
||||
&self,
|
||||
app_state: Arc<AppState>,
|
||||
session_key: &str,
|
||||
mode_id: &str,
|
||||
) -> Arc<AgentRuntime> {
|
||||
if let Some(mut entry) = self.entries.get_mut(session_key) {
|
||||
entry.last_used = Instant::now();
|
||||
return entry.runtime.clone();
|
||||
}
|
||||
|
||||
// DashMap entry API 避免创建竞态
|
||||
let entry = self
|
||||
.entries
|
||||
.entry(session_key.to_string())
|
||||
.or_insert_with(|| {
|
||||
info!(
|
||||
"[SessionRegistry] 为会话 {} 创建 runtime (mode={})",
|
||||
session_key, mode_id
|
||||
);
|
||||
SessionRuntimeEntry {
|
||||
runtime: Arc::new(AgentRuntime::new_for_session(
|
||||
app_state,
|
||||
session_key,
|
||||
mode_id,
|
||||
)),
|
||||
turn_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
last_used: Instant::now(),
|
||||
mode_id: mode_id.to_string(),
|
||||
}
|
||||
});
|
||||
entry.runtime.clone()
|
||||
}
|
||||
|
||||
/// 获取会话的 turn 串行化锁（不存在时返回 None）
|
||||
pub fn turn_lock(&self, session_key: &str) -> Option<Arc<tokio::sync::Mutex<()>>> {
|
||||
self.entries.get(session_key).map(|e| e.turn_lock.clone())
|
||||
}
|
||||
|
||||
/// 移除会话条目（会话删除时调用）
|
||||
pub fn remove(&self, session_key: &str) {
|
||||
if self.entries.remove(session_key).is_some() {
|
||||
info!("[SessionRegistry] 移除会话 {} 的 runtime", session_key);
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前缓存的会话数
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.is_empty()
|
||||
}
|
||||
|
||||
/// 回收空闲超过 TTL 的条目（后台周期调用）
|
||||
pub fn sweep_idle(&self) -> usize {
|
||||
let now = Instant::now();
|
||||
let mut expired: Vec<String> = Vec::new();
|
||||
for entry in self.entries.iter() {
|
||||
if now.duration_since(entry.last_used) > self.idle_ttl {
|
||||
expired.push(entry.key().clone());
|
||||
}
|
||||
}
|
||||
let count = expired.len();
|
||||
for key in expired {
|
||||
self.entries.remove(&key);
|
||||
info!("[SessionRegistry] 回收空闲会话 runtime: {}", key);
|
||||
}
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
// ── 会话级压缩守卫（替代进程级全局静态） ──
|
||||
//
|
||||
// 历史实现是单个进程级 AtomicBool：并发两个会话压缩时互相跳过。
|
||||
// 改为按 session_id 的集合守卫：同会话递归压缩被拦截，不同会话互不影响。
|
||||
|
||||
static COMPACTING_SESSIONS: std::sync::Mutex<Option<HashSet<String>>> = std::sync::Mutex::new(None);
|
||||
|
||||
/// 尝试进入压缩临界区。返回 false 表示该会话已有进行中的压缩（递归守卫触发）。
|
||||
pub fn try_begin_compaction(session_id: &str) -> bool {
|
||||
let mut guard = COMPACTING_SESSIONS
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
let set = guard.get_or_insert_with(HashSet::new);
|
||||
if set.contains(session_id) {
|
||||
false
|
||||
} else {
|
||||
set.insert(session_id.to_string());
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// 退出压缩临界区。会话不存在时静默（防御性）。
|
||||
pub fn end_compaction(session_id: &str) {
|
||||
let mut guard = COMPACTING_SESSIONS
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(set) = guard.as_mut() {
|
||||
set.remove(session_id);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_compaction_guard_per_session() {
|
||||
assert!(try_begin_compaction("session-a"));
|
||||
// 同会话递归被拦截
|
||||
assert!(!try_begin_compaction("session-a"));
|
||||
// 不同会话互不影响
|
||||
assert!(try_begin_compaction("session-b"));
|
||||
end_compaction("session-a");
|
||||
// 释放后可重新进入
|
||||
assert!(try_begin_compaction("session-a"));
|
||||
end_compaction("session-a");
|
||||
end_compaction("session-b");
|
||||
end_compaction("never-began"); // 防御性：不 panic
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,6 @@
|
||||
// 支持并发取消检测，累积推理内容、文本增量和工具调用。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::error;
|
||||
|
||||
use crate::clients::llm::{
|
||||
@@ -40,7 +39,7 @@ pub async fn process_llm_stream(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
tool_defs: &[ToolDefinition],
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
tap: &crate::agent::engine::EventTap,
|
||||
step: usize,
|
||||
session_id: &str,
|
||||
cancelled_runs: Arc<dashmap::DashMap<String, ()>>,
|
||||
@@ -51,7 +50,7 @@ pub async fn process_llm_stream(
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[Streaming] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
tap.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
return StreamOutput {
|
||||
@@ -95,15 +94,15 @@ pub async fn process_llm_stream(
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
tap.send(AgentStreamEvent::Thought {
|
||||
content: tap.thought(&accumulated_reasoning),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
tap.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
tool_call_id: None,
|
||||
});
|
||||
|
||||
@@ -1,929 +0,0 @@
|
||||
// src/agent/runtime/streaming_executor.rs
|
||||
//
|
||||
// 流式工具执行器。
|
||||
// 参考 Claude Code StreamingToolExecutor 设计。
|
||||
//
|
||||
// 当 LLM 流式输出 tool_use 块时，立即开始执行并发安全的工具。
|
||||
// 非并发安全的工具排队等待。结果按流中顺序 yield。
|
||||
//
|
||||
// 与 Claude Code 的对齐改进 (2026-06-22)：
|
||||
// 1. 真正的流式调度 — on_tool_use 中对并发安全工具立即 spawn tokio task
|
||||
// 2. 并发分区 — 自动分组连续只读工具并行执行
|
||||
// 3. Progress 流式 — 长操作进度消息即时 yield
|
||||
// 4. Sibling Abort — 副效应工具报错时级联中止兄弟姐妹
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 流式工具执行状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TrackedToolStatus {
|
||||
/// 工具调用已从 LLM 流中接收到，等待调度
|
||||
Queued,
|
||||
/// 正在执行中（spawned tokio task 运行中）
|
||||
Executing,
|
||||
/// 执行完成，结果就绪等待 yield
|
||||
Completed,
|
||||
/// 结果已 yield 给调用方
|
||||
Yielded,
|
||||
}
|
||||
|
||||
/// Sibling Abort 原因
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AbortReason {
|
||||
/// 兄弟工具出错触发的级联取消
|
||||
SiblingError { description: String },
|
||||
/// 用户主动中断
|
||||
UserInterrupted,
|
||||
}
|
||||
|
||||
/// 单次工具执行的结果
|
||||
#[derive(Debug)]
|
||||
struct ToolExecutionResult {
|
||||
tool_name: String,
|
||||
output: ToolOutput,
|
||||
}
|
||||
|
||||
/// 跟踪中的工具执行
|
||||
struct TrackedTool {
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
status: TrackedToolStatus,
|
||||
/// 执行完成后的输出
|
||||
output: Option<ToolOutput>,
|
||||
/// 并发安全的工具在 spawn 后的 JoinHandle
|
||||
handle: Option<JoinHandle<ToolExecutionResult>>,
|
||||
}
|
||||
|
||||
/// 流式工具执行器。
|
||||
///
|
||||
/// 参考 Claude Code `StreamingToolExecutor` (531 行 TypeScript)，
|
||||
/// 关键改进：并发安全工具立即 spawn tokio task，不等待 flush。
|
||||
pub struct StreamingToolExecutor {
|
||||
/// 所有跟踪中的工具（按 LLM 流中到达顺序）
|
||||
tracked: Vec<TrackedTool>,
|
||||
/// 工具注册表
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
/// 工具上下文（按需 clone 给每个 spawn 的 task）
|
||||
tool_context: ToolContext,
|
||||
/// Sibling Abort 广播通道 (tx)
|
||||
abort_tx: broadcast::Sender<AbortReason>,
|
||||
/// Sibling Abort 广播通道 (rx) — 保留以保持 channel 存活，
|
||||
/// 实际使用时通过 `abort_tx.subscribe()` 获取新接收端。
|
||||
#[allow(dead_code)]
|
||||
abort_rx: broadcast::Receiver<AbortReason>,
|
||||
/// 当前是否已发生错误（触发 sibling abort）
|
||||
has_errored: bool,
|
||||
/// 出错工具的描述（如 "bash(git push)"）
|
||||
errored_tool_desc: String,
|
||||
/// 最大并发数（预留，当前使用 executing_non_concurrent 控制）
|
||||
#[allow(dead_code)]
|
||||
max_concurrency: usize,
|
||||
/// 最大工具输出字符数
|
||||
max_output_chars: usize,
|
||||
/// 当前正在执行的非并发安全工具数（0 或 1）
|
||||
executing_non_concurrent: bool,
|
||||
/// 已完成但尚未 yield 的结果队列（按流顺序）
|
||||
completed_queue: VecDeque<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,
|
||||
tool_context,
|
||||
abort_tx,
|
||||
abort_rx,
|
||||
has_errored: false,
|
||||
errored_tool_desc: String::new(),
|
||||
max_concurrency,
|
||||
max_output_chars,
|
||||
executing_non_concurrent: false,
|
||||
completed_queue: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取 abort 广播发送端（供外部注入取消信号）。
|
||||
pub fn abort_sender(&self) -> broadcast::Sender<AbortReason> {
|
||||
self.abort_tx.clone()
|
||||
}
|
||||
|
||||
/// 当 LLM 流产生一个新的 tool_use 时调用。
|
||||
///
|
||||
/// 如果是并发安全工具且当前没有非并发安全工具在执行，立即 spawn tokio task。
|
||||
/// 否则加入队列等待调度。
|
||||
///
|
||||
/// 返回 true 表示该工具已立即开始执行，false 表示排队。
|
||||
pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool {
|
||||
let is_concurrency_safe = self
|
||||
.tool_registry
|
||||
.get(&name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let mut tool = TrackedTool {
|
||||
tool_call_id: call_id.clone(),
|
||||
tool_name: name.clone(),
|
||||
args: args.clone(),
|
||||
status: TrackedToolStatus::Queued,
|
||||
output: None,
|
||||
handle: None,
|
||||
};
|
||||
|
||||
let idx = self.tracked.len();
|
||||
let can_start_now = is_concurrency_safe && !self.executing_non_concurrent;
|
||||
|
||||
if can_start_now {
|
||||
// 立即 spawn tokio task（参考 Claude Code: addTool 立即 processQueue）
|
||||
info!(
|
||||
"[StreamingExecutor] 立即 spawn 并发安全工具: {} (id={})",
|
||||
name, call_id
|
||||
);
|
||||
let handle = self.spawn_tool_task(idx, call_id.clone(), name.clone(), args.clone());
|
||||
tool.handle = Some(handle);
|
||||
tool.status = TrackedToolStatus::Executing;
|
||||
} else {
|
||||
info!(
|
||||
"[StreamingExecutor] 排队工具: {} (concurrent={}, executing_non_concurrent={})",
|
||||
name, is_concurrency_safe, self.executing_non_concurrent
|
||||
);
|
||||
}
|
||||
|
||||
if !is_concurrency_safe {
|
||||
self.executing_non_concurrent = true;
|
||||
}
|
||||
|
||||
self.tracked.push(tool);
|
||||
can_start_now
|
||||
}
|
||||
|
||||
/// LLM 流结束后调用，等待所有剩余排队工具完成。
|
||||
pub async fn flush(&mut self) {
|
||||
let queued_count = self
|
||||
.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Queued)
|
||||
.count();
|
||||
|
||||
info!(
|
||||
"[StreamingExecutor] flush: {} tracked, {} queued, {} executing",
|
||||
self.tracked.len(),
|
||||
queued_count,
|
||||
self.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Executing)
|
||||
.count()
|
||||
);
|
||||
|
||||
// 启动所有还在排队的工具
|
||||
self.start_all_queued();
|
||||
|
||||
// 等待所有执行中的工具完成
|
||||
self.await_all_executing().await;
|
||||
}
|
||||
|
||||
/// 按流顺序获取下一个完成的结果（非阻塞）。
|
||||
///
|
||||
/// 对于已完成的任务，如果其 handle 已就绪则收集结果。
|
||||
/// 返回按到达顺序的第一个已完成结果。
|
||||
pub fn next_result(&mut self) -> Option<(String, ToolOutput)> {
|
||||
// 先尝试收集任何已完成的 async task 结果
|
||||
self.collect_completed_tasks();
|
||||
|
||||
// 从 completed_queue 中按序取
|
||||
while let Some(&idx) = self.completed_queue.front() {
|
||||
self.completed_queue.pop_front();
|
||||
let tool = &mut self.tracked[idx];
|
||||
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)
|
||||
|| !self.completed_queue.is_empty()
|
||||
}
|
||||
|
||||
/// 是否有未完成的工具（仍在排队或执行中）。
|
||||
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)> {
|
||||
self.collect_completed_tasks();
|
||||
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));
|
||||
}
|
||||
tool.status = TrackedToolStatus::Yielded;
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
// ── 内部方法 ──
|
||||
|
||||
/// Spawn 一个 tokio task 执行单个工具调用。
|
||||
fn spawn_tool_task(
|
||||
&self,
|
||||
_idx: usize,
|
||||
_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
) -> JoinHandle<ToolExecutionResult> {
|
||||
let tool_registry = self.tool_registry.clone();
|
||||
let tool_context = self.tool_context.clone();
|
||||
let max_output_chars = self.max_output_chars;
|
||||
let mut abort_rx = self.abort_tx.subscribe();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// tokio::select! 在工具执行和 Sibling Abort 之间竞速
|
||||
tokio::select! {
|
||||
result = async {
|
||||
match tool_registry.get(&tool_name) {
|
||||
Some(tool) => {
|
||||
tool.execute_with_progress(args, &tool_context, None).await
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
}
|
||||
} => {
|
||||
// 截断输出
|
||||
let truncated = if result.content.len() > max_output_chars {
|
||||
let t: String = result.content.chars().take(max_output_chars).collect();
|
||||
ToolOutput {
|
||||
content: format!(
|
||||
"{}...\n[输出已截断，原始长度: {} 字符]",
|
||||
t,
|
||||
result.content.len()
|
||||
),
|
||||
is_error: result.is_error,
|
||||
metadata: result.metadata,
|
||||
skip_persist: result.skip_persist,
|
||||
}
|
||||
} else {
|
||||
result
|
||||
};
|
||||
|
||||
ToolExecutionResult {
|
||||
tool_name,
|
||||
output: truncated,
|
||||
}
|
||||
}
|
||||
Ok(reason) = abort_rx.recv() => {
|
||||
let msg = match reason {
|
||||
AbortReason::SiblingError { description } => {
|
||||
format!("取消：并行工具 {} 出错，已级联取消", description)
|
||||
}
|
||||
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
|
||||
};
|
||||
ToolExecutionResult {
|
||||
tool_name,
|
||||
output: ToolOutput::error(msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// 尝试收集所有已完成 tokio task 的结果（非阻塞）。
|
||||
fn collect_completed_tasks(&mut self) {
|
||||
for idx in 0..self.tracked.len() {
|
||||
if self.tracked[idx].status != TrackedToolStatus::Executing {
|
||||
continue;
|
||||
}
|
||||
if self.tracked[idx].handle.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 检查 JoinHandle 是否已完成（非阻塞）
|
||||
let handle = self.tracked[idx].handle.take().unwrap();
|
||||
if handle.is_finished() {
|
||||
// is_finished=true 保证 .await 会立即返回
|
||||
// 使用 tokio::task::yield_now 之后的 poll 可能也成功，
|
||||
// 这里直接在同步上下文中检查后放入完成队列
|
||||
// 等下次 async 上下文中通过 await_all_executing 处理
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
// 标记为需要收集 — 将在 flush/await 中处理
|
||||
} else {
|
||||
// 放回未完成的 handle
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动所有排队的工具。
|
||||
fn start_all_queued(&mut self) {
|
||||
// 收集需要启动的工具索引（避免借用冲突）
|
||||
let to_start: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.filter(|(_, t)| {
|
||||
let is_safe = self
|
||||
.tool_registry
|
||||
.get(&t.tool_name)
|
||||
.map(|reg_tool| reg_tool.is_concurrency_safe(&t.args))
|
||||
.unwrap_or(false);
|
||||
// 并发安全工具可随时启动，非并发安全的需要独占
|
||||
is_safe || !self.executing_non_concurrent
|
||||
})
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in to_start {
|
||||
let tool = &self.tracked[idx];
|
||||
let call_id = tool.tool_call_id.clone();
|
||||
let tool_name = tool.tool_name.clone();
|
||||
let args = tool.args.clone();
|
||||
|
||||
let is_safe = self
|
||||
.tool_registry
|
||||
.get(&tool_name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let handle = self.spawn_tool_task(idx, call_id, tool_name.clone(), args);
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
|
||||
if !is_safe {
|
||||
self.executing_non_concurrent = true;
|
||||
// 非并发安全工具启动后停止（独占执行）
|
||||
break;
|
||||
}
|
||||
|
||||
info!("[StreamingExecutor] 启动排队工具: {}", tool_name);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等待所有执行中的工具完成。
|
||||
async fn await_all_executing(&mut self) {
|
||||
// 收集所有剩余 JoinHandles
|
||||
let mut handles: Vec<(usize, JoinHandle<ToolExecutionResult>)> = Vec::new();
|
||||
for idx in 0..self.tracked.len() {
|
||||
if self.tracked[idx].status == TrackedToolStatus::Executing {
|
||||
if let Some(handle) = self.tracked[idx].handle.take() {
|
||||
handles.push((idx, handle));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 并发等待所有任务
|
||||
for (idx, handle) in handles {
|
||||
match handle.await {
|
||||
Ok(result) => {
|
||||
let tool_name = result.tool_name.clone();
|
||||
let is_error = result.output.is_error;
|
||||
|
||||
self.tracked[idx].output = Some(result.output);
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
self.completed_queue.push_back(idx);
|
||||
|
||||
if is_error {
|
||||
self.check_sibling_abort(idx, &tool_name);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[StreamingExecutor] tokio task 异常: {}", e);
|
||||
self.tracked[idx].output =
|
||||
Some(ToolOutput::error(format!("工具执行异常: {}", e)));
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
self.completed_queue.push_back(idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.executing_non_concurrent = false;
|
||||
}
|
||||
|
||||
/// 检查错误工具是否触发 Sibling Abort。
|
||||
fn check_sibling_abort(&mut self, idx: usize, tool_name: &str) {
|
||||
let causes_abort = self
|
||||
.tool_registry
|
||||
.get(tool_name)
|
||||
.map(|t| t.causes_sibling_abort())
|
||||
.unwrap_or(false);
|
||||
|
||||
if causes_abort && !self.has_errored {
|
||||
warn!(
|
||||
"[StreamingExecutor] 工具 {} 出错，触发 sibling abort",
|
||||
tool_name
|
||||
);
|
||||
self.has_errored = true;
|
||||
self.errored_tool_desc = self.get_tool_description(idx);
|
||||
let _ = self.abort_tx.send(AbortReason::SiblingError {
|
||||
description: self.errored_tool_desc.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取工具的人类可读描述（用于错误消息）。
|
||||
fn get_tool_description(&self, idx: usize) -> String {
|
||||
let tool = &self.tracked[idx];
|
||||
let summary = tool
|
||||
.args
|
||||
.get("command")
|
||||
.or_else(|| tool.args.get("file_path"))
|
||||
.or_else(|| tool.args.get("pattern"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
if summary.is_empty() {
|
||||
tool.tool_name.clone()
|
||||
} else {
|
||||
let truncated: String = summary.chars().take(40).collect();
|
||||
if summary.len() > 40 {
|
||||
format!("{}({}…)", tool.tool_name, truncated)
|
||||
} else {
|
||||
format!("{}({})", tool.tool_name, summary)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::tools::AgentTool;
|
||||
use crate::agent::tools::ToolContext;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// 可配置的 Mock 工具，用于测试 StreamingToolExecutor 状态机。
|
||||
struct MockAgentTool {
|
||||
name_str: &'static str,
|
||||
concurrency_safe: bool,
|
||||
causes_abort: bool,
|
||||
/// 执行返回的内容
|
||||
result_content: &'static str,
|
||||
/// 执行是否返回错误
|
||||
result_is_error: bool,
|
||||
/// 可选：执行后设置此标志（用于验证工具是否被调用）
|
||||
executed: AtomicBool,
|
||||
}
|
||||
|
||||
impl MockAgentTool {
|
||||
fn new(name: &'static str) -> Self {
|
||||
MockAgentTool {
|
||||
name_str: name,
|
||||
concurrency_safe: false,
|
||||
causes_abort: false,
|
||||
result_content: "mock result",
|
||||
result_is_error: false,
|
||||
executed: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn concurrency_safe(mut self, v: bool) -> Self {
|
||||
self.concurrency_safe = v;
|
||||
self
|
||||
}
|
||||
|
||||
fn causes_abort(mut self, v: bool) -> Self {
|
||||
self.causes_abort = v;
|
||||
self
|
||||
}
|
||||
|
||||
fn result(mut self, content: &'static str, is_error: bool) -> Self {
|
||||
self.result_content = content;
|
||||
self.result_is_error = is_error;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for MockAgentTool {
|
||||
fn name(&self) -> &str {
|
||||
self.name_str
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"mock tool for testing"
|
||||
}
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({})
|
||||
}
|
||||
async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
self.executed.store(true, Ordering::SeqCst);
|
||||
if self.result_is_error {
|
||||
ToolOutput::error(self.result_content)
|
||||
} else {
|
||||
ToolOutput::success(self.result_content, serde_json::json!({}))
|
||||
}
|
||||
}
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
self.concurrency_safe
|
||||
}
|
||||
fn causes_sibling_abort(&self) -> bool {
|
||||
self.causes_abort
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建测试用的 ToolContext。
|
||||
async fn make_test_tool_context() -> ToolContext {
|
||||
use crate::agent::memory::MemoryManager;
|
||||
use crate::agent::skills::SkillRegistry;
|
||||
use crate::api::AppState;
|
||||
use crate::clients::ads::AdsClient;
|
||||
use crate::clients::arxiv::ArxivClient;
|
||||
use crate::clients::cds::vizier::VizierClient;
|
||||
use crate::clients::llm::{EmbeddingClient, LlmClient};
|
||||
use crate::clients::qiniu::QiniuClient;
|
||||
use crate::services::batch::asset::AssetBatchStatus;
|
||||
use crate::services::batch::meta::MetaSyncStatus;
|
||||
use crate::services::download::Downloader;
|
||||
use crate::services::translation::Dictionary;
|
||||
use crate::Config;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
|
||||
sqlx::migrate!("./migrations").run(&pool).await.unwrap();
|
||||
|
||||
let config = Config::from_env();
|
||||
let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into()).unwrap();
|
||||
let embedding =
|
||||
EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into()).unwrap();
|
||||
let ads = AdsClient::new("tk".into()).unwrap();
|
||||
let arxiv = ArxivClient::new().unwrap();
|
||||
let vizier =
|
||||
VizierClient::new("https://tapvizier.cds.unistra.fr/TAPVizieR/tap", 60).unwrap();
|
||||
let qiniu = QiniuClient::new(
|
||||
"ak".into(),
|
||||
"sk".into(),
|
||||
"b".into(),
|
||||
"http://localhost".into(),
|
||||
);
|
||||
|
||||
let app_state = Arc::new(AppState {
|
||||
config,
|
||||
db: pool,
|
||||
dict: Dictionary::default(),
|
||||
qiniu,
|
||||
downloader: Downloader::new().expect("downloader"),
|
||||
http_client: reqwest::Client::new(),
|
||||
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
|
||||
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
|
||||
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from(
|
||||
"/tmp/sk",
|
||||
)))),
|
||||
sse_broadcast: None,
|
||||
memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from(
|
||||
"/tmp/test_mem",
|
||||
)))),
|
||||
login_rate_limiter: Arc::new(dashmap::DashMap::new()),
|
||||
upload_rate_limiter: Arc::new(dashmap::DashMap::new()),
|
||||
llm: crate::api::LlmState {
|
||||
primary: llm.clone(),
|
||||
medium: llm.clone(),
|
||||
fast: llm.clone(),
|
||||
vision: None,
|
||||
embedding,
|
||||
},
|
||||
sources: crate::api::DataSourceState {
|
||||
ads,
|
||||
arxiv,
|
||||
vizier,
|
||||
lamost: crate::clients::lamost::LamostClient::new("https://www.lamost.org", 60)
|
||||
.unwrap(),
|
||||
gaia: crate::clients::gaia::GaiaClient::new(
|
||||
"https://gea.esac.esa.int/tap-server/tap",
|
||||
"https://gea.esac.esa.int/data-server",
|
||||
90,
|
||||
)
|
||||
.unwrap(),
|
||||
sdss: crate::clients::sdss::SdssClient::new(
|
||||
"https://datalab.noirlab.edu/tap/sync",
|
||||
90,
|
||||
)
|
||||
.unwrap(),
|
||||
desi: crate::clients::desi::DesiClient::new(
|
||||
"https://datalab.noirlab.edu/tap/sync",
|
||||
120,
|
||||
)
|
||||
.unwrap(),
|
||||
irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60)
|
||||
.unwrap(),
|
||||
mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(),
|
||||
observation_registry: std::sync::Arc::new(
|
||||
crate::services::observation::ObservationRegistry::default(),
|
||||
),
|
||||
},
|
||||
session: crate::api::SessionState {
|
||||
sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
|
||||
session_last_active: Arc::new(dashmap::DashMap::new()),
|
||||
cancelled_runs: Arc::new(dashmap::DashMap::new()),
|
||||
session_permission_checkers: Arc::new(dashmap::DashMap::new()),
|
||||
pending_questions: Arc::new(tokio::sync::Mutex::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
pending_permissions: Arc::new(tokio::sync::Mutex::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
},
|
||||
});
|
||||
|
||||
ToolContext::new(app_state)
|
||||
}
|
||||
|
||||
/// 创建包含指定 Mock 工具的 StreamingToolExecutor。
|
||||
async fn make_executor(tools: Vec<MockAgentTool>) -> StreamingToolExecutor {
|
||||
let skill_registry = Arc::new(tokio::sync::RwLock::new(
|
||||
crate::agent::skills::SkillRegistry::new(std::path::PathBuf::from("/tmp/sk")),
|
||||
));
|
||||
let mut registry = ToolRegistry::new(skill_registry);
|
||||
for tool in tools {
|
||||
registry.add_tool(Box::new(tool));
|
||||
}
|
||||
let tool_context = make_test_tool_context().await;
|
||||
StreamingToolExecutor::new(Arc::new(registry), tool_context, 4, 4000)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_spawns_concurrency_safe_immediately() {
|
||||
let mut executor =
|
||||
make_executor(vec![MockAgentTool::new("safe_tool").concurrency_safe(true)]).await;
|
||||
|
||||
let spawned = executor.on_tool_use(
|
||||
"call_1".into(),
|
||||
"safe_tool".into(),
|
||||
serde_json::json!({"key": "val"}),
|
||||
);
|
||||
assert!(spawned, "concurrency-safe tool should spawn immediately");
|
||||
assert_eq!(executor.tracked.len(), 1);
|
||||
assert_eq!(
|
||||
executor.tracked[0].status,
|
||||
TrackedToolStatus::Executing,
|
||||
"safe tool should be executing"
|
||||
);
|
||||
assert!(
|
||||
!executor.executing_non_concurrent,
|
||||
"safe tool does not lock executor"
|
||||
);
|
||||
assert!(
|
||||
executor.has_unfinished(),
|
||||
"should have unfinished tool (executing)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_queues_non_concurrency_safe() {
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("unsafe_tool").concurrency_safe(false)
|
||||
])
|
||||
.await;
|
||||
|
||||
let spawned =
|
||||
executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({}));
|
||||
assert!(!spawned, "non-concurrency-safe tool should be queued");
|
||||
assert_eq!(executor.tracked.len(), 1);
|
||||
assert_eq!(executor.tracked[0].status, TrackedToolStatus::Queued);
|
||||
assert!(
|
||||
executor.executing_non_concurrent,
|
||||
"non-concurrent tool sets the lock flag"
|
||||
);
|
||||
assert!(
|
||||
executor.has_unfinished(),
|
||||
"queued tool counts as unfinished"
|
||||
);
|
||||
assert!(!executor.has_pending_results(), "no results yet");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_queues_safe_tool_when_non_concurrent_executing() {
|
||||
// 先加入一个非并发安全工具（设为 executing_non_concurrent=true），
|
||||
// 再尝试加入一个并发安全工具，应排队而非立即执行
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("unsafe_tool").concurrency_safe(false),
|
||||
MockAgentTool::new("safe_tool").concurrency_safe(true),
|
||||
])
|
||||
.await;
|
||||
|
||||
// 第一个：非并发安全，排队但设置 executing_non_concurrent
|
||||
executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({}));
|
||||
// 第二个：并发安全但 executor 被非并发工具锁定，应排队
|
||||
let spawned =
|
||||
executor.on_tool_use("call_2".into(), "safe_tool".into(), serde_json::json!({}));
|
||||
|
||||
assert!(
|
||||
!spawned,
|
||||
"safe tool should be queued when executor is locked"
|
||||
);
|
||||
assert_eq!(executor.tracked.len(), 2);
|
||||
assert_eq!(executor.tracked[1].status, TrackedToolStatus::Queued);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_tool_description_parses_args() {
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("bash")]).await;
|
||||
// 手动添加 tracked tool（绕过 on_tool_use 的 spawn）
|
||||
executor.tracked.push(TrackedTool {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "bash".into(),
|
||||
args: serde_json::json!({"command": "git push origin main"}),
|
||||
status: TrackedToolStatus::Completed,
|
||||
output: None,
|
||||
handle: None,
|
||||
});
|
||||
|
||||
let desc = executor.get_tool_description(0);
|
||||
assert!(
|
||||
desc.contains("bash"),
|
||||
"description should contain tool name, got: {}",
|
||||
desc
|
||||
);
|
||||
assert!(
|
||||
desc.contains("git push"),
|
||||
"description should contain command arg, got: {}",
|
||||
desc
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_tool_description_falls_back_to_name() {
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("unknown_tool")]).await;
|
||||
executor.tracked.push(TrackedTool {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "unknown_tool".into(),
|
||||
args: serde_json::json!({}), // no command/file_path/pattern
|
||||
status: TrackedToolStatus::Completed,
|
||||
output: None,
|
||||
handle: None,
|
||||
});
|
||||
|
||||
let desc = executor.get_tool_description(0);
|
||||
assert_eq!(desc, "unknown_tool");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_flush_starts_queued_and_awaits_completion() {
|
||||
// 使用并发安全工具测试 flush：on_tool_use 立即 spawn，flush 等待完成
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.result("done", false)])
|
||||
.await;
|
||||
|
||||
let spawned = executor.on_tool_use("call_1".into(), "tool_a".into(), serde_json::json!({}));
|
||||
assert!(spawned, "concurrency-safe tool should spawn immediately");
|
||||
assert_eq!(executor.tracked[0].status, TrackedToolStatus::Executing);
|
||||
|
||||
executor.flush().await;
|
||||
|
||||
assert_eq!(
|
||||
executor.tracked[0].status,
|
||||
TrackedToolStatus::Completed,
|
||||
"after flush, tool should be completed"
|
||||
);
|
||||
assert!(
|
||||
executor.has_pending_results(),
|
||||
"completed tool = pending result"
|
||||
);
|
||||
assert!(!executor.has_unfinished(), "nothing queued or executing");
|
||||
|
||||
let result = executor.next_result();
|
||||
assert!(
|
||||
result.is_some(),
|
||||
"next_result should return the completed tool"
|
||||
);
|
||||
let (call_id, output) = result.unwrap();
|
||||
assert_eq!(call_id, "call_1");
|
||||
assert!(!output.is_error, "tool should succeed");
|
||||
assert_eq!(output.content, "done");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_next_result_yields_in_insertion_order() {
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.result("result_a", false),
|
||||
MockAgentTool::new("tool_b")
|
||||
.concurrency_safe(true)
|
||||
.result("result_b", false),
|
||||
])
|
||||
.await;
|
||||
|
||||
// 两个并发安全工具，都应该立即 spawn
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({}));
|
||||
|
||||
// flush 等待它们完成
|
||||
executor.flush().await;
|
||||
|
||||
// 结果应按插入顺序产出
|
||||
let result_a = executor.next_result();
|
||||
assert!(result_a.is_some());
|
||||
assert_eq!(result_a.as_ref().unwrap().0, "call_a");
|
||||
assert_eq!(result_a.as_ref().unwrap().1.content, "result_a");
|
||||
|
||||
let result_b = executor.next_result();
|
||||
assert!(result_b.is_some());
|
||||
assert_eq!(result_b.as_ref().unwrap().0, "call_b");
|
||||
assert_eq!(result_b.as_ref().unwrap().1.content, "result_b");
|
||||
|
||||
// 第三次调用返回 None（全部已 yield）
|
||||
let result_c = executor.next_result();
|
||||
assert!(result_c.is_none());
|
||||
assert!(!executor.has_pending_results());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sibling_abort_broadcast_on_error() {
|
||||
// 注册两个工具：tool_a 会报错且触发 sibling abort，tool_b 并发执行中被取消
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.causes_abort(true)
|
||||
.result("critical failure", true),
|
||||
MockAgentTool::new("tool_b")
|
||||
.concurrency_safe(true)
|
||||
.result("should be aborted", false),
|
||||
])
|
||||
.await;
|
||||
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({}));
|
||||
|
||||
executor.flush().await;
|
||||
|
||||
// tool_a 的结果应该是错误
|
||||
let result_a = executor.next_result();
|
||||
assert!(result_a.is_some());
|
||||
assert!(result_a.unwrap().1.is_error, "tool_a should have errored");
|
||||
|
||||
// tool_b 可能被取消（sibling abort）或正常完成（取决于竞态）
|
||||
let result_b = executor.next_result();
|
||||
assert!(result_b.is_some(), "tool_b should also have a result");
|
||||
|
||||
// 验证 has_errored 被设置
|
||||
assert!(
|
||||
executor.has_errored,
|
||||
"has_errored should be set after sibling abort"
|
||||
);
|
||||
assert!(
|
||||
!executor.errored_tool_desc.is_empty(),
|
||||
"errored_tool_desc should be populated"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_all_results_mut_drains_all_outputs() {
|
||||
let mut executor =
|
||||
make_executor(vec![MockAgentTool::new("tool_a").result("result_a", false)]).await;
|
||||
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
// 不 flush — all_results_mut 是同步方法，只收集已完成的结果
|
||||
// 工具可能仍在执行，所以不能保证一定有结果
|
||||
let results = executor.all_results_mut();
|
||||
// 无论有没有结果，调用后 tracked tool 被标记为 Yielded
|
||||
for tool in &executor.tracked {
|
||||
assert_eq!(tool.status, TrackedToolStatus::Yielded);
|
||||
assert!(tool.output.is_none(), "output should be taken");
|
||||
}
|
||||
// 验证返回的 results 和 tracked 一致
|
||||
let _ = results; // 如果 joinhandle 还没完成，results 可能是空的
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_abort_sender_clone_works() {
|
||||
let executor = make_executor(vec![]).await;
|
||||
let sender = executor.abort_sender();
|
||||
// 验证 sender 可用
|
||||
assert_eq!(sender.receiver_count(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_executor_has_no_pending_or_unfinished() {
|
||||
let executor = make_executor(vec![]).await;
|
||||
assert!(!executor.has_pending_results());
|
||||
assert!(!executor.has_unfinished());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user