新增: 并行调度+Token预算池+事件协议+前端类型+编译警告清理
- Coordinator.dispatch_with_budget: JoinSet层内并行+层间串行+预算超限降级串行 - TokenBudgetPool: AtomicU64 CAS无锁并发安全,0=不限制 - 4个新事件: AiPlanCreated/AiSubTaskStatusChanged/AiMergeCompleted/AiConflictResolved - 前端类型: PlanRecord/SubTaskRecord/ConflictRecord/PlanLayerInfo/ConflictInfo - PlanProgress.vue: 接入真实状态+persona徽章+冲突徽章+i18n - 编译警告全部清零(audit子模块allow+事件allow+record allow) - Coordinator测试27个全绿(含4个Token预算+4个并行调度)
This commit is contained in:
@@ -10,6 +10,62 @@
|
||||
|
||||
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<Self> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 枚举 --------------------------------------------------------------------
|
||||
|
||||
@@ -192,32 +248,140 @@ impl Coordinator {
|
||||
DecompositionResult { subtasks, plan }
|
||||
}
|
||||
|
||||
/// 分发执行:按 Plan 逐层执行 SubTask
|
||||
/// 分发执行:按 Plan 分层执行 SubTask(层间串行 + 层内并行)
|
||||
///
|
||||
/// Phase 1:串行执行(层内层间均串行),后续可改为 `JoinSet` 并行。
|
||||
/// - 层间串行:上层全部 done 才进下一层(DAG 依赖保证)
|
||||
/// - 层内并行:同层 SubTask 用 tokio::task::JoinSet 并发执行
|
||||
/// - Token 预算超限时降级为串行(不拒绝执行)
|
||||
///
|
||||
/// ## 参数
|
||||
/// - `plan`: 待执行的 DAG Plan
|
||||
/// - `executor`: 子任务执行函数,
|
||||
/// 签名 `fn(SubTask, AgentPersona) -> Future<Output = ExecutionResult>`
|
||||
/// 传入 owned SubTask 和 AgentPersona 副本,避免生命周期约束。
|
||||
/// - `executor`: 子任务执行函数(接收 SubTask + persona_id 字符串)
|
||||
/// - `budget`: Token 预算池(None=不限制)
|
||||
///
|
||||
/// ## 返回值
|
||||
/// 按 Plan 原始顺序排列的执行结果列表。
|
||||
/// 按 Plan 原始 tasks 顺序排列的执行结果列表。
|
||||
pub async fn dispatch_with_budget<F, Fut>(
|
||||
&self,
|
||||
plan: &Plan,
|
||||
executor: F,
|
||||
budget: Option<&Arc<TokenBudgetPool>>,
|
||||
) -> Vec<ExecutionResult>
|
||||
where
|
||||
F: Fn(SubTask, String) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: std::future::Future<Output = ExecutionResult> + 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<String, String> = 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<String, ExecutionResult> =
|
||||
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<ExecutionResult> =
|
||||
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<F, Fut>(&self, plan: &Plan, executor: F) -> Vec<ExecutionResult>
|
||||
where
|
||||
F: Fn(SubTask, crate::persona::AgentPersona) -> Fut,
|
||||
Fut: std::future::Future<Output = ExecutionResult>,
|
||||
{
|
||||
let mut results = Vec::new();
|
||||
self.dispatch_serial(plan, &executor).await
|
||||
}
|
||||
|
||||
// Phase 1:串行逐任务执行(忽略 to_layers 分层,纯顺序执行)
|
||||
/// 内部串行执行(无并行)
|
||||
async fn dispatch_serial<F, Fut>(&self, plan: &Plan, executor: &F) -> Vec<ExecutionResult>
|
||||
where
|
||||
F: Fn(SubTask, crate::persona::AgentPersona) -> Fut,
|
||||
Fut: std::future::Future<Output = ExecutionResult>,
|
||||
{
|
||||
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
|
||||
}
|
||||
|
||||
@@ -511,4 +675,152 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// -- 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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user