crabmate 0.5.0

//! DAG 并行调度：就绪检测、choice 剪枝、`for_each` 运行时展开、信号量主循环。

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use futures_util::FutureExt;
use futures_util::StreamExt;
use futures_util::stream::FuturesUnordered;
use tokio::sync::Semaphore;

use super::super::for_each_expand::expand_pending_for_each;
use super::super::model::{WorkflowNodeSpec, WorkflowSpec};
use super::super::run_if::{node_deps_resolved, node_run_if_satisfied};
use super::super::types::{NodeRunResult, NodeRunStatus};
use super::node::run_node;
use super::trace::{WorkflowTracePush, workflow_trace_push};
use super::{WorkflowApprovalMode, WorkflowToolExecCtx};

/// `execute_workflow_dag` 主调度循环结束后的聚合状态。
pub(crate) struct DagExecutionProgress {
    pub(crate) completed: HashMap<String, NodeRunResult>,
    pub(crate) started: HashSet<String>,
    pub(crate) completion_order: Vec<String>,
    pub(crate) first_failure: Option<NodeRunResult>,
}

/// `inflight` 为空时的处理：正常结束 / fail_fast 退出 / 活锁检测。
/// 返回值：`true` 表示应 break 退出主循环，`false` 表示 continue。
#[allow(clippy::too_many_arguments)]
fn dag_handle_empty_inflight(
    spec: &WorkflowSpec,
    active_nodes: &[WorkflowNodeSpec],
    completed: &mut HashMap<String, NodeRunResult>,
    started: &mut HashSet<String>,
    for_each_pending: &[super::super::model::ForEachPendingSpec],
    first_failure: &Option<NodeRunResult>,
    tool_exec_ctx: &WorkflowToolExecCtx,
    stall_count: &mut u32,
    max_stall: u32,
) -> bool {
    // 正常调度完成
    if dag_schedule_finished(active_nodes, completed, for_each_pending) {
        return true; // break
    }
    // fail_fast 首失败
    if spec.fail_fast && first_failure.is_some() {
        dag_mark_remaining_nodes_fail_fast_skipped(active_nodes, completed, started, tool_exec_ctx);
        return true; // break
    }
    // 活锁检测（P0-3）
    *stall_count += 1;
    if *stall_count > max_stall {
        let err_msg = format!(
            "workflow 调度活锁：连续 {} 次迭代无节点可调度（可能是运行时循环依赖），已强制终止",
            stall_count,
        );
        log::error!(target: "crabmate", "{}", err_msg);
        for node in active_nodes.iter() {
            if !completed.contains_key(&node.id) {
                started.insert(node.id.clone());
                completed.insert(
                    node.id.clone(),
                    NodeRunResult {
                        id: node.id.clone(),
                        status: NodeRunStatus::Failed,
                        output: err_msg.clone().into(),
                        workspace_changed: false,
                        exit_code: None,
                        error_code: Some("workflow_livelock".to_string()),
                        attempt: 0,
                    },
                );
            }
        }
        return true; // break
    }
    false // continue
}

fn node_panic_result(node_id: String, panic_payload: Box<dyn std::any::Any + Send>) -> NodeRunResult {
    let msg = if let Some(s) = panic_payload.downcast_ref::<&str>() {
        format!("workflow 节点 panic：{}", s)
    } else if let Some(s) = panic_payload.downcast_ref::<String>() {
        format!("workflow 节点 panic：{}", s)
    } else {
        "workflow 节点 panic（原因未知）".to_string()
    };
    log::error!(
        target: "crabmate",
        "workflow 节点 panic node_id={} msg={}",
        node_id,
        msg,
    );
    NodeRunResult {
        id: node_id,
        status: NodeRunStatus::Failed,
        output: msg.into(),
        workspace_changed: false,
        exit_code: None,
        error_code: Some("workflow_node_panic".to_string()),
        attempt: 1,
    }
}

async fn run_node_with_permit(
    node: WorkflowNodeSpec,
    approval_mode: WorkflowApprovalMode,
    exec_ctx: WorkflowToolExecCtx,
    completed_snapshot: HashMap<String, NodeRunResult>,
    inject_max_chars: usize,
    permit_sem: Arc<Semaphore>,
) -> NodeRunResult {
    let node_id = node.id.clone();
    let _permit = match permit_sem.acquire_owned().await {
        Ok(p) => p,
        Err(_) => {
            return NodeRunResult {
                id: node_id,
                status: NodeRunStatus::Failed,
                output: "workflow 并发控制异常（semaphore closed）".into(),
                workspace_changed: false,
                exit_code: None,
                error_code: Some("workflow_semaphore_closed".to_string()),
                attempt: 1,
            };
        }
    };
    let node_fut = run_node(
        node,
        approval_mode,
        exec_ctx,
        completed_snapshot,
        inject_max_chars,
        "main",
    );
    match std::panic::AssertUnwindSafe(node_fut).catch_unwind().await {
        Ok(res) => res,
        Err(panic_payload) => node_panic_result(node_id, panic_payload),
    }
}

fn mark_run_if_skipped(
    node: &WorkflowNodeSpec,
    started: &mut HashSet<String>,
    completed: &mut HashMap<String, NodeRunResult>,
    tool_exec_ctx: &WorkflowToolExecCtx,
) {
    started.insert(node.id.clone());
    completed.insert(
        node.id.clone(),
        NodeRunResult {
            id: node.id.clone(),
            status: NodeRunStatus::Skipped,
            output: "choice: run_if not satisfied".into(),
            workspace_changed: false,
            exit_code: None,
            error_code: Some("workflow_choice_skipped".to_string()),
            attempt: 0,
        },
    );
    workflow_trace_push(WorkflowTracePush {
        trace: &tool_exec_ctx.trace_events,
        workflow_run_id: tool_exec_ctx.workflow_run_id,
        event: "node_choice_skipped",
        node_id: Some(node.id.as_str()),
        detail: node.run_if.as_ref().map(|_| "run_if=false".to_string()),
        attempt: None,
        status: Some("skipped"),
        elapsed_ms: None,
        error_code: Some("workflow_choice_skipped"),
        tool_name: Some(node.tool_name.as_str()),
        phase: Some("main"),
    });
}

fn trace_for_each_expanded(tool_exec_ctx: &WorkflowToolExecCtx, expanded: &[String]) {
    for id in expanded.iter() {
        workflow_trace_push(WorkflowTracePush {
            trace: &tool_exec_ctx.trace_events,
            workflow_run_id: tool_exec_ctx.workflow_run_id,
            event: "for_each_expanded",
            node_id: Some(id.as_str()),
            detail: None,
            attempt: None,
            status: None,
            elapsed_ms: None,
            error_code: None,
            tool_name: None,
            phase: Some("main"),
        });
    }
}

enum NodeScheduleDecision {
    Ignore,
    SkipRunIf,
    Start,
}

fn node_schedule_decision(
    node: &WorkflowNodeSpec,
    started: &HashSet<String>,
    completed: &HashMap<String, NodeRunResult>,
) -> NodeScheduleDecision {
    if started.contains(&node.id) || completed.contains_key(&node.id) {
        return NodeScheduleDecision::Ignore;
    }
    if !node_deps_resolved(&node.deps, completed) {
        return NodeScheduleDecision::Ignore;
    }
    if !node_run_if_satisfied(node.run_if.as_ref(), completed) {
        return NodeScheduleDecision::SkipRunIf;
    }
    NodeScheduleDecision::Start
}

/// 并行调度就绪节点并等待全部 inflight 完成。
pub(super) async fn dag_run_parallel_schedule_loop(
    spec: &WorkflowSpec,
    approval_mode: WorkflowApprovalMode,
    tool_exec_ctx: WorkflowToolExecCtx,
) -> DagExecutionProgress {
    let mut active_nodes = spec.nodes.clone();
    let mut for_each_pending = spec.for_each_pending.clone();
    let mut completed: HashMap<String, NodeRunResult> = HashMap::new();
    let mut started: HashSet<String> = HashSet::new();
    let mut completion_order: Vec<String> = Vec::new();
    let mut first_failure: Option<NodeRunResult> = None;

    let max_parallelism = spec.max_parallelism.max(1);
    let semaphore = Arc::new(Semaphore::new(max_parallelism));
    let mut inflight: FuturesUnordered<_> = FuturesUnordered::new();

    // P0-3: 运行时活锁检测
    let max_stall = (max_parallelism.max(4) * 2) as u32;
    let mut stall_count: u32 = 0;

    loop {
        let expanded =
            expand_pending_for_each(&mut for_each_pending, &mut active_nodes, &completed);
        trace_for_each_expanded(&tool_exec_ctx, &expanded);

        if !(spec.fail_fast && first_failure.is_some()) {
            for node in active_nodes.iter() {
                match node_schedule_decision(node, &started, &completed) {
                    NodeScheduleDecision::Ignore => {}
                    NodeScheduleDecision::SkipRunIf => {
                        mark_run_if_skipped(node, &mut started, &mut completed, &tool_exec_ctx);
                    }
                    NodeScheduleDecision::Start => {
                        started.insert(node.id.clone());
                        inflight.push(run_node_with_permit(
                            node.clone(),
                            approval_mode.clone(),
                            tool_exec_ctx.clone(),
                            completed.clone(),
                            spec.output_inject_max_chars,
                            semaphore.clone(),
                        ));
                    }
                }
            }
        }

        if inflight.is_empty() {
            if dag_handle_empty_inflight(
                spec,
                &active_nodes,
                &mut completed,
                &mut started,
                &for_each_pending,
                &first_failure,
                &tool_exec_ctx,
                &mut stall_count,
                max_stall,
            ) {
                break;
            }
            continue;
        }

        stall_count = 0;
        let Some(res) = inflight.next().await else {
            continue;
        };
        dag_record_node_completion(
            &res,
            &mut completed,
            &mut completion_order,
            &mut first_failure,
        );
    }

    DagExecutionProgress {
        completed,
        started,
        completion_order,
        first_failure,
    }
}

fn dag_record_node_completion(
    res: &NodeRunResult,
    completed: &mut HashMap<String, NodeRunResult>,
    completion_order: &mut Vec<String>,
    first_failure: &mut Option<NodeRunResult>,
) {
    if res.status == NodeRunStatus::Passed {
        completion_order.push(res.id.clone());
        completed.insert(res.id.clone(), res.clone());
        return;
    }
    if res.status == NodeRunStatus::Skipped {
        completed.insert(res.id.clone(), res.clone());
        return;
    }
    if first_failure.is_none() {
        *first_failure = Some(res.clone());
    }
    completed.insert(
        res.id.clone(),
        NodeRunResult {
            id: res.id.clone(),
            status: NodeRunStatus::Failed,
            output: res.output.clone(),
            workspace_changed: res.workspace_changed,
            exit_code: res.exit_code,
            error_code: res.error_code.clone(),
            attempt: res.attempt,
        },
    );
}

fn dag_schedule_finished(
    nodes: &[WorkflowNodeSpec],
    completed: &HashMap<String, NodeRunResult>,
    for_each_pending: &[super::super::model::ForEachPendingSpec],
) -> bool {
    nodes.iter().all(|n| completed.contains_key(&n.id)) && for_each_pending.is_empty()
}

/// `fail_fast` 且已有首失败后：将尚未完成的节点标为跳过，避免调度器在 `inflight` 为空时 tight-loop。
fn dag_mark_remaining_nodes_fail_fast_skipped(
    nodes: &[WorkflowNodeSpec],
    completed: &mut HashMap<String, NodeRunResult>,
    started: &mut HashSet<String>,
    tool_exec_ctx: &WorkflowToolExecCtx,
) {
    for node in nodes {
        if completed.contains_key(&node.id) {
            continue;
        }
        started.insert(node.id.clone());
        completed.insert(
            node.id.clone(),
            NodeRunResult {
                id: node.id.clone(),
                status: NodeRunStatus::Skipped,
                output: "fail_fast: workflow aborted after first failure".into(),
                workspace_changed: false,
                exit_code: None,
                error_code: Some("workflow_fail_fast_aborted".to_string()),
                attempt: 0,
            },
        );
        workflow_trace_push(WorkflowTracePush {
            trace: &tool_exec_ctx.trace_events,
            workflow_run_id: tool_exec_ctx.workflow_run_id,
            event: "node_fail_fast_skipped",
            node_id: Some(node.id.as_str()),
            detail: Some("fail_fast after upstream failure".to_string()),
            attempt: None,
            status: Some("skipped"),
            elapsed_ms: None,
            error_code: Some("workflow_fail_fast_aborted"),
            tool_name: Some(node.tool_name.as_str()),
            phase: Some("main"),
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cm_workflow::model::WorkflowNodeSpec;
    use std::collections::HashMap;

    fn sample_nodes() -> Vec<WorkflowNodeSpec> {
        vec![
            WorkflowNodeSpec {
                id: "a".into(),
                tool_name: "tool_a".into(),
                tool_args: serde_json::json!({}),
                deps: vec![],
                requires_approval: false,
                timeout_secs: None,
                compensate_with: vec![],
                max_retries: 0,
                node_tool_role: None,
                run_if: None,
            },
            WorkflowNodeSpec {
                id: "b".into(),
                tool_name: "tool_b".into(),
                tool_args: serde_json::json!({}),
                deps: vec!["a".into()],
                requires_approval: false,
                timeout_secs: None,
                compensate_with: vec![],
                max_retries: 0,
                node_tool_role: None,
                run_if: None,
            },
        ]
    }

    #[test]
    fn dag_schedule_finished_requires_all_nodes_completed() {
        let nodes = sample_nodes();
        let mut completed = HashMap::new();
        assert!(!dag_schedule_finished(&nodes, &completed, &[]));
        completed.insert(
            "a".into(),
            NodeRunResult {
                id: "a".into(),
                status: NodeRunStatus::Failed,
                output: "fail".into(),
                workspace_changed: false,
                exit_code: None,
                error_code: None,
                attempt: 1,
            },
        );
        assert!(!dag_schedule_finished(&nodes, &completed, &[]));
    }
}