DevFlow/crates/df-ai/src/coordinator.rs

//! 多 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<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)
    }
}

// ---- 枚举 --------------------------------------------------------------------

/// 调度策略枚举
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<SubTask>,
    /// 对应 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<ConflictItem>,
}

/// 冲突项
pub struct ConflictItem {
    /// 冲突涉及的文件
    pub file: String,
    /// 冲突描述
    pub description: String,
    /// 冲突双方 SubTask id
    pub subtask_a: Option<String>,
    pub subtask_b: Option<String>,
    /// Reviewer 仲裁推荐(Phase 2 LLM 仲裁填充)
    pub recommendation: Option<ConflictResolution>,
}

/// 冲突解决方案
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConflictResolution {
    /// 接受 A 的改动
    AcceptA,
    /// 接受 B 的改动
    AcceptB,
    /// 合并两者
    Merged,
    /// 用户手动处理
    Manual,
}

impl ConflictResolution {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::AcceptA => "a",
            Self::AcceptB => "b",
            Self::Merged => "merged",
            Self::Manual => "manual",
        }
    }
}

// ---- Coordinator 实现 --------------------------------------------------------

/// 多 Agent 协作调度器
///
/// 当前实现:
/// - 规则驱动 decompose(关键词匹配)
/// - dispatch_with_budget: JoinSet 层内并行 + token 预算管控
/// - dispatch_serial: 串行降级路径(无 'static 约束)
/// - merge: 拼接产出 + 同文件冲突检测
///
/// Phase 2 预留:
/// - LLM 驱动 decompose
/// - Reviewer Agent 仲裁
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<F, Fut>(
        &self,
        plan: &Plan,
        executor: F,
        budget: Option<&Arc<TokenBudgetPool>>,
    ) -> Vec<ExecutionResult>
    where
        // 'static + Clone 要求:JoinSet::spawn 需要 owned 闭包(不能借用 self/db)。
        // 调用方需用 Arc 包裹共享状态(db/session)传入闭包。
        // 不需要并行时用 dispatch_serial(无 'static 约束)。
        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(layer.len() as u64 * 10_000)) // 每层按 SubTask 数累加
                .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>,
    {
        self.dispatch_serial(plan, &executor).await
    }

    /// 内部串行执行(无并行)
    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
    }

    /// 合并:汇总子结果 → 合并产出 + 冲突检测
    ///
    /// - 成功的 SubTask 产出拼接为 merged_output
    /// - 失败的 SubTask 跳过(不参与合并)
    /// - 冲突检测:多个 SubTask 写同一文件路径 → 标记冲突
    pub fn merge(&self, results: &[ExecutionResult]) -> MergeResult {
        let success_results: Vec<&ExecutionResult> =
            results.iter().filter(|r| r.success).collect();

        let parts: Vec<String> = success_results
            .iter()
            .map(|r| {
                format!(
                    "## SubTask: {} (Persona: {})\n\n{}",
                    r.subtask_id, r.persona_id, r.output
                )
            })
            .collect();

        // 冲突检测:检查是否有多个 SubTask 改了同一文件
        let conflicts = self.detect_file_conflicts(results);

        MergeResult {
            merged_output: parts.join("\n\n---\n\n"),
            conflicts,
        }
    }

    /// 检测文件冲突:从 ExecutionResult 中提取写入的文件路径,同路径 → 冲突
    /// Phase 1:基于 output 文本中的文件路径关键词(简化检测)
    /// Phase 2:从 Git worktree 的 git diff 提取精确路径
    fn detect_file_conflicts(&self, results: &[ExecutionResult]) -> Vec<ConflictItem> {
        use std::collections::HashMap;

        // 收集每个 SubTask 写入的文件路径(从 output 中提取)
        let mut file_map: HashMap<String, Vec<String>> = HashMap::new();
        for r in results {
            if !r.success {
                continue;
            }
            let files = extract_written_files(&r.output);
            for f in files {
                file_map.entry(f).or_default().push(r.subtask_id.clone());
            }
        }

        // 同文件被多个 SubTask 写 → 冲突
        let mut conflicts = Vec::new();
        for (file, subtasks) in &file_map {
            if subtasks.len() > 1 {
                conflicts.push(ConflictItem {
                    file: file.clone(),
                    description: format!(
                        "文件 {} 被 {} 个子任务同时修改: {}",
                        file,
                        subtasks.len(),
                        subtasks.join(", ")
                    ),
                    subtask_a: subtasks.first().cloned(),
                    subtask_b: subtasks.get(1).cloned(),
                    recommendation: None,
                });
            }
        }
        conflicts
    }

    /// 仲裁冲突:对每个冲突给出推荐解决方案
    ///
    /// Phase 1:规则驱动(简单启发式)
    /// - 只有一方成功的冲突 → 推荐成功方
    /// - 两方都成功 → 推荐 Merged(需人工确认)
    /// - 无法判断 → Manual(留给用户)
    ///
    /// Phase 2 预留:LLM 驱动(读 diff → 推荐方案 + 理由)
    pub fn arbitrate_conflicts(
        &self,
        conflicts: &mut [ConflictItem],
        results: &[ExecutionResult],
    ) {
        let success_ids: std::collections::HashSet<&str> =
            results.iter().filter(|r| r.success).map(|r| r.subtask_id.as_str()).collect();

        for c in conflicts.iter_mut() {
            let a_ok = c.subtask_a.as_deref().map(|id| success_ids.contains(id)).unwrap_or(false);
            let b_ok = c.subtask_b.as_deref().map(|id| success_ids.contains(id)).unwrap_or(false);

            c.recommendation = Some(match (a_ok, b_ok) {
                (true, false) => ConflictResolution::AcceptA,
                (false, true) => ConflictResolution::AcceptB,
                (true, true) => ConflictResolution::Merged,
                (false, false) => ConflictResolution::Manual,
            });

            tracing::info!(
                file = %c.file,
                recommendation = ?c.recommendation,
                "[REVIEWER] 冲突仲裁推荐"
            );
        }
    }
}

/// 从执行产出文本中提取被写入的文件路径(简化检测:匹配 write_file/patch_file 后的路径)
fn extract_written_files(output: &str) -> Vec<String> {
    let mut files = Vec::new();
    for line in output.lines() {
        let trimmed = line.trim();
        // 匹配 "write_file: path" 或 "patch_file: path" 或 "写入: path" 模式
        for prefix in ["write_file:", "patch_file:", "写入:", "修改:"] {
            if let Some(rest) = trimmed.strip_prefix(prefix) {
                let path = rest.trim().split_whitespace().next().unwrap_or("");
                if !path.is_empty() {
                    files.push(path.to_string());
                }
            }
        }
    }
    files
}

// ---- 单元测试 ---------------------------------------------------------------

#[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());
    }

    #[test]
    fn merge_skips_failed_results() {
        let coord = make_coord();
        let results = vec![
            ExecutionResult {
                subtask_id: "ok".into(),
                persona_id: "coder".into(),
                output: "成功产出".into(),
                success: true,
            },
            ExecutionResult {
                subtask_id: "fail".into(),
                persona_id: "coder".into(),
                output: "执行失败".into(),
                success: false,
            },
        ];
        let merged = coord.merge(&results);
        assert!(merged.merged_output.contains("成功产出"), "应含成功产出");
        assert!(!merged.merged_output.contains("执行失败"), "不应含失败产出");
    }

    #[test]
    fn merge_detects_same_file_conflict() {
        let coord = make_coord();
        let results = vec![
            ExecutionResult {
                subtask_id: "A".into(),
                persona_id: "coder".into(),
                output: "write_file: src/main.rs\n新增 auth 模块".into(),
                success: true,
            },
            ExecutionResult {
                subtask_id: "B".into(),
                persona_id: "coder".into(),
                output: "write_file: src/main.rs\n新增 payment 模块".into(),
                success: true,
            },
        ];
        let merged = coord.merge(&results);
        assert_eq!(merged.conflicts.len(), 1, "同文件应检测到 1 个冲突");
        assert!(merged.conflicts[0].file.contains("main.rs"));
        assert!(merged.conflicts[0].description.contains("A"));
        assert!(merged.conflicts[0].description.contains("B"));
    }

    #[test]
    fn merge_no_conflict_different_files() {
        let coord = make_coord();
        let results = vec![
            ExecutionResult {
                subtask_id: "A".into(),
                persona_id: "coder".into(),
                output: "write_file: src/auth.rs".into(),
                success: true,
            },
            ExecutionResult {
                subtask_id: "B".into(),
                persona_id: "coder".into(),
                output: "write_file: src/payment.rs".into(),
                success: true,
            },
        ];
        let merged = coord.merge(&results);
        assert!(merged.conflicts.is_empty(), "不同文件不应有冲突");
    }

    #[test]
    fn merge_conflict_ignored_for_failed() {
        let coord = make_coord();
        // 失败的 SubTask 不参与冲突检测
        let results = vec![
            ExecutionResult {
                subtask_id: "A".into(),
                persona_id: "coder".into(),
                output: "write_file: src/main.rs".into(),
                success: true,
            },
            ExecutionResult {
                subtask_id: "B".into(),
                persona_id: "coder".into(),
                output: "write_file: src/main.rs".into(),
                success: false, // 失败
            },
        ];
        let merged = coord.merge(&results);
        assert!(merged.conflicts.is_empty(), "失败的不应参与冲突检测");
    }

    #[test]
    fn extract_files_from_output() {
        let files = extract_written_files("write_file: src/main.rs\n其他内容");
        assert_eq!(files, vec!["src/main.rs"]);

        let files = extract_written_files("patch_file: lib/utils.rs\n修改完成");
        assert!(files.contains(&"lib/utils.rs".to_string()));

        let files = extract_written_files("没有写操作");
        assert!(files.is_empty());
    }

    // -- Reviewer 仲裁 --

    #[test]
    fn arb_01_both_success_recommends_merged() {
        let coord = make_coord();
        let mut conflicts = vec![ConflictItem {
            file: "main.rs".into(),
            description: "冲突".into(),
            subtask_a: Some("A".into()),
            subtask_b: Some("B".into()),
            recommendation: None,
        }];
        let results = vec![
            ExecutionResult { subtask_id: "A".into(), persona_id: "coder".into(), output: "".into(), success: true },
            ExecutionResult { subtask_id: "B".into(), persona_id: "coder".into(), output: "".into(), success: true },
        ];
        coord.arbitrate_conflicts(&mut conflicts, &results);
        assert_eq!(conflicts[0].recommendation, Some(ConflictResolution::Merged));
    }

    #[test]
    fn arb_02_a_failed_recommends_b() {
        let coord = make_coord();
        let mut conflicts = vec![ConflictItem {
            file: "main.rs".into(),
            description: "冲突".into(),
            subtask_a: Some("A".into()),
            subtask_b: Some("B".into()),
            recommendation: None,
        }];
        let results = vec![
            ExecutionResult { subtask_id: "A".into(), persona_id: "coder".into(), output: "".into(), success: false },
            ExecutionResult { subtask_id: "B".into(), persona_id: "coder".into(), output: "".into(), success: true },
        ];
        coord.arbitrate_conflicts(&mut conflicts, &results);
        assert_eq!(conflicts[0].recommendation, Some(ConflictResolution::AcceptB));
    }

    #[test]
    fn arb_03_both_failed_recommends_manual() {
        let coord = make_coord();
        let mut conflicts = vec![ConflictItem {
            file: "main.rs".into(),
            description: "冲突".into(),
            subtask_a: Some("A".into()),
            subtask_b: Some("B".into()),
            recommendation: None,
        }];
        let results = vec![
            ExecutionResult { subtask_id: "A".into(), persona_id: "coder".into(), output: "".into(), success: false },
            ExecutionResult { subtask_id: "B".into(), persona_id: "coder".into(), output: "".into(), success: false },
        ];
        coord.arbitrate_conflicts(&mut conflicts, &results);
        assert_eq!(conflicts[0].recommendation, Some(ConflictResolution::Manual));
    }

    // -- 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");
    }
}