//! 多 Agent 协作调度中心 — Coordinator //! //! Phase 1 实现,Phase 2 LLM 扩展 //! //! 职责拆解(4 步): //! 1. **decompose**:接收用户意图 + 消息 → 拆解为 SubTask 列表 //! 2. **分配**:为每个 SubTask 分配人设(PersonaRegistry.recommend_for_intent) //! 3. **dispatch**:并行执行 SubTask(当前 Phase 1 串行,预留 JoinSet 并行) //! 4. **merge**:汇总子结果 → 合并产出 → 处理冲突 use crate::persona::PersonaRegistry; use crate::planner::{Plan, SubTask}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; // ---- Token 预算池 ------------------------------------------------------------ /// 全局 Token 预算池(CAS 无锁并发安全) /// 多个 SubTask 启动前向预算池申请估算额度,超限时降级串行(不拒绝执行)。 #[derive(Debug)] pub struct TokenBudgetPool { total: AtomicU64, consumed: AtomicU64, } impl TokenBudgetPool { /// 创建预算池。total=0 表示不限制(等价无限)。 pub fn new(total: u64) -> Arc { Arc::new(Self { total: AtomicU64::new(total), consumed: AtomicU64::new(0), }) } /// 尝试预占额度。成功返回 true,超限返回 false。 /// total=0 时不限制,始终返回 true。 pub fn try_reserve(&self, estimate: u64) -> bool { let total = self.total.load(Ordering::SeqCst); if total == 0 { return true; } let mut consumed = self.consumed.load(Ordering::SeqCst); loop { if consumed + estimate > total { return false; } match self.consumed.compare_exchange_weak( consumed, consumed + estimate, Ordering::SeqCst, Ordering::SeqCst, ) { Ok(_) => return true, Err(actual) => consumed = actual, } } } /// 已消耗额度 pub fn consumed(&self) -> u64 { self.consumed.load(Ordering::SeqCst) } /// 总预算(0=不限制) pub fn total(&self) -> u64 { self.total.load(Ordering::SeqCst) } } // ---- 枚举 -------------------------------------------------------------------- /// 调度策略枚举 pub enum DispatchStrategy { /// 规则驱动(关键词匹配,Phase 1 MVP) RuleBased, /// LLM 驱动(LLM 生成 Plan,Phase 2 预留) LLMDriven, } impl DispatchStrategy { /// 策略名称标签(日志 / 审计用) pub fn name(&self) -> &'static str { match self { Self::RuleBased => "rule_based", Self::LLMDriven => "llm_driven", } } } // ---- 数据结构 ---------------------------------------------------------------- /// 拆解结果 pub struct DecompositionResult { /// 拆解出的子任务列表 pub subtasks: Vec, /// 对应 Plan 结构(含依赖关系) pub plan: Plan, } /// 执行结果 pub struct ExecutionResult { /// 对应 SubTask id pub subtask_id: String, /// 分配的人设 id pub persona_id: String, /// 执行产出(工具结果摘要 / LLM 回复) pub output: String, /// 执行是否成功 pub success: bool, } /// 合并结果 pub struct MergeResult { /// 合并后的最终产出 pub merged_output: String, /// 检测到的冲突列表(Phase 2 由 reviewer Agent 填充) pub conflicts: Vec, } /// 冲突项 pub struct ConflictItem { /// 冲突涉及的文件 pub file: String, /// 冲突描述 pub description: String, } // ---- Coordinator 实现 -------------------------------------------------------- /// 多 Agent 协作调度器 /// /// # Phase 1(当前) /// /// 规则驱动 decompose + 串行 dispatch + 简单拼接 merge。 /// 适用于「读后写」「搜索后分析」等可预测的意图组合。 /// /// # Phase 2(预留) /// /// - LLM 驱动 decompose(LLM 生成 Plan 结构) /// - JoinSet 并行 dispatch(层内并行) /// - Reviewer Agent 仲裁 merge(冲突检测 + 自动解决) pub struct Coordinator { /// 人设注册表(内置 5 人设 + 自定义) registry: PersonaRegistry, } impl Coordinator { /// 创建 Coordinator,绑定人设注册表 pub fn new(registry: PersonaRegistry) -> Self { Self { registry } } /// 规则拆解:intent + text → SubTask 列表 + Plan /// /// Phase 1 规则(关键词匹配): /// - 检测 text 中是否含 `read`/`写`/`查看`/`分析`/`search` 等关键词 → 读子任务 /// - 检测 text 中是否含 `write`/`修改`/`生成`/`create`/`fix` 等关键词 → 写子任务 /// - 同时含读+写关键词 → 两个子任务,写依赖读(read → write) /// - 每 SubTask 分配人设(`PersonaRegistry::recommend_for_intent`) /// - 兜底:无法匹配任何关键词时产单个 default 子任务 /// /// ## Phase 2 扩展点 /// `DispatchStrategy::LLMDriven` 时改为调 LLM 生成 Plan 整体结构, /// 本方法当前仅走 `RuleBased` 路径。 pub fn decompose(&self, intent: &str, text: &str) -> DecompositionResult { let lower = text.to_lowercase(); // 规则 1: 读操作关键词 let has_read = lower.contains("read") || lower.contains("查看") || lower.contains("读取") || lower.contains("分析") || lower.contains("search") || lower.contains("搜索") || lower.contains("审查") || lower.contains("review"); // 规则 2: 写操作关键词 let has_write = lower.contains("write") || lower.contains("写") || lower.contains("修改") || lower.contains("生成") || lower.contains("创建") || lower.contains("implement") || lower.contains("fix") || lower.contains("修复") || lower.contains("新增"); let mut subtasks = Vec::new(); match (has_read, has_write) { (true, true) => { // 读 + 写 → 两个子任务:read → write subtasks.push( SubTask::new("read", "读取/分析现有代码") .with_tools(vec![ "read_file".into(), "search_files".into(), "list_directory".into(), ]) .with_deps(vec![]), ); subtasks.push( SubTask::new("write", "修改/生成代码") .with_tools(vec![ "write_file".into(), "patch_file".into(), "edit_file".into(), ]) .with_deps(vec!["read".into()]), ); } (true, false) => { // 纯读 → 单个读子任务 subtasks.push( SubTask::new("read", "读取/分析代码") .with_tools(vec![ "read_file".into(), "search_files".into(), "list_directory".into(), ]) .with_deps(vec![]), ); } (false, true) => { // 纯写 → 单个写子任务 subtasks.push( SubTask::new("write", "修改/生成代码") .with_tools(vec![ "write_file".into(), "patch_file".into(), "edit_file".into(), ]) .with_deps(vec![]), ); } (false, false) => { // 兜底:按 intent 分配单个默认子任务 subtasks.push( SubTask::new("default", intent.to_string()) .with_tools(vec![]) .with_deps(vec![]), ); } } let plan = Plan::from_tasks(subtasks.clone()); DecompositionResult { subtasks, plan } } /// 分发执行:按 Plan 分层执行 SubTask(层间串行 + 层内并行) /// /// - 层间串行:上层全部 done 才进下一层(DAG 依赖保证) /// - 层内并行:同层 SubTask 用 tokio::task::JoinSet 并发执行 /// - Token 预算超限时降级为串行(不拒绝执行) /// /// ## 参数 /// - `plan`: 待执行的 DAG Plan /// - `executor`: 子任务执行函数(接收 SubTask + persona_id 字符串) /// - `budget`: Token 预算池(None=不限制) /// /// ## 返回值 /// 按 Plan 原始 tasks 顺序排列的执行结果列表。 pub async fn dispatch_with_budget( &self, plan: &Plan, executor: F, budget: Option<&Arc>, ) -> Vec where F: Fn(SubTask, String) -> Fut + Clone + Send + Sync + 'static, Fut: std::future::Future + Send, { let layers = match plan.to_layers() { Ok(l) => l, Err(_) => { tracing::warn!("[COORDINATOR] DAG 分层失败,降级串行"); let mut results = Vec::new(); for task in &plan.tasks { let persona_id = self.registry.recommend_for_intent(&task.intent).id.clone(); let result = executor(task.clone(), persona_id).await; results.push(result); } return results; } }; // 预分配每个 SubTask 的 persona_id(避免在 JoinSet 内借用 self) let task_persona: std::collections::HashMap = plan .tasks .iter() .map(|t| { (t.id.clone(), self.registry.recommend_for_intent(&t.intent).id.clone()) }) .collect(); let mut results_map: std::collections::HashMap = std::collections::HashMap::new(); for (layer_idx, layer) in layers.iter().enumerate() { if layer.len() <= 1 { for task in layer { let pid = task_persona.get(&task.id).cloned().unwrap_or_default(); let result = executor(task.clone(), pid).await; results_map.insert(result.subtask_id.clone(), result); } continue; } let can_parallel = budget .map(|b| b.try_reserve(10_000)) .unwrap_or(true); if can_parallel { tracing::info!( layer = layer_idx, count = layer.len(), "[COORDINATOR] 层 {} 并行执行 {} 个子任务", layer_idx, layer.len() ); let mut join_set: tokio::task::JoinSet = tokio::task::JoinSet::new(); for task in layer { let pid = task_persona.get(&task.id).cloned().unwrap_or_default(); let exec = executor.clone(); let task = task.clone(); join_set.spawn(async move { exec(task, pid).await }); } while let Some(join_result) = join_set.join_next().await { match join_result { Ok(result) => { results_map.insert(result.subtask_id.clone(), result); } Err(e) => { tracing::error!("[COORDINATOR] 子任务 panic: {}", e); } } } } else { tracing::info!( layer = layer_idx, "[COORDINATOR] Token 预算超限,层 {} 降级串行", layer_idx ); for task in layer { let pid = task_persona.get(&task.id).cloned().unwrap_or_default(); let result = executor(task.clone(), pid).await; results_map.insert(result.subtask_id.clone(), result); } } } plan.tasks .iter() .filter_map(|t| results_map.remove(&t.id)) .collect() } /// 串行 dispatch(降级/Phase 1 兼容路径) pub async fn dispatch(&self, plan: &Plan, executor: F) -> Vec where F: Fn(SubTask, crate::persona::AgentPersona) -> Fut, Fut: std::future::Future, { self.dispatch_serial(plan, &executor).await } /// 内部串行执行(无并行) async fn dispatch_serial(&self, plan: &Plan, executor: &F) -> Vec where F: Fn(SubTask, crate::persona::AgentPersona) -> Fut, Fut: std::future::Future, { let mut results = Vec::new(); for task in &plan.tasks { let persona = self.registry.recommend_for_intent(&task.intent); let result = executor(task.clone(), persona.clone()).await; results.push(result); } results } /// 合并:汇总子结果 → 合并产出 /// /// Phase 1:简单拼接各 SubTask 产出,每段带标题标识来源。 /// `conflicts` 返回空 vec(Phase 2 由 reviewer Agent 仲裁填充)。 pub fn merge(&self, results: &[ExecutionResult]) -> MergeResult { let parts: Vec = results .iter() .map(|r| { format!( "## SubTask: {} (Persona: {})\n\n{}", r.subtask_id, r.persona_id, r.output ) }) .collect(); MergeResult { merged_output: parts.join("\n\n---\n\n"), conflicts: Vec::new(), } } } // ---- 单元测试 --------------------------------------------------------------- #[cfg(test)] mod tests { use super::*; use crate::persona::{PersonaRegistry, PERSONA_CODER}; fn make_coord() -> Coordinator { Coordinator::new(PersonaRegistry::new()) } // -- decompose 路径覆盖 -- #[test] fn decompose_read_then_write_chinese() { let coord = make_coord(); let result = coord.decompose("modify", "请读取当前代码,然后修改它"); assert_eq!(result.subtasks.len(), 2, "读+写应拆为两个子任务"); assert_eq!(result.subtasks[0].id, "read", "第一个子任务应为读"); assert_eq!(result.subtasks[1].id, "write", "第二个子任务应为写"); assert!( result.subtasks[1].deps.contains(&"read".to_string()), "写子任务应依赖读子任务" ); assert_eq!(result.plan.len(), 2); } #[test] fn decompose_read_then_write_english() { let coord = make_coord(); let result = coord.decompose("modify", "read the file and then write changes"); assert_eq!(result.subtasks.len(), 2); assert_eq!(result.subtasks[0].id, "read"); assert_eq!(result.subtasks[1].id, "write"); assert!(result.subtasks[1].deps.contains(&"read".to_string())); } #[test] fn decompose_read_only_chinese() { let coord = make_coord(); let result = coord.decompose("review", "帮我查看这段代码有什么问题"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "read"); } #[test] fn decompose_read_only_search() { let coord = make_coord(); let result = coord.decompose("search", "搜索项目中的 TODO 注释"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "read"); } #[test] fn decompose_write_only_generate() { let coord = make_coord(); let result = coord.decompose("implement", "请生成一个 Rust 模块"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "write"); } #[test] fn decompose_write_only_fix() { let coord = make_coord(); let result = coord.decompose("fix", "修复这个 bug"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "write"); } #[test] fn decompose_fallback_default() { let coord = make_coord(); let result = coord.decompose("chat", "你好,今天天气怎么样"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "default"); } #[test] fn decompose_create_keyword() { let coord = make_coord(); let result = coord.decompose("create", "创建新的 API 端点"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "write"); } // -- dispatch -- #[test] fn dispatch_serial_execution() { let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("read", "read code").with_deps(vec![]), SubTask::new("write", "write code").with_deps(vec!["read".into()]), ]); let results = futures::executor::block_on(coord.dispatch(&plan, |task, persona| { async move { ExecutionResult { subtask_id: task.id, persona_id: persona.id, output: format!("执行 {} 完毕", task.intent), success: true, } } })); assert_eq!(results.len(), 2); assert!(results[0].success); assert!(results[1].success); assert_eq!(results[0].subtask_id, "read"); assert_eq!(results[1].subtask_id, "write"); assert_eq!(results[0].persona_id, PERSONA_CODER); } #[test] fn dispatch_empty_plan() { let coord = make_coord(); let plan = Plan::new(); let results = futures::executor::block_on(coord.dispatch(&plan, |task, persona| { async move { ExecutionResult { subtask_id: task.id, persona_id: persona.id, output: String::new(), success: true, } } })); assert!(results.is_empty(), "空 Plan 应产生空结果"); } #[test] fn dispatch_persona_assigned_by_intent() { let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("code", "implement feature").with_deps(vec![]), ]); let results = futures::executor::block_on(coord.dispatch(&plan, |task, persona| { async move { ExecutionResult { subtask_id: task.id, persona_id: persona.id, output: String::new(), success: true, } } })); // "implement" 不触发特殊人设 → 默认 coder assert_eq!(results[0].persona_id, PERSONA_CODER); } // -- merge -- #[test] fn merge_single_result() { let coord = make_coord(); let results = vec![ExecutionResult { subtask_id: "read".into(), persona_id: "reviewer".into(), output: "分析完成,发现 3 个问题".into(), success: true, }]; let merged = coord.merge(&results); assert!(merged.merged_output.contains("分析完成")); assert!(merged.conflicts.is_empty(), "Phase 1 不应有冲突"); } #[test] fn merge_multiple_results() { let coord = make_coord(); let results = vec![ ExecutionResult { subtask_id: "read".into(), persona_id: "reviewer".into(), output: "代码分析结果".into(), success: true, }, ExecutionResult { subtask_id: "write".into(), persona_id: "coder".into(), output: "修改后的代码".into(), success: true, }, ]; let merged = coord.merge(&results); assert!(merged.merged_output.contains("代码分析结果")); assert!(merged.merged_output.contains("修改后的代码")); assert!(merged.merged_output.contains("SubTask: read")); assert!(merged.merged_output.contains("SubTask: write")); assert!(merged.conflicts.is_empty(), "Phase 1 不应有冲突"); } #[test] fn merge_empty_results() { let coord = make_coord(); let merged = coord.merge(&[]); assert!(merged.merged_output.is_empty(), "空输入应产生空输出"); assert!(merged.conflicts.is_empty()); } // -- DispatchStrategy -- #[test] fn dispatch_strategy_name() { assert_eq!(DispatchStrategy::RuleBased.name(), "rule_based"); assert_eq!(DispatchStrategy::LLMDriven.name(), "llm_driven"); } // -- Edge cases: 关键词混合边界 -- #[test] fn decompose_review_triggers_read() { let coord = make_coord(); // "review" 同时触发读关键词 let result = coord.decompose("review", "review this PR for me"); assert_eq!(result.subtasks.len(), 1); assert_eq!(result.subtasks[0].id, "read"); } #[test] fn decompose_create_and_review_triggers_both() { let coord = make_coord(); let result = coord.decompose("implement", "请先分析代码结构然后创建新模块"); assert_eq!( result.subtasks.len(), 2, "含分析和创建关键词应拆为两个子任务" ); } #[test] fn decompose_plan_not_empty() { let coord = make_coord(); let result = coord.decompose("chat", "随便聊聊"); assert!(!result.plan.is_empty(), "兜底场景也应有 Plan"); } #[test] fn decompose_result_holds_plan_consistent() { let coord = make_coord(); let result = coord.decompose("modify", "读取配置文件然后更新它"); assert_eq!( result.subtasks.len(), result.plan.len(), "DecompositionResult 的 subtasks 应与 plan 一致" ); for st in &result.subtasks { assert!( result.plan.tasks.iter().any(|t| t.id == st.id), "plan 应包含所有 subtask: {}", st.id ); } } // -- Token 预算池 -- #[test] fn tok_01_budget_sufficient() { let pool = TokenBudgetPool::new(100_000); assert!(pool.try_reserve(10_000), "预算充足时应允许"); assert_eq!(pool.consumed(), 10_000); assert!(pool.try_reserve(50_000), "剩余充足时应允许"); assert_eq!(pool.consumed(), 60_000); } #[test] fn tok_02_budget_exceeded() { let pool = TokenBudgetPool::new(100_000); assert!(pool.try_reserve(98_000)); assert!(!pool.try_reserve(5_000), "超限应拒绝"); assert_eq!(pool.consumed(), 98_000, "拒绝后 consumed 不增"); } #[test] fn tok_03_concurrent_reserve() { let pool = TokenBudgetPool::new(15_000); // 模拟两个并发申请各 10k,总额 20k > 15k,只有第一个应成功 let pool1 = pool.clone(); let pool2 = pool.clone(); let r1 = pool1.try_reserve(10_000); let r2 = pool2.try_reserve(10_000); // 至少一个成功一个失败(顺序由调度决定) assert!(r1 || r2, "至少一个成功"); assert!(!r1 || !r2, "不能两个都成功(总超预算)"); } #[test] fn tok_04_zero_means_unlimited() { let pool = TokenBudgetPool::new(0); assert!(pool.try_reserve(999_999_999), "total=0 不限制"); } // -- 并行调度 -- #[tokio::test] async fn par_01_diamond_dependency() { // A → {B, C} → D,B/C 同层应并行,D 等 B+C 都 done let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("A", "base"), SubTask::new("B", "branch1").with_deps(vec!["A".into()]), SubTask::new("C", "branch2").with_deps(vec!["A".into()]), SubTask::new("D", "final").with_deps(vec!["B".into(), "C".into()]), ]); let results = coord.dispatch_with_budget(&plan, |task, persona_id: String| { async move { ExecutionResult { subtask_id: task.id, persona_id, output: format!("executed {}", task.intent), success: true, } } }, None).await; assert_eq!(results.len(), 4, "全部 4 个子任务应有结果"); // 结果按 plan.tasks 原始顺序 assert_eq!(results[0].subtask_id, "A"); assert_eq!(results[1].subtask_id, "B"); assert_eq!(results[2].subtask_id, "C"); assert_eq!(results[3].subtask_id, "D"); } #[tokio::test] async fn par_02_parallel_layer_executed() { // 3 个无依赖任务应在同层并行执行 let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("t1", "read file1"), SubTask::new("t2", "read file2"), SubTask::new("t3", "read file3"), ]); let results = coord.dispatch_with_budget(&plan, |task, persona_id: String| { async move { ExecutionResult { subtask_id: task.id.clone(), persona_id, output: format!("done {}", task.id), success: true, } } }, None).await; assert_eq!(results.len(), 3); for r in &results { assert!(r.success, "所有子任务应成功"); } } #[tokio::test] async fn par_03_budget_exceeded_degrades_to_serial() { // Token 预算仅够 1 个并行,第 2 个应降级串行(不拒绝) let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("t1", "task1"), SubTask::new("t2", "task2"), SubTask::new("t3", "task3"), ]); let pool = TokenBudgetPool::new(5_000); // 不够 3×10k let results = coord.dispatch_with_budget(&plan, |task, persona_id: String| { async move { ExecutionResult { subtask_id: task.id, persona_id, output: String::new(), success: true, } } }, Some(&pool)).await; // 预算超限应降级串行,不拒绝执行 assert_eq!(results.len(), 3, "降级串行也应全部执行"); } #[tokio::test] async fn par_04_single_task_layer_serial() { // 线性链(每层 1 任务)不走并行路径 let coord = make_coord(); let plan = Plan::from_tasks(vec![ SubTask::new("a", "step1"), SubTask::new("b", "step2").with_deps(vec!["a".into()]), ]); let results = coord.dispatch_with_budget(&plan, |task, persona_id: String| { async move { ExecutionResult { subtask_id: task.id, persona_id, output: String::new(), success: true, } } }, None).await; assert_eq!(results.len(), 2); assert_eq!(results[0].subtask_id, "a"); assert_eq!(results[1].subtask_id, "b"); } }