zesdex/crates/zesdex-backend/src/tool/spawn.rs

//! `spawn_agents` tool — simple interface for the main agent to fan out work
//! to multiple subagents running in parallel.
//!
//! Unlike `workflow_run` (which requires a JSON-encoded `WorkflowScript`),
//! `spawn_agents` accepts a plain list of prompt strings and automatically
//! runs them as a `Parallel` workflow. The agent just says what each
//! subagent should do, not how to encode the script.
//!
//! Also provides a pipeline variant: `spawn_pipeline` runs agents
//! sequentially so each stage sees the previous stage's findings.
use super::{Tool, ToolCtx};
use crate::app::workflow::engine::PrimitiveCtx;
use crate::app::workflow::script::{ScriptOptions, ScriptPrimitive, WorkflowScript};
use anyhow::{anyhow, Result};
use serde_json::{json, Value};
use std::collections::HashMap;

/// Fan out a list of prompts to independent parallel subagents.
pub struct SpawnAgents;

impl Tool for SpawnAgents {
    fn name(&self) -> &'static str {
        "spawn_agents"
    }

    fn description(&self) -> &'static str {
        "Fan out independent subtasks to multiple Hive nodes running in PARALLEL. \
        Pass a list of prompt strings — each becomes one autonomous node with \
        access to all tools. Use this whenever a task has independent parts that do \
        not need each other's output (e.g. analysing multiple files simultaneously, \
        writing multiple independent modules, parallel verification). \
        Results from all nodes are returned together. \
        Use spawn_pipeline instead when each stage needs the previous stage's output."
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "agents": {
                    "type": "array",
                    "description": "List of prompt strings, one per Hive node. Each node runs independently and in parallel.",
                    "items": { "type": "string" },
                    "minItems": 2
                },
                "max_concurrency": {
                    "type": "integer",
                    "description": "Maximum number of agents to run simultaneously (default: 10, max: 10).",
                    "default": 10
                }
            },
            "required": ["agents"]
        })
    }

    fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
        use std::sync::{Arc, Mutex};
        let agents: Vec<String> = args
            .get("agents")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow!("missing required argument: agents"))?
            .iter()
            .filter_map(|v| v.as_str().map(std::string::ToString::to_string))
            .collect();

        if agents.is_empty() {
            return Err(anyhow!("agents list must not be empty"));
        }
        if agents.len() == 1 {
            return Err(anyhow!(
                "use a single agent tool call for one task; spawn_agents is for 2+ parallel tasks"
            ));
        }

        let max_concurrency = args
            .get("max_concurrency")
            .and_then(serde_json::Value::as_u64)
            .map_or(10, |v| v.min(10) as usize);

        let agent_count = agents.len();
        let primitives: Vec<ScriptPrimitive> =
            agents.into_iter().map(ScriptPrimitive::Agent).collect();

        let wf = WorkflowScript {
            name: format!("parallel-{agent_count}-agents"),
            description: format!("Auto-spawned parallel workflow with {agent_count} agents"),
            script: ScriptPrimitive::Parallel(primitives),
            options: ScriptOptions {
                max_concurrency,
                continue_on_error: true,
                timeout_ms: None,
            },
        };

        let live: Option<crate::app::workflow::engine::LiveStateFn> =
            ctx.turn_events.as_ref().map(|turn_events| {
                let turn_events = turn_events.clone();
                let f: crate::app::workflow::engine::LiveStateFn =
                    Arc::new(move |agent_id: String, agent_name: String, status| {
                        if let Ok(mut q) = turn_events.lock() {
                            q.push_back(
                                crate::app::state::runtime::TurnEvent::WorkflowAgentUpdate {
                                    agent_id,
                                    agent_name,
                                    status,
                                },
                            );
                        }
                    });
                f
            });

        // Create a per-invocation findings scope so subagents spawned
        // by this tool call are isolated from any other concurrent
        // spawn_agents or workflow_run invocations.
        let findings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
        let results = crate::app::workflow::engine::execute_primitive(PrimitiveCtx {
            primitive: &wf.script,
            args: &HashMap::new(),
            concurrency_cap: max_concurrency,
            continue_on_error: true,
            abort_flag: &no_abort,
            live: live.as_ref(),
            session_dir: &ctx.session_dir,
            workspaces: &ctx.workspaces,
            findings: &findings,
            timeout_ms: None,
        })?;
        Ok(format_results(&results, "parallel"))
    }
}

/// Run agents sequentially in a pipeline — each stage sees previous findings.
pub struct SpawnPipeline;

impl Tool for SpawnPipeline {
    fn name(&self) -> &'static str {
        "spawn_pipeline"
    }

    fn description(&self) -> &'static str {
        "Run Hive nodes SEQUENTIALLY in a pipeline — each stage sees findings \
        shared by previous stages via note_finding. Use when stages build on each \
        other (e.g. 'research -> plan -> implement -> test'). \
        Use spawn_agents instead when tasks are truly independent and order does not matter."
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "stages": {
                    "type": "array",
                    "description": "Ordered list of prompt strings — each stage is a Hive node that runs after the previous one completes. Stages can call note_finding() to pass data to later stages.",
                    "items": { "type": "string" },
                    "minItems": 2
                }
            },
            "required": ["stages"]
        })
    }

    fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
        use std::sync::{Arc, Mutex};
        let stages: Vec<String> = args
            .get("stages")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow!("missing required argument: stages"))?
            .iter()
            .filter_map(|v| v.as_str().map(std::string::ToString::to_string))
            .collect();

        if stages.is_empty() {
            return Err(anyhow!("stages list must not be empty"));
        }

        let primitives: Vec<ScriptPrimitive> =
            stages.into_iter().map(ScriptPrimitive::Agent).collect();

        let wf = WorkflowScript {
            name: "pipeline".to_string(),
            description: "Auto-spawned pipeline workflow".to_string(),
            script: ScriptPrimitive::Pipeline(primitives),
            options: ScriptOptions {
                max_concurrency: 1,
                continue_on_error: false,
                timeout_ms: None,
            },
        };

        let live: Option<crate::app::workflow::engine::LiveStateFn> =
            ctx.turn_events.as_ref().map(|turn_events| {
                let turn_events = turn_events.clone();
                let f: crate::app::workflow::engine::LiveStateFn =
                    Arc::new(move |agent_id: String, agent_name: String, status| {
                        if let Ok(mut q) = turn_events.lock() {
                            q.push_back(
                                crate::app::state::runtime::TurnEvent::WorkflowAgentUpdate {
                                    agent_id,
                                    agent_name,
                                    status,
                                },
                            );
                        }
                    });
                f
            });

        // Per-invocation findings scope isolates this pipeline from any
        // other concurrent spawn_agents / spawn_pipeline / workflow_run.
        let findings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
        let results = crate::app::workflow::engine::execute_primitive(PrimitiveCtx {
            primitive: &wf.script,
            args: &HashMap::new(),
            concurrency_cap: 1,
            continue_on_error: false,
            abort_flag: &no_abort,
            live: live.as_ref(),
            session_dir: &ctx.session_dir,
            workspaces: &ctx.workspaces,
            findings: &findings,
            timeout_ms: None,
        })?;
        Ok(format_results(&results, "pipeline"))
    }
}

/// Format a list of agent results into a readable summary string.
fn format_results(results: &[String], mode: &str) -> String {
    if results.is_empty() {
        return format!("{mode} workflow completed with no output");
    }
    let formatted: Vec<String> = results
        .iter()
        .enumerate()
        .map(|(i, r)| format!("=== Agent {} ===\n{}", i + 1, r.trim()))
        .collect();
    formatted.join("\n\n")
}