clawmates/crates/cm-api/src/swarm.rs

//! The self-verifying swarm loop: Opus 4.8 decomposes a GOAL into worker tasks →
//! a worker swarm executes each (grounded with web search) → Opus verifies each
//! output against the per-task checklist (and that cited URLs resolve) → rejected
//! tasks requeue with the reason → loop until nothing fails (or MAX_PASSES).
//!
//! It runs as a durable run (tier `swarm`) so it streams into the Runs view and
//! survives restarts. Each worker output and each verdict is journaled as a
//! `StepRecord` in the checkpoint, which `run_events_sse` emits as `step` events.

use cm_orchestrator::{RunMetrics, RunProgress, RunRecord, StepPhase, StepRecord};
use cm_runtime::Runtime;
use cm_topology::TopologyKind;
use serde::Deserialize;
use sqlx::PgPool;
use uuid::Uuid;

use crate::routes::claws::extract_json;

/// The swarm config (stored in the durable run's `graph` JSONB).
#[derive(Deserialize, Default)]
pub struct SwarmJob {
    pub goal: String,
    #[serde(default)]
    pub checklist: Vec<String>,
    #[serde(default)]
    pub task_count: Option<usize>,
    #[serde(default)]
    pub worker_model: String,
}

const MAX_PASSES: usize = 3;

const PLAN_SYSTEM: &str = "You are the planner for a self-verifying agent swarm. Decompose the GOAL into a list \
of independent, concrete worker tasks — one unit of work each (e.g. one company, one file, one question). Each \
task must be self-contained and instruct the worker to cite resolvable source URLs. Respond with STRICT JSON \
ONLY: {\"tasks\":[\"task 1\",\"task 2\", ...]}. Aim for the requested count if one is given, else pick a sensible number.";

fn checklist_lines(checklist: &[String]) -> String {
    checklist
        .iter()
        .map(|c| format!("- {c}"))
        .collect::<Vec<_>>()
        .join("\n")
}

fn worker_system(checklist: &[String]) -> String {
    format!(
        "You are a worker in a research swarm. Complete the TASK precisely and concisely. Every factual claim \
MUST include a resolvable source URL (use web search to find real sources). Your output will be verified against \
this checklist — satisfy ALL of it:\n{}",
        checklist_lines(checklist)
    )
}

fn verify_system(checklist: &[String]) -> String {
    format!(
        "You are a STRICT verifier. Check the worker OUTPUT for its TASK against this checklist:\n{}\n\nAlso \
confirm any cited source URLs are real/resolvable and use web search to spot-check the key claims. If ANYTHING \
fails, reject it. Respond with STRICT JSON ONLY: {{\"passed\": true|false, \"reason\": \"one concise sentence\"}}.",
        checklist_lines(checklist)
    )
}

/// Resolve a worker-model spec to something the runtime can call. "auto"/unknown
/// → cheap, always-available Claude Haiku (workers are cheap; Opus verifies).
/// An explicit `name:model` (e.g. `kimi:kimi-k2.6`) is honored as-is.
fn resolve_worker_model(requested: &str) -> String {
    let r = requested.trim();
    if r.is_empty() || r.eq_ignore_ascii_case("auto") {
        "claude-haiku-4-5-20251001".to_string()
    } else {
        r.to_string()
    }
}

fn step(
    node_id: impl Into<String>,
    role: impl Into<String>,
    phase: StepPhase,
    output: impl Into<String>,
) -> StepRecord {
    StepRecord {
        node_id: node_id.into(),
        role: role.into(),
        phase,
        output: output.into(),
        gated: Vec::new(),
        tokens: 0,
        spend: Default::default(),
    }
}

async fn ckpt(pool: &PgPool, id: Uuid, records: &[StepRecord], totals: &RunMetrics) {
    let prog = RunProgress {
        completed: records.len(),
        outputs: Vec::new(),
        records: records.to_vec(),
        totals: *totals,
    };
    if let Ok(v) = serde_json::to_value(&prog) {
        let _ = cm_db::repo::topology_runs::checkpoint(pool, id, &v, records.len() as i64).await;
    }
}

/// Run a swarm job to completion, journaling every worker output + verdict.
pub async fn run_swarm_job(
    pool: &PgPool,
    runtime: &Runtime,
    id: Uuid,
    job: SwarmJob,
    goal: &str,
) -> Result<RunRecord, String> {
    let mut records: Vec<StepRecord> = Vec::new();
    let totals = RunMetrics::default();
    let checklist = if job.checklist.is_empty() {
        vec![
            "output is accurate and complete".to_string(),
            "every claim cites a resolvable source URL".to_string(),
        ]
    } else {
        job.checklist.clone()
    };
    let worker_model = resolve_worker_model(&job.worker_model);

    // 1) PLAN — Opus decomposes the goal into worker tasks.
    // This record is written BEFORE the call, so it cannot name the model that
    // answers. The record after the call can, and does.
    records.push(step(
        "planner",
        "planner",
        StepPhase::Plan,
        format!("Planning tasks for: {goal}"),
    ));
    ckpt(pool, id, &records, &totals).await;
    let want = job
        .task_count
        .map(|n| format!("\n\nDesired number of tasks: {n}."))
        .unwrap_or_default();
    let plan_user = format!(
        "GOAL:\n{goal}\n\nCHECKLIST each task's output must satisfy:\n{}{want}",
        checklist_lines(&checklist)
    );
    // The recorded role says which model ANSWERED. When opus is capped the
    // chain steps down, and a step labelled "planner:opus" that GLM wrote is a
    // lie in the one place an operator looks to explain a bad decomposition.
    let (plan_raw, plan_model) = crate::subscription::complete_with_fallback(
        runtime,
        PLAN_SYSTEM,
        &plan_user,
        "claude-opus-5",
        4000,
        false,
    )
    .await?;
    let tasks: Vec<String> = extract_json(&plan_raw)
        .and_then(|v| {
            v.get("tasks").and_then(|t| t.as_array()).map(|a| {
                a.iter()
                    .filter_map(|x| x.as_str().map(String::from))
                    .collect()
            })
        })
        .unwrap_or_default();
    if tasks.is_empty() {
        return Err("planner produced no tasks".to_string());
    }
    records.push(step(
        "planner",
        format!("planner:{plan_model}"),
        StepPhase::Plan,
        format!(
            "Decomposed into {} tasks. Workers: {worker_model}. Verifier: claude-opus-4-8.",
            tasks.len()
        ),
    ));
    ckpt(pool, id, &records, &totals).await;

    // 2) LOOP — run pending tasks, verify each, requeue failures until clean.
    let mut pending: Vec<(usize, String)> = tasks.iter().cloned().enumerate().collect();
    let mut results: std::collections::BTreeMap<usize, String> = std::collections::BTreeMap::new();
    let wsys = worker_system(&checklist);
    let vsys = verify_system(&checklist);

    for pass in 1..=MAX_PASSES {
        if pending.is_empty() {
            break;
        }
        let count = pending.len();
        let mut still: Vec<(usize, String)> = Vec::new();
        let mut rejected = 0usize;
        for (idx, task) in pending.iter() {
            // `worker_model` may be a `name:model` spec the operator chose;
            // `complete_or` passes those straight through untouched.
            let out =
                crate::subscription::complete_or(runtime, &wsys, task, &worker_model, 4000, true)
                    .await
                    .unwrap_or_else(|e| format!("worker error: {e}"));
            records.push(step(
                format!("task-{idx}"),
                format!("worker:{worker_model}"),
                StepPhase::Work,
                out.clone(),
            ));
            ckpt(pool, id, &records, &totals).await;

            let vuser = format!("TASK:\n{task}\n\nWORKER OUTPUT:\n{out}");
            let v_raw = crate::subscription::complete_with_fallback(
                runtime,
                &vsys,
                &vuser,
                "claude-opus-5",
                1200,
                true,
            )
            .await
            .map(|(text, _)| text)
            .unwrap_or_default();
            let v = extract_json(&v_raw);
            let passed = v
                .as_ref()
                .and_then(|x| x.get("passed").and_then(|p| p.as_bool()))
                .unwrap_or(false);
            let reason = v
                .as_ref()
                .and_then(|x| x.get("reason").and_then(|r| r.as_str()))
                .unwrap_or("no verifier response")
                .to_string();
            records.push(step(
                format!("verify-{idx}"),
                "verifier:opus",
                StepPhase::Aggregate,
                format!("{} — {reason}", if passed { "✓ PASS" } else { "✗ REJECT" }),
            ));
            ckpt(pool, id, &records, &totals).await;

            if passed {
                results.insert(*idx, out);
            } else {
                rejected += 1;
                still.push((
                    *idx,
                    format!(
                        "{task}\n\n(Your previous attempt was REJECTED: {reason}. Correct it.)"
                    ),
                ));
            }
        }
        records.push(step(
            "verifier",
            "verifier:opus",
            StepPhase::Aggregate,
            format!("Verify pass {pass}: checked {count}, rejected {rejected}."),
        ));
        ckpt(pool, id, &records, &totals).await;
        pending = still;
    }

    // 3) Assemble the report.
    let mut report = format!(
        "# Swarm result — {goal}\n\n{} of {} tasks verified clean.\n",
        results.len(),
        tasks.len()
    );
    for (idx, task) in tasks.iter().enumerate() {
        report.push_str(&format!("\n## Task {}\n", idx + 1));
        match results.get(&idx) {
            Some(out) => report.push_str(out),
            None => report.push_str(&format!("(unresolved after {MAX_PASSES} passes)\n{task}")),
        }
        report.push('\n');
    }

    Ok(RunRecord {
        kind: TopologyKind::Swarm,
        steps: records,
        final_output: report,
        totals,
    })
}