Agent-JAE/default-extensions/swarm.ts

import type { ExtensionAPI } from "@jaeswift/jae-coding-agent";
import { execSync, spawn } from "node:child_process";
import { writeFileSync, readFileSync, mkdirSync, existsSync } from "node:fs";
import { join } from "node:path";
import { Type } from "@sinclair/typebox";

interface SwarmTask {
  id: number;
  description: string;
  status: "pending" | "running" | "done" | "error";
  result: string;
}

export default function (pi: ExtensionAPI) {
  let taskCounter = 0;

  async function runSubtask(
    description: string,
    cwd: string,
  ): Promise<{ success: boolean; output: string }> {
    return new Promise((resolve) => {
      const timeout = 120_000; // 2 minutes max per subtask
      let output = "";
      let resolved = false;

      try {
        const child = spawn(
          "jae",
          ["-p", description],
          {
            cwd,
            stdio: ["ignore", "pipe", "pipe"],
            timeout,
            env: { ...process.env },
          },
        );

        child.stdout?.on("data", (data: Buffer) => {
          output += data.toString();
        });

        child.stderr?.on("data", (data: Buffer) => {
          output += data.toString();
        });

        child.on("close", (code) => {
          if (!resolved) {
            resolved = true;
            resolve({ success: code === 0, output: output.trim() || "(no output)" });
          }
        });

        child.on("error", (err) => {
          if (!resolved) {
            resolved = true;
            resolve({ success: false, output: `Process error: ${err.message}` });
          }
        });

        setTimeout(() => {
          if (!resolved) {
            resolved = true;
            try { child.kill("SIGTERM"); } catch {}
            resolve({ success: false, output: output.trim() + "\n(timed out)" });
          }
        }, timeout);
      } catch (err: any) {
        resolve({ success: false, output: `Spawn error: ${err.message}` });
      }
    });
  }

  pi.registerTool({
    name: "swarm_task",
    label: "Swarm Task",
    description:
      "Split a complex task into subtasks and run them in parallel using multiple JAE instances. Provide the main task and a list of subtask descriptions.",
    parameters: Type.Object({
      task: Type.String({ description: "Main task description" }),
      subtasks: Type.Array(Type.String(), {
        description: "List of subtask descriptions to run in parallel",
      }),
    }),
    execute: async (args, ctx) => {
      const { task, subtasks } = args;
      if (!subtasks || subtasks.length === 0) {
        return { error: "No subtasks provided" };
      }

      const tasks: SwarmTask[] = subtasks.map((desc: string) => ({
        id: ++taskCounter,
        description: desc,
        status: "pending" as const,
        result: "",
      }));

      ctx.ui.notify(
        `\u{1F41D} Swarm: Launching ${tasks.length} subtasks for: ${task}`,
        "info",
      );

      // Run all in parallel
      const promises = tasks.map(async (t) => {
        t.status = "running";
        const res = await runSubtask(t.description, ctx.cwd);
        t.status = res.success ? "done" : "error";
        t.result = res.output;
        return t;
      });

      const results = await Promise.all(promises);

      const summary = results
        .map(
          (t) =>
            `### Subtask #${t.id}: ${t.status === "done" ? "\u2705" : "\u274C"} ${t.description}\n${t.result.substring(0, 500)}`,
        )
        .join("\n\n");

      const successCount = results.filter((t) => t.status === "done").length;

      return {
        output: `# Swarm Results\n\nTask: ${task}\nCompleted: ${successCount}/${tasks.length}\n\n${summary}`,
      };
    },
  });

  pi.registerCommand("swarm", {
    description:
      "Run parallel subtasks: /swarm <task description> (JAE will decompose into subtasks)",
    handler: async (args, ctx) => {
      const task = args.trim();
      if (!task) {
        ctx.ui.notify(
          "Usage: /swarm <task description>\n\nJAE will decompose the task into parallel subtasks and run them simultaneously.",
          "info",
        );
        return;
      }

      ctx.ui.notify(
        `\u{1F41D} Swarm mode: Submit this task to JAE and it will use the swarm_task tool to parallelize it.\n\nTask: ${task}`,
        "info",
      );
    },
  });
}