100 lines
3.6 KiB
TypeScript
100 lines
3.6 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 { monitorRun } from "./src/monitor.ts";
|
|
import { listAvailableModelReferences } from "./src/models.ts";
|
|
import { createTmuxSingleRunner } from "./src/runner.ts";
|
|
import {
|
|
buildCurrentWindowArgs,
|
|
buildKillPaneArgs,
|
|
buildSplitWindowArgs,
|
|
buildWrapperShellCommand,
|
|
isInsideTmux,
|
|
} from "./src/tmux.ts";
|
|
import { createSubagentParamsSchema } from "./src/schema.ts";
|
|
import { createSubagentTool } from "./src/tool.ts";
|
|
|
|
const packageRoot = dirname(fileURLToPath(import.meta.url));
|
|
const wrapperPath = join(packageRoot, "src", "wrapper", "cli.mjs");
|
|
|
|
export default function tmuxSubagentExtension(pi: ExtensionAPI) {
|
|
if (process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR === "agent") {
|
|
pi.registerProvider("github-copilot", {
|
|
headers: { "X-Initiator": "agent" },
|
|
});
|
|
}
|
|
|
|
// In wrapper/child sessions spawned by the tmux runner we must not register the
|
|
// subagent tool (that would cause nested subagent registrations). Skip all
|
|
// subagent-tool registration logic when PI_TMUX_SUBAGENT_CHILD is set. Provider
|
|
// overrides (above) are still allowed in child runs, so the guard is placed
|
|
// after provider registration.
|
|
if (process.env.PI_TMUX_SUBAGENT_CHILD === "1") {
|
|
return;
|
|
}
|
|
|
|
let lastRegisteredModelsKey: string | undefined;
|
|
|
|
const runSingleTask = 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 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);
|
|
});
|
|
}
|