car_multi/patterns/

swarm.rs

//! Swarm — N agents working on the same problem.
//!
//! Modes:
//! - **Parallel**: all agents run concurrently, then a synthesizer combines results.
//! - **Sequential**: agents run one after another, each seeing prior agents' outputs.
//! - **Debate**: two rounds — initial answers, then critique, then a judge picks the best.

use crate::error::MultiError;
use crate::mailbox::Mailbox;
use crate::runner::AgentRunner;
use crate::shared::SharedInfra;
use crate::types::{AgentOutput, AgentSpec};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Instant;
use tracing::instrument;

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SwarmMode {
    Parallel,
    Sequential,
    Debate,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SwarmResult {
    pub task: String,
    pub outputs: Vec<AgentOutput>,
    pub final_summary: String,
}

pub struct Swarm {
    pub agents: Vec<AgentSpec>,
    pub mode: SwarmMode,
    pub synthesizer: Option<AgentSpec>,
    /// When true, each agent gets an isolated state overlay.
    /// Writes go to a per-agent local store; reads fall through to the shared parent.
    /// On success, local state is merged back to the parent.
    pub isolated: bool,
    /// When set (parallel mode), each agent gets an isolated filesystem
    /// workspace, advertised to the runner via `AgentSpec.metadata["workspace"]`.
    pub workspaces: Option<crate::workspace::WorkspaceConfig>,
}

impl Swarm {
    pub fn new(agents: Vec<AgentSpec>, mode: SwarmMode) -> Self {
        Self {
            agents,
            mode,
            synthesizer: None,
            isolated: false,
            workspaces: None,
        }
    }

    pub fn with_synthesizer(mut self, spec: AgentSpec) -> Self {
        self.synthesizer = Some(spec);
        self
    }

    /// Enable per-agent state isolation for this swarm.
    pub fn with_isolation(mut self) -> Self {
        self.isolated = true;
        self
    }

    /// Provision an isolated filesystem workspace per agent (parallel mode). Each
    /// agent's [`AgentSpec`] gets a `workspace` metadata entry with its directory;
    /// the runner is expected to run its file tools there. Workspaces are removed
    /// when the run completes. Prevents parallel file-mutating agents from
    /// clobbering one another.
    pub fn with_workspaces(mut self, config: crate::workspace::WorkspaceConfig) -> Self {
        self.workspaces = Some(config);
        self
    }

    #[instrument(name = "multi.swarm", skip_all)]
    pub fn run<'a>(
        &'a self,
        task: &'a str,
        runner: &'a Arc<dyn AgentRunner>,
        infra: &'a SharedInfra,
    ) -> futures::future::BoxFuture<'a, Result<SwarmResult, MultiError>> {
        Box::pin(async move {
            match self.mode {
                SwarmMode::Parallel => self.run_parallel(task, runner, infra).await,
                SwarmMode::Sequential => self.run_sequential(task, runner, infra).await,
                SwarmMode::Debate => self.run_debate(task, runner, infra).await,
            }
        })
    }

    async fn run_parallel(
        &self,
        task: &str,
        runner: &Arc<dyn AgentRunner>,
        infra: &SharedInfra,
    ) -> Result<SwarmResult, MultiError> {
        let mailbox = Arc::new(Mailbox::default());

        // Each agent slot is either spawned (index into `handles`) or pre-empted
        // by the coordination budget. Keeping per-agent slots preserves output
        // order even when some agents are skipped.
        enum Slot {
            Spawned(usize),
            Skipped(AgentOutput),
        }

        // When isolated, each handle returns (Result, Option<AgentContext>) so we
        // can merge state back on success.  When not isolated, the context is None.
        let mut handles: Vec<
            tokio::task::JoinHandle<(
                Result<AgentOutput, MultiError>,
                Option<crate::task_context::AgentContext>,
            )>,
        > = Vec::new();
        let mut slots: Vec<Slot> = Vec::new();

        for spec in &self.agents {
            // Provision an isolated filesystem workspace if configured, and
            // advertise its path to the runner via the spec's metadata. Done
            // BEFORE the budget reservation so a provisioning failure doesn't
            // burn a (non-refundable) agent slot. On failure, fail this agent
            // closed rather than running it unisolated and risk clobbering a
            // sibling.
            let workspace = match &self.workspaces {
                Some(cfg) => match crate::workspace::AgentWorkspace::provision(cfg, &spec.name) {
                    Ok(ws) => Some(ws),
                    Err(e) => {
                        slots.push(Slot::Skipped(AgentOutput {
                            name: spec.name.clone(),
                            answer: String::new(),
                            turns: 0,
                            tool_calls: 0,
                            duration_ms: 0.0,
                            error: Some(format!("workspace provisioning failed: {e}")),
                            outcome: None,
                            tokens: None,
                        }));
                        continue;
                    }
                },
                None => None,
            };

            // Budget pre-flight: a crossed token/cost ceiling or the agent cap
            // stops further spawns. The whole parallel batch is launched at once,
            // so this gates the batch rather than metering mid-batch. On denial
            // the just-provisioned `workspace` guard drops here and cleans up.
            if let Err(e) = infra.begin_agent() {
                slots.push(Slot::Skipped(crate::budget::budget_skipped_output(
                    &spec.name, &e,
                )));
                continue;
            }

            let runner = Arc::clone(runner);
            let mut spec = spec.clone();
            if let Some(ws) = &workspace {
                spec = ws.inject(spec);
            }
            let task = task.to_string();
            let mailbox = Arc::clone(&mailbox);

            if self.isolated {
                let (rt, ctx) = infra.make_isolated_runtime(&spec.name);
                for tool in &spec.tools {
                    rt.register_tool(tool).await;
                }
                let ctx_clone = ctx.clone();
                handles.push(tokio::spawn(async move {
                    // Hold the workspace guard for the agent's lifetime; dropped
                    // (cleaned up) when the task finishes.
                    let _workspace = workspace;
                    let result = crate::task_context::TaskScope::run(ctx_clone, async {
                        runner.run(&spec, &task, &rt, &mailbox).await
                    })
                    .await;
                    (result, Some(ctx))
                }));
            } else {
                let rt = infra.make_runtime();
                for tool in &spec.tools {
                    rt.register_tool(tool).await;
                }
                handles.push(tokio::spawn(async move {
                    let _workspace = workspace;
                    let result = runner.run(&spec, &task, &rt, &mailbox).await;
                    (result, None)
                }));
            }
            slots.push(Slot::Spawned(handles.len() - 1));
        }

        // Move owned join results out by handle index as each slot is visited.
        let mut results: Vec<Option<_>> = futures::future::join_all(handles)
            .await
            .into_iter()
            .map(Some)
            .collect();
        let mut outputs = Vec::new();
        for (i, slot) in slots.into_iter().enumerate() {
            let handle_idx = match slot {
                Slot::Skipped(output) => {
                    outputs.push(output);
                    continue;
                }
                Slot::Spawned(idx) => idx,
            };
            match results.get_mut(handle_idx).and_then(Option::take) {
                Some(Ok((Ok(output), ctx))) => {
                    // Merge isolated state back to parent on success
                    if let Some(ctx) = ctx {
                        ctx.merge_to_parent();
                    }
                    // Record reported spend against the coordination budget.
                    infra.record_output(&output);
                    // Write to shared state
                    infra.state.set(
                        &format!("agent.{}.answer", output.name),
                        serde_json::Value::String(output.answer.clone()),
                        &format!("swarm.{}", output.name),
                    );
                    outputs.push(output);
                }
                Some(Ok((Err(e), _ctx))) => {
                    // Note: an agent that spent tokens before returning Err has
                    // that spend dropped — the error path carries no token
                    // payload, so the budget can under-count failed work.
                    outputs.push(AgentOutput {
                        name: self.agents[i].name.clone(),
                        answer: String::new(),
                        turns: 0,
                        tool_calls: 0,
                        duration_ms: 0.0,
                        error: Some(e.to_string()),
                        outcome: None,
                        tokens: None,
                    });
                }
                Some(Err(e)) => {
                    outputs.push(AgentOutput {
                        name: self.agents[i].name.clone(),
                        answer: String::new(),
                        turns: 0,
                        tool_calls: 0,
                        duration_ms: 0.0,
                        error: Some(format!("join error: {}", e)),
                        outcome: None,
                        tokens: None,
                    });
                }
                None => {
                    outputs.push(AgentOutput {
                        name: self.agents[i].name.clone(),
                        answer: String::new(),
                        turns: 0,
                        tool_calls: 0,
                        duration_ms: 0.0,
                        error: Some("internal: missing join result".to_string()),
                        outcome: None,
                        tokens: None,
                    });
                }
            }
        }

        let summary = self.synthesize(task, &outputs, runner, infra).await;

        Ok(SwarmResult {
            task: task.to_string(),
            outputs,
            final_summary: summary,
        })
    }

    async fn run_sequential(
        &self,
        task: &str,
        runner: &Arc<dyn AgentRunner>,
        infra: &SharedInfra,
    ) -> Result<SwarmResult, MultiError> {
        let mailbox = Arc::new(Mailbox::default());
        let mut outputs = Vec::new();

        for spec in &self.agents {
            // Budget gate before each agent. In a sequential chain this is real
            // between-agent enforcement: once a prior agent's reported spend
            // crosses a limit, the remaining agents are skipped.
            if let Err(e) = infra.begin_agent() {
                outputs.push(crate::budget::budget_skipped_output(&spec.name, &e));
                continue;
            }

            // Enrich task with prior results
            let enriched = if outputs.is_empty() {
                task.to_string()
            } else {
                let prior: Vec<String> = outputs
                    .iter()
                    .filter_map(|o: &AgentOutput| {
                        if o.succeeded() {
                            Some(format!("- {}: {}", o.name, truncate(&o.answer, 300)))
                        } else {
                            None
                        }
                    })
                    .collect();
                format!("{}\n\nPrior agents' findings:\n{}", task, prior.join("\n"))
            };

            let rt = infra.make_runtime();
            for tool in &spec.tools {
                rt.register_tool(tool).await;
            }

            let start = Instant::now();
            match runner.run(spec, &enriched, &rt, &mailbox).await {
                Ok(output) => {
                    infra.record_output(&output);
                    infra.state.set(
                        &format!("agent.{}.answer", output.name),
                        serde_json::Value::String(output.answer.clone()),
                        &format!("swarm.{}", output.name),
                    );
                    outputs.push(output);
                }
                Err(e) => {
                    outputs.push(AgentOutput {
                        name: spec.name.clone(),
                        answer: String::new(),
                        turns: 0,
                        tool_calls: 0,
                        duration_ms: start.elapsed().as_secs_f64() * 1000.0,
                        error: Some(e.to_string()),
                        outcome: None,
                        tokens: None,
                    });
                }
            }
        }

        let summary = self.synthesize(task, &outputs, runner, infra).await;

        Ok(SwarmResult {
            task: task.to_string(),
            outputs,
            final_summary: summary,
        })
    }

    async fn run_debate(
        &self,
        task: &str,
        runner: &Arc<dyn AgentRunner>,
        infra: &SharedInfra,
    ) -> Result<SwarmResult, MultiError> {
        // Round 1: independent answers
        let round1 = Swarm::new(self.agents.clone(), SwarmMode::Parallel)
            .run(task, runner, infra)
            .await?;

        // Round 2: each agent critiques the others
        let mut critique_specs = Vec::new();
        for spec in &self.agents {
            let others: Vec<String> = round1
                .outputs
                .iter()
                .filter(|o| o.name != spec.name && o.succeeded())
                .map(|o| format!("- {}: {}", o.name, truncate(&o.answer, 300)))
                .collect();

            let critique_prompt = format!(
                "{}\n\nOriginal task: {}\n\nOther agents' answers:\n{}\n\n\
                 Critique these answers and provide your improved response.",
                spec.system_prompt,
                task,
                others.join("\n")
            );

            let mut critique_spec = spec.clone();
            critique_spec.name = format!("{}_critique", spec.name);
            critique_spec.system_prompt = critique_prompt;
            critique_specs.push(critique_spec);
        }

        let round2 = Swarm::new(critique_specs, SwarmMode::Parallel)
            .run(task, runner, infra)
            .await?;

        // Combine both rounds
        let mut all_outputs = round1.outputs;
        all_outputs.extend(round2.outputs);

        let summary = self.synthesize(task, &all_outputs, runner, infra).await;

        Ok(SwarmResult {
            task: task.to_string(),
            outputs: all_outputs,
            final_summary: summary,
        })
    }

    async fn synthesize(
        &self,
        task: &str,
        outputs: &[AgentOutput],
        runner: &Arc<dyn AgentRunner>,
        infra: &SharedInfra,
    ) -> String {
        let answers: Vec<&AgentOutput> = outputs.iter().filter(|o| o.succeeded()).collect();
        if answers.is_empty() {
            return "[no agent produced an answer]".to_string();
        }
        if answers.len() == 1 {
            return answers[0].answer.clone();
        }

        if let Some(synth_spec) = &self.synthesizer {
            let summaries: Vec<String> = answers
                .iter()
                .map(|o| format!("- {}: {}", o.name, truncate(&o.answer, 500)))
                .collect();

            let synth_task = format!(
                "Original task: {}\n\nAgent outputs:\n{}\n\nSynthesize these into a single coherent answer.",
                task,
                summaries.join("\n")
            );

            // Gate the synthesizer on the budget too; on denial fall through to
            // the default concatenation rather than failing the whole run.
            if infra.begin_agent().is_ok() {
                let mailbox = Mailbox::default();
                let rt = infra.make_runtime();
                if let Ok(output) = runner.run(synth_spec, &synth_task, &rt, &mailbox).await {
                    infra.record_output(&output);
                    return output.answer;
                }
            }
        }

        // Default: concatenate with headers
        answers
            .iter()
            .map(|o| format!("## {}\n{}", o.name, o.answer))
            .collect::<Vec<_>>()
            .join("\n\n")
    }
}

fn truncate(s: &str, max_len: usize) -> &str {
    if s.len() <= max_len {
        return s;
    }
    let mut end = max_len;
    while end > 0 && !s.is_char_boundary(end) {
        end -= 1;
    }
    &s[..end]
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::MultiError;
    use crate::mailbox::Mailbox;
    use crate::runner::AgentRunner;
    use crate::types::{AgentOutput, AgentSpec};
    use car_engine::Runtime;
    use std::sync::atomic::{AtomicU32, Ordering};

    struct MockRunner {
        call_count: AtomicU32,
    }

    #[async_trait::async_trait]
    impl AgentRunner for MockRunner {
        async fn run(
            &self,
            spec: &AgentSpec,
            task: &str,
            _runtime: &Runtime,
            _mailbox: &Mailbox,
        ) -> Result<AgentOutput, MultiError> {
            let _n = self.call_count.fetch_add(1, Ordering::SeqCst);
            Ok(AgentOutput {
                name: spec.name.clone(),
                answer: format!(
                    "answer from {} for: {}",
                    spec.name,
                    &task[..task.len().min(50)]
                ),
                turns: 1,
                tool_calls: 0,
                duration_ms: 10.0,
                error: None,
                outcome: None,
                tokens: None,
            })
        }
    }

    #[tokio::test]
    async fn test_parallel_swarm() {
        let agents = vec![
            AgentSpec::new("alice", "You are Alice"),
            AgentSpec::new("bob", "You are Bob"),
        ];
        let runner: Arc<dyn AgentRunner> = Arc::new(MockRunner {
            call_count: AtomicU32::new(0),
        });
        let infra = SharedInfra::new();

        let result = Swarm::new(agents, SwarmMode::Parallel)
            .run("test task", &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.outputs.len(), 2);
        assert!(result.outputs.iter().all(|o| o.succeeded()));

        // Check shared state was written
        assert!(infra.state.get("agent.alice.answer").is_some());
        assert!(infra.state.get("agent.bob.answer").is_some());
    }

    #[tokio::test]
    async fn test_sequential_swarm() {
        let agents = vec![
            AgentSpec::new("first", "Go first"),
            AgentSpec::new("second", "Go second"),
        ];
        let runner: Arc<dyn AgentRunner> = Arc::new(MockRunner {
            call_count: AtomicU32::new(0),
        });
        let infra = SharedInfra::new();

        let result = Swarm::new(agents, SwarmMode::Sequential)
            .run("sequential task", &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.outputs.len(), 2);
        // Second agent should see first agent's output in enriched task
        assert!(result.outputs[1].answer.contains("Prior agents"));
    }

    /// Reports a fixed token spend per call so a budget can meter it.
    struct TokenRunner {
        per_call_total: u64,
    }

    #[async_trait::async_trait]
    impl AgentRunner for TokenRunner {
        async fn run(
            &self,
            spec: &AgentSpec,
            _task: &str,
            _runtime: &Runtime,
            _mailbox: &Mailbox,
        ) -> Result<AgentOutput, MultiError> {
            Ok(AgentOutput {
                name: spec.name.clone(),
                answer: format!("answer from {}", spec.name),
                turns: 1,
                tool_calls: 0,
                duration_ms: 1.0,
                error: None,
                outcome: None,
                tokens: Some(crate::types::TokenAccounting::new(
                    self.per_call_total,
                    0,
                    0.0,
                )),
            })
        }
    }

    #[tokio::test]
    async fn sequential_budget_stops_chain_when_tokens_exhausted() {
        // Three agents, each reporting 100 tokens; a 150-token ceiling lets the
        // first two run (cumulative 200 crosses 150 only after the second) and
        // denies the third.
        let agents = vec![
            AgentSpec::new("a", ""),
            AgentSpec::new("b", ""),
            AgentSpec::new("c", ""),
        ];
        let runner: Arc<dyn AgentRunner> = Arc::new(TokenRunner { per_call_total: 100 });
        let infra = SharedInfra::new().with_budget(crate::BudgetLimits {
            max_total_tokens: Some(150),
            ..Default::default()
        });

        let result = Swarm::new(agents, SwarmMode::Sequential)
            .run("task", &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.outputs.len(), 3);
        assert!(result.outputs[0].succeeded());
        assert!(result.outputs[1].succeeded());
        assert!(!result.outputs[2].succeeded());
        assert!(crate::is_budget_skipped(&result.outputs[2]));
        assert_eq!(infra.budget.snapshot().total_tokens, 200);
    }

    /// Records the `workspace` metadata each agent was handed.
    struct WorkspaceProbeRunner {
        seen: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
    }

    #[async_trait::async_trait]
    impl AgentRunner for WorkspaceProbeRunner {
        async fn run(
            &self,
            spec: &AgentSpec,
            _task: &str,
            _runtime: &Runtime,
            _mailbox: &Mailbox,
        ) -> Result<AgentOutput, MultiError> {
            let ws = spec
                .metadata
                .get(crate::workspace::WORKSPACE_METADATA_KEY)
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();
            self.seen.lock().unwrap().push(ws.clone());
            // The directory must exist while the agent runs.
            assert!(!ws.is_empty() && std::path::Path::new(&ws).is_dir());
            Ok(AgentOutput {
                name: spec.name.clone(),
                answer: "ok".into(),
                turns: 1,
                tool_calls: 0,
                duration_ms: 1.0,
                error: None,
                outcome: None,
                tokens: None,
            })
        }
    }

    #[tokio::test]
    async fn parallel_workspaces_are_provisioned_and_distinct() {
        let base = std::env::temp_dir().join(format!("car-swarm-ws-{}", std::process::id()));
        let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
        let runner: Arc<dyn AgentRunner> = Arc::new(WorkspaceProbeRunner { seen: seen.clone() });
        let infra = SharedInfra::new();

        let agents = vec![AgentSpec::new("alice", ""), AgentSpec::new("bob", "")];
        let result = Swarm::new(agents, SwarmMode::Parallel)
            .with_workspaces(crate::workspace::WorkspaceConfig::directory(&base))
            .run("task", &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.outputs.len(), 2);
        assert!(result.outputs.iter().all(|o| o.succeeded()));
        let paths = seen.lock().unwrap().clone();
        assert_eq!(paths.len(), 2);
        assert_ne!(paths[0], paths[1], "each agent gets a distinct workspace");
        // Cleaned up after the run.
        for p in &paths {
            assert!(!std::path::Path::new(p).exists(), "workspace removed on drop");
        }
        let _ = std::fs::remove_dir_all(&base);
    }

    #[tokio::test]
    async fn parallel_budget_agent_cap_skips_excess() {
        // Five agents, cap of 2: exactly two run, three are skipped.
        let agents: Vec<AgentSpec> = (0..5)
            .map(|i| AgentSpec::new(&format!("a{}", i), ""))
            .collect();
        let runner: Arc<dyn AgentRunner> = Arc::new(MockRunner {
            call_count: AtomicU32::new(0),
        });
        let infra = SharedInfra::new().with_budget(crate::BudgetLimits {
            max_agents: Some(2),
            ..Default::default()
        });

        let result = Swarm::new(agents, SwarmMode::Parallel)
            .run("task", &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.outputs.len(), 5);
        let ran = result.outputs.iter().filter(|o| o.succeeded()).count();
        let skipped = result
            .outputs
            .iter()
            .filter(|o| crate::is_budget_skipped(o))
            .count();
        assert_eq!(ran, 2);
        assert_eq!(skipped, 3);
    }

    #[tokio::test]
    async fn test_debate_swarm() {
        let agents = vec![
            AgentSpec::new("debater_a", "Argue for"),
            AgentSpec::new("debater_b", "Argue against"),
        ];
        let runner: Arc<dyn AgentRunner> = Arc::new(MockRunner {
            call_count: AtomicU32::new(0),
        });
        let infra = SharedInfra::new();

        let result = Swarm::new(agents, SwarmMode::Debate)
            .run("debate topic", &runner, &infra)
            .await
            .unwrap();

        // 2 agents x 2 rounds = 4 outputs
        assert_eq!(result.outputs.len(), 4);
    }
}