pi-subagents/index.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import { createRunArtifacts } from "./src/artifacts.ts";
import { loadSubagentsConfig } from "./src/config.ts";
import { monitorRun } from "./src/monitor.ts";
import { listAvailableModelReferences } from "./src/models.ts";
import { createProcessSingleRunner } from "./src/process-runner.ts";
import { createConfiguredRunSingleTask } from "./src/runner.ts";
import { createSubagentParamsSchema } from "./src/schema.ts";
import { createSubagentTool } from "./src/tool.ts";
import { createTmuxSingleRunner } from "./src/tmux-runner.ts";
import {
  buildCurrentWindowArgs,
  buildKillPaneArgs,
  buildSplitWindowArgs,
  buildWrapperShellCommand,
  isInsideTmux,
} from "./src/tmux.ts";

const packageRoot = dirname(fileURLToPath(import.meta.url));
const wrapperPath = join(packageRoot, "src", "wrapper", "cli.mjs");

export default function subagentsExtension(pi: ExtensionAPI) {
  if (process.env.PI_SUBAGENTS_GITHUB_COPILOT_INITIATOR === "agent") {
    pi.registerProvider("github-copilot", {
      headers: { "X-Initiator": "agent" },
    });
  }

  // In wrapper/child sessions spawned by subagent runners we must not register
  // the subagent tool again. Provider overrides above are still allowed in child
  // runs, so the guard stays after provider registration.
  if (process.env.PI_SUBAGENTS_CHILD === "1") {
    return;
  }

  let lastRegisteredModelsKey: string | undefined;

  const tmuxRunner = createTmuxSingleRunner({
    assertInsideTmux() {
      if (!isInsideTmux()) throw new Error('tmux runner requires pi to be running inside tmux.');
    },
    async getCurrentWindowId() {
      const result = await pi.exec("tmux", buildCurrentWindowArgs());
      return result.stdout.trim();
    },
    createArtifacts: createRunArtifacts,
    buildWrapperCommand(metaPath: string) {
      return buildWrapperShellCommand({ nodePath: process.execPath, wrapperPath, metaPath });
    },
    async createPane(input) {
      const result = await pi.exec("tmux", buildSplitWindowArgs(input));
      return result.stdout.trim();
    },
    monitorRun,
    async killPane(paneId: string) {
      await pi.exec("tmux", buildKillPaneArgs(paneId));
    },
  });

  const processRunner = createProcessSingleRunner({
    createArtifacts: createRunArtifacts,
    buildWrapperSpawn(metaPath: string) {
      return { command: process.execPath, args: [wrapperPath, metaPath] };
    },
    monitorRun,
  });

  // background registry and helpers
  const registry = (typeof createBackgroundRegistry === 'function') ? createBackgroundRegistry() : undefined;
  let latestUi: { notify(message: string, type?: "info" | "warning" | "error"): void; setStatus(key: string, text: string | undefined): void } | undefined;

  function renderCounts() {
    if (!registry) return undefined;
    const counts = registry.getCounts();
    return counts.total === 0 ? undefined : `bg: ${counts.running} running / ${counts.total} total`;
  }

  function updateStatus() {
    latestUi?.setStatus("pi-subagents", renderCounts());
  }

  async function watchBackgroundRun(runId: string) {
    if (!registry) return;
    const run = registry.getRun(runId);
    if (!run || !run.paths || !run.paths.eventsPath || !run.paths.resultPath) return;
    try {
      const result = await monitorRun({ eventsPath: run.paths.eventsPath, resultPath: run.paths.resultPath });
      const status = result.stopReason === "aborted" ? "aborted" : result.exitCode === 0 ? "completed" : "failed";
      registry.recordUpdate(runId, { status, exitCode: result.exitCode, finalText: result.finalText, stopReason: result.stopReason });
      try {
        pi.appendEntry("pi-subagents:bg-update", { runId, status, finalText: result.finalText, exitCode: result.exitCode });
      } catch (e) {}
      updateStatus();
      latestUi?.notify(`Background agent ${run.preset} finished: ${result.finalText || status}`, status === "completed" ? "info" : "error");
      try {
        pi.sendMessage({
          customType: "pi-subagents:bg-complete",
          content: `Background agent ${run.preset} completed (${runId}): ${result.finalText || status}`,
          display: true,
          details: { runId, status },
        }, { triggerTurn: false });
      } catch (e) {}
    } catch (e) {
      // monitorRun errors are non-fatal here
    }
  }

  const runSingleTask = createConfiguredRunSingleTask({
    loadConfig: (cwd) => loadSubagentsConfig(cwd),
    processRunner,
    tmuxRunner,
  });

  const registerSubagentTool = (availableModels: string[]) => {
    // Do not register a tool when no models are available. Remember that the
    // last-registered key is different from the empty sentinel so that a later
    // non-empty list will still trigger registration.
    if (!availableModels || availableModels.length === 0) {
      const emptyKey = "\u0000";
      if (lastRegisteredModelsKey === emptyKey) return;
      lastRegisteredModelsKey = emptyKey;
      return;
    }

    // Create a deduplication key that is independent of the order of
    // availableModels by sorting a lowercase copy. Do not mutate
    // availableModels itself since we want to preserve the original order for
    // schema enum values.
    const key = [...availableModels].map((s) => s.toLowerCase()).sort().join("\u0000");
    if (key === lastRegisteredModelsKey) return;

    lastRegisteredModelsKey = key;
    pi.registerTool(
      createSubagentTool({
        parameters: createSubagentParamsSchema(availableModels),
        runSingleTask,
      }),
    );

    // also register background tools when models available
    try {
      pi.registerTool(createBackgroundStatusTool({ registry }));
      pi.registerTool(createBackgroundAgentTool({
        parameters: createBackgroundAgentSchema(availableModels),
        discoverSubagentPresets,
        launchDetachedTask: processRunner.launchDetachedTask,
        registerBackgroundRun(entry: any) {
          if (!registry) return { running: 0, completed: 0, failed: 0, aborted: 0, total: 0 };
          registry.recordLaunch({ runId: entry.runId, preset: entry.preset, task: entry.task, requestedModel: entry.requestedModel, resolvedModel: entry.resolvedModel, paths: entry.paths, meta: entry.meta });
          return registry.getCounts();
        },
        watchBackgroundRun(runId: string) {
          void watchBackgroundRun(runId);
        },
      }));
    } catch (e) {}
  };

  const syncSubagentTool = (ctx: { modelRegistry: { getAvailable(): Array<{ provider: string; id: string }> } }) => {
    registerSubagentTool(listAvailableModelReferences(ctx.modelRegistry));
  };

  pi.on("session_start", (_event, ctx) => {
    latestUi = ctx.ui;
    // replay persisted runs if session manager provides entries
    try {
      const entries = ctx.sessionManager?.getEntries?.() ?? [];
      const bgEntries = entries.filter((e: any) => e.type === "pi-subagents:bg-run" || e.type === "pi-subagents:bg-update");
      // convert to registry runs if possible
      const runs = bgEntries.map((be: any) => be.data?.run).filter(Boolean);
      if (registry && runs.length) registry.replay(runs);
    } catch (e) {}

    updateStatus();

    // reattach watchers for running runs
    try {
      if (registry) {
        for (const r of registry.getSnapshot({ includeCompleted: true })) {
          if (r.status === "running") void watchBackgroundRun(r.runId);
        }
      }
    } catch (e) {}

    syncSubagentTool(ctx);
  });

  pi.on("before_agent_start", (_event, ctx) => {
    latestUi = ctx.ui;
    syncSubagentTool(ctx);
  });
}