dispatch.index.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { spawn } from "node:child_process";
import { createWriteStream, mkdirSync, rmSync, existsSync, readdirSync, unlinkSync } from "node:fs";
import { join } from "node:path";
import { homedir } from "node:os";

const DISPATCH_LOG_DIR = join(homedir(), ".pi", "agent", "dispatch");

function ensureLogDir(): void {
  mkdirSync(DISPATCH_LOG_DIR, { recursive: true });
}

function logPaths(pid: number): { stdout: string; stderr: string } {
  return {
    stdout: join(DISPATCH_LOG_DIR, `${pid}-stdout.log`),
    stderr: join(DISPATCH_LOG_DIR, `${pid}-stderr.log`),
  };
}

function cleanupLogs(pid: number): void {
  const paths = logPaths(pid);
  try { unlinkSync(paths.stdout); } catch {}
  try { unlinkSync(paths.stderr); } catch {}
}

function cleanupAllLogs(): void {
  if (!existsSync(DISPATCH_LOG_DIR)) return;
  try {
    for (const f of readdirSync(DISPATCH_LOG_DIR)) {
      try { unlinkSync(join(DISPATCH_LOG_DIR, f)); } catch {}
    }
  } catch {}
}

function normalizeCommand(command: string): { effectiveCommand: string; rewriteNote?: string } {
  const loginBashPattern = /\bbash\s+(-lc|-l\s+-c)\b/g;
  if (!loginBashPattern.test(command)) {
    return { effectiveCommand: command };
  }

  const effectiveCommand = command.replace(loginBashPattern, "bash --noprofile --norc -c");
  return {
    effectiveCommand,
    rewriteNote: "Rewrote `bash -lc`/`bash -l -c` to `bash --noprofile --norc -c` to avoid profile side effects.",
  };
}

export default function (pi: ExtensionAPI) {
  const active = new Map<
    number,
    { originalCommand: string; effectiveCommand: string; proc: ReturnType<typeof spawn> }
  >();

  pi.registerTool({
    name: "dispatch",
    label: "Dispatch",
    description:
      "Run a shell command in the background. Returns immediately with a PID and log file paths. " +
      "When the process exits, you are notified with its output via a new turn.",
    promptSnippet: "Run a background command; notified on completion via triggerTurn",
    promptGuidelines: [
      "Use dispatch for long-running commands (builds, deploys, pipeline runs) instead of bash with long timeouts.",
      "dispatch returns immediately — do NOT poll or wait. You will be woken up when the process exits.",
      "Avoid login shells (`bash -lc`) in dispatched commands; use `bash -c` or `bash --noprofile --norc -c`.",
      "dispatch writes stdout/stderr to log files at `~/.pi/agent/dispatch/{pid}-stdout.log` and `{pid}-stderr.log`. Use `read` on these paths to check on a running process mid-flight.",
      "You can dispatch sub-agent pi processes for background research and side tasks: `pi -p \"<prompt>\" --model claude-sonnet-4-6`. The sub-agent runs headless with full tool access (read/edit/bash) and returns its final response to stdout. Use this for parallel investigations, code analysis, or generating content that doesn't need interactive back-and-forth.",
      "For parallel work, dispatch multiple sub-agents at once — each gets its own log files. You'll be notified as each completes, and can `read` their logs to check progress or collect results.",
    ],
    parameters: Type.Object({
      command: Type.String({ description: "Shell command to run" }),
      timeout: Type.Optional(Type.Number({ description: "Timeout in seconds (optional)" })),
    }),

    async execute(_toolCallId, params) {
      const { command, timeout } = params;
      const { effectiveCommand, rewriteNote } = normalizeCommand(command);

      ensureLogDir();

      const proc = spawn("/bin/sh", ["-c", effectiveCommand], {
        stdio: ["ignore", "pipe", "pipe"],
        cwd: pi.cwd,
      });

      if (!proc.pid) {
        return {
          content: [{ type: "text", text: "Failed to spawn process" }],
          details: {},
          isError: true,
        };
      }

      const pid = proc.pid;
      const paths = logPaths(pid);
      const stdoutStream = createWriteStream(paths.stdout);
      const stderrStream = createWriteStream(paths.stderr);

      let stdout = "";
      let stderr = "";
      proc.stdout!.on("data", (d: Buffer) => {
        const chunk = d.toString();
        stdout += chunk;
        stdoutStream.write(d);
      });
      proc.stderr!.on("data", (d: Buffer) => {
        const chunk = d.toString();
        stderr += chunk;
        stderrStream.write(d);
      });

      active.set(pid, { originalCommand: command, effectiveCommand, proc });

      let timedOut = false;
      let timer: ReturnType<typeof setTimeout> | undefined;
      if (timeout && timeout > 0) {
        timer = setTimeout(() => {
          timedOut = true;
          proc.kill();
        }, timeout * 1000);
      }

      proc.on("close", (exitCode, signal) => {
        if (timer) clearTimeout(timer);
        active.delete(pid);
        stdoutStream.end();
        stderrStream.end();

        const tail = (s: string) => (s.length > 4096 ? "…" + s.slice(-4096) : s);

        const parts = [
          `**dispatch completed** (pid ${pid})`,
          `Command: \`${command}\`${effectiveCommand !== command ? `\nEffective command: \`${effectiveCommand}\`` : ""}`,
          `Exit: ${exitCode}${signal ? ` (signal: ${signal})` : ""}${timedOut ? " (timed out)" : ""}`,
          `Logs: \`${paths.stdout}\` | \`${paths.stderr}\``,
        ];
        if (rewriteNote) parts.splice(1, 0, rewriteNote);
        if (stdout.trim()) parts.push(`**stdout (tail):**\n\`\`\`\n${tail(stdout.trim())}\n\`\`\``);
        if (stderr.trim()) parts.push(`**stderr (tail):**\n\`\`\`\n${tail(stderr.trim())}\n\`\`\``);

        pi.sendMessage(
          {
            customType: "dispatch-complete",
            content: parts.join("\n\n"),
            display: true,
            details: { pid, command, effectiveCommand, exitCode, signal, timedOut, rewriteNote, logPaths: paths },
          },
          { triggerTurn: true },
        );
      });

      const responseLines = [
        rewriteNote ?? "",
        `Dispatched (pid ${pid}): \`${effectiveCommand}\``,
        `stdout log: \`${paths.stdout}\``,
        `stderr log: \`${paths.stderr}\``,
      ].filter(Boolean);

      return {
        content: [{ type: "text", text: responseLines.join("\n") }],
        details: { pid, command, effectiveCommand, rewriteNote, logPaths: paths },
      };
    },
  });

  pi.on("session_shutdown", async () => {
    for (const [, { proc }] of active) {
      proc.kill();
    }
    active.clear();
    cleanupAllLogs();
  });
}