Files
pi-subagents/index.ts

115 lines
4.0 KiB
TypeScript

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-backed subagents require 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,
});
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,
}),
);
};
const syncSubagentTool = (ctx: { modelRegistry: { getAvailable(): Array<{ provider: string; id: string }> } }) => {
registerSubagentTool(listAvailableModelReferences(ctx.modelRegistry));
};
pi.on("session_start", (_event, ctx) => {
syncSubagentTool(ctx);
});
pi.on("before_agent_start", (_event, ctx) => {
syncSubagentTool(ctx);
});
}