zesdex/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 serde_json::{json, Value};
use anyhow::{Result, anyhow};
use std::collections::HashMap;
use super::{Tool, ToolCtx};
use crate::app::workflow::script::{ScriptPrimitive, ScriptOptions, WorkflowScript};

/// 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 subagents running in PARALLEL. \
        Pass a list of prompt strings — each becomes one autonomous subagent 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 agents 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 subagent. Each subagent 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 results = crate::app::workflow::engine::execute_primitive(
            &wf.script,
            &HashMap::new(),
            max_concurrency,
            true,
            live.as_ref(),
            &ctx.session_dir,
            &ctx.workspaces,
            &findings,
            None, // no per-agent timeout for spawn_agents
        )?;
        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 subagents 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 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 results = crate::app::workflow::engine::execute_primitive(
            &wf.script,
            &HashMap::new(),
            1,
            false,
            live.as_ref(),
            &ctx.session_dir,
            &ctx.workspaces,
            &findings,
            None, // no per-agent timeout for spawn_pipeline
        )?;
        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")
}