clawmates/crates/cm-orchestrator/src/plan.rs

//! Pure, synchronous planners: turn a [`TopologyGraph`] into an ordered list
//! of steps. Separated from execution so the control flow is unit-testable
//! without any async/LLM machinery.

use std::collections::{HashMap, HashSet, VecDeque};

use cm_topology::TopologyGraph;

use crate::{OrchestratorError, StepPhase};

/// One planned step: which node acts, in what phase, and which prior step
/// outputs (plus optionally the top-level task) form its context.
pub(crate) struct PlanStep {
    pub node_idx: usize,
    pub phase: StepPhase,
    pub use_task: bool,
    pub ctx_from: Vec<usize>,
}

fn index_map(g: &TopologyGraph) -> HashMap<&str, usize> {
    g.nodes
        .iter()
        .enumerate()
        .map(|(i, n)| (n.id.as_str(), i))
        .collect()
}

fn in_degrees(g: &TopologyGraph, idx: &HashMap<&str, usize>) -> Vec<usize> {
    let mut indeg = vec![0usize; g.nodes.len()];
    let mut seen = HashSet::new();
    for e in &g.edges {
        if let (Some(&a), Some(&b)) = (idx.get(e.from.as_str()), idx.get(e.to.as_str())) {
            if a != b && seen.insert((a, b)) {
                indeg[b] += 1;
            }
        }
    }
    indeg
}

/// Hierarchical: root plans (top-down), direct children work, root synthesizes
/// (bottom-up). Deeper trees are flattened to one level for v1.
pub(crate) fn hierarchical(g: &TopologyGraph) -> Result<Vec<PlanStep>, OrchestratorError> {
    let idx = index_map(g);
    let n = g.nodes.len();
    let indeg = in_degrees(g, &idx);
    let root = (0..n).find(|&i| indeg[i] == 0).unwrap_or(0);

    let mut children = Vec::new();
    let mut seen = HashSet::new();
    for e in &g.edges {
        if e.from == g.nodes[root].id {
            if let Some(&c) = idx.get(e.to.as_str()) {
                if c != root && seen.insert(c) {
                    children.push(c);
                }
            }
        }
    }

    let mut steps = Vec::new();
    if children.is_empty() {
        steps.push(PlanStep {
            node_idx: root,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![],
        });
        return Ok(steps);
    }
    steps.push(PlanStep {
        node_idx: root,
        phase: StepPhase::Plan,
        use_task: true,
        ctx_from: vec![],
    });
    let mut child_steps = Vec::new();
    for c in children {
        steps.push(PlanStep {
            node_idx: c,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![0],
        });
        child_steps.push(steps.len() - 1);
    }
    steps.push(PlanStep {
        node_idx: root,
        phase: StepPhase::Synth,
        use_task: false,
        ctx_from: child_steps,
    });
    Ok(steps)
}

/// Pipeline: topological order; each stage takes the previous stage's output.
pub(crate) fn pipeline(g: &TopologyGraph) -> Result<Vec<PlanStep>, OrchestratorError> {
    let idx = index_map(g);
    let n = g.nodes.len();
    let mut succ = vec![Vec::new(); n];
    let mut indeg = vec![0usize; n];
    let mut seen = HashSet::new();
    for e in &g.edges {
        if let (Some(&a), Some(&b)) = (idx.get(e.from.as_str()), idx.get(e.to.as_str())) {
            if a != b && seen.insert((a, b)) {
                succ[a].push(b);
                indeg[b] += 1;
            }
        }
    }
    let mut queue: VecDeque<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
    let mut order = Vec::new();
    while let Some(u) = queue.pop_front() {
        order.push(u);
        for &v in &succ[u] {
            indeg[v] -= 1;
            if indeg[v] == 0 {
                queue.push_back(v);
            }
        }
    }
    // Any nodes left in a cycle: append in node order so we still run them.
    for i in 0..n {
        if !order.contains(&i) {
            order.push(i);
        }
    }

    let steps = order
        .iter()
        .enumerate()
        .map(|(i, &node)| PlanStep {
            node_idx: node,
            phase: StepPhase::Work,
            use_task: i == 0,
            ctx_from: if i == 0 { vec![] } else { vec![i - 1] },
        })
        .collect();
    Ok(steps)
}

/// Swarm: every node attempts the task independently, then an aggregator
/// (a "coordinator" role, else a source node) synthesizes all outputs.
pub(crate) fn swarm(g: &TopologyGraph) -> Result<Vec<PlanStep>, OrchestratorError> {
    let n = g.nodes.len();
    if n == 0 {
        return Err(OrchestratorError::Malformed("empty topology".into()));
    }
    let idx = index_map(g);
    let indeg = in_degrees(g, &idx);

    let mut steps: Vec<PlanStep> = (0..n)
        .map(|i| PlanStep {
            node_idx: i,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![],
        })
        .collect();

    let aggregator = (0..n)
        .find(|&i| g.nodes[i].role.to_lowercase().contains("coordinator"))
        .or_else(|| (0..n).find(|&i| indeg[i] == 0))
        .unwrap_or(0);

    steps.push(PlanStep {
        node_idx: aggregator,
        phase: StepPhase::Aggregate,
        use_task: false,
        ctx_from: (0..n).collect(),
    });
    Ok(steps)
}

/// Mesh / blackboard: two rounds of peer exchange (round 2 sees every round-1
/// output), then the first peer aggregates — models iterative convergence.
pub(crate) fn mesh(g: &TopologyGraph) -> Result<Vec<PlanStep>, OrchestratorError> {
    let n = g.nodes.len();
    if n == 0 {
        return Err(OrchestratorError::Malformed("empty topology".into()));
    }
    let mut steps = Vec::new();
    for i in 0..n {
        steps.push(PlanStep {
            node_idx: i,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![],
        });
    }
    let round1: Vec<usize> = (0..n).collect();
    for i in 0..n {
        steps.push(PlanStep {
            node_idx: i,
            phase: StepPhase::Work,
            use_task: false,
            ctx_from: round1.clone(),
        });
    }
    let round2: Vec<usize> = (n..2 * n).collect();
    steps.push(PlanStep {
        node_idx: 0,
        phase: StepPhase::Aggregate,
        use_task: false,
        ctx_from: round2,
    });
    Ok(steps)
}

/// Debate: proposer drafts, critic critiques, proposer revises, judge decides.
pub(crate) fn debate(g: &TopologyGraph) -> Result<Vec<PlanStep>, OrchestratorError> {
    let n = g.nodes.len();
    if n == 0 {
        return Err(OrchestratorError::Malformed("empty topology".into()));
    }
    if n == 1 {
        return Ok(vec![PlanStep {
            node_idx: 0,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![],
        }]);
    }
    let proposer = 0;
    let critic = 1;
    let judge = if n >= 3 { 2 } else { 0 };
    Ok(vec![
        PlanStep {
            node_idx: proposer,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![],
        },
        PlanStep {
            node_idx: critic,
            phase: StepPhase::Work,
            use_task: true,
            ctx_from: vec![0],
        },
        PlanStep {
            node_idx: proposer,
            phase: StepPhase::Synth,
            use_task: false,
            ctx_from: vec![0, 1],
        },
        PlanStep {
            node_idx: judge,
            phase: StepPhase::Aggregate,
            use_task: false,
            ctx_from: vec![2, 1],
        },
    ])
}