305 lines
12 KiB
TypeScript
305 lines
12 KiB
TypeScript
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
import { Type } from "typebox";
|
|
import { loadAgents } from "./agents.ts";
|
|
import { loadConfig, resolveSpawn, type Diagnostics } from "./config.ts";
|
|
import { SubprocessRpcRunner } from "./runner.ts";
|
|
import { Supervisor } from "./supervisor.ts";
|
|
import { milestoneNotification } from "./status.ts";
|
|
import type { SpawnRequest, SubagentStatus } from "./types.ts";
|
|
import { widget } from "./ui.ts";
|
|
|
|
let supervisor: Supervisor | undefined;
|
|
let lastDiagnostics: Diagnostics = { warnings: [] };
|
|
let lastStatuses: SubagentStatus[] = [];
|
|
let uiExpanded = false;
|
|
|
|
export default function subagents(pi: ExtensionAPI) {
|
|
const getSupervisor = (ctx: ExtensionContext): Supervisor => {
|
|
if (supervisor) return supervisor;
|
|
const diagnostics: Diagnostics = { warnings: [] };
|
|
const cwd = cwdOf(ctx);
|
|
const config = loadConfig(cwd, isProjectTrusted(ctx), diagnostics);
|
|
lastDiagnostics = diagnostics;
|
|
uiExpanded = config.ui.defaultExpanded;
|
|
supervisor = new Supervisor(new SubprocessRpcRunner(), cwd, {
|
|
maxConcurrent: config.maxConcurrent,
|
|
recentTerminalTtlMs: config.recentTerminalTtlMs,
|
|
onMilestone: (status, event) => {
|
|
pi.appendEntry("subagent_milestone", { event, status });
|
|
const notification = milestoneNotification(status, event);
|
|
if (notification) ctx.ui?.notify?.(notification.message, notification.level);
|
|
},
|
|
onChange: (statuses) => {
|
|
lastStatuses = statuses;
|
|
updateUi(ctx, config.ui.enabled);
|
|
},
|
|
});
|
|
updateUi(ctx, config.ui.enabled);
|
|
return supervisor;
|
|
};
|
|
|
|
const resolve = (ctx: ExtensionContext, request: SpawnRequest): SpawnRequest => {
|
|
const diagnostics: Diagnostics = { warnings: [] };
|
|
const cwd = cwdOf(ctx);
|
|
const trusted = isProjectTrusted(ctx);
|
|
const config = loadConfig(cwd, trusted, diagnostics);
|
|
const agents = loadAgents(cwd, trusted, diagnostics);
|
|
lastDiagnostics = diagnostics;
|
|
const resolved = resolveSpawn(request, config, agents);
|
|
if (resolved.context === "fork") resolved.parentSessionFile = ctx.sessionManager.getSessionFile();
|
|
return resolved;
|
|
};
|
|
|
|
pi.registerTool({
|
|
name: "subagent_spawn",
|
|
label: "Spawn subagent",
|
|
description: "Start one ad hoc independent subagent and return immediately with its child id",
|
|
parameters: Type.Object({
|
|
prompt: Type.String({ description: "Prompt for the delegated subagent" }),
|
|
label: Type.Optional(Type.String({ description: "Human-readable label for this work item" })),
|
|
agent: Type.Optional(Type.String({ description: "Named agent definition to use" })),
|
|
context: Type.Optional(Type.Union([Type.Literal("independent"), Type.Literal("fork")])),
|
|
model: Type.Optional(Type.String({ description: "Optional model selector for the child" })),
|
|
thinking: Type.Optional(Type.String({ description: "Optional thinking level for the child" })),
|
|
tools: Type.Optional(Type.String({ description: "Tool profile name" })),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
const accepted = getSupervisor(ctx).spawn(resolve(ctx, params as SpawnRequest));
|
|
ctx.ui?.notify?.(`Started subagent ${accepted.label}`, "info");
|
|
return textResult(accepted);
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_batch",
|
|
label: "Spawn subagent batch",
|
|
description: "Start multiple subagents and return immediately with accepted child ids and per-entry failures",
|
|
parameters: Type.Object({
|
|
subagents: Type.Array(
|
|
Type.Object({
|
|
prompt: Type.String({ description: "Prompt for the delegated subagent" }),
|
|
label: Type.Optional(Type.String({ description: "Human-readable label for this work item" })),
|
|
agent: Type.Optional(Type.String({ description: "Named agent definition to use" })),
|
|
context: Type.Optional(Type.Union([Type.Literal("independent"), Type.Literal("fork")])),
|
|
model: Type.Optional(Type.String({ description: "Optional model selector for the child" })),
|
|
thinking: Type.Optional(Type.String({ description: "Optional thinking level for the child" })),
|
|
tools: Type.Optional(Type.String({ description: "Tool profile name" })),
|
|
}),
|
|
),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
const requests = Array.isArray((params as { subagents?: unknown }).subagents) ? ((params as { subagents: SpawnRequest[] }).subagents) : [];
|
|
const accepted: SpawnRequest[] = [];
|
|
const failed: Array<{ index: number; error: string }> = [];
|
|
requests.forEach((request, index) => {
|
|
try {
|
|
accepted.push(resolve(ctx, request));
|
|
} catch (error) {
|
|
failed.push({ index, error: error instanceof Error ? error.message : String(error) });
|
|
}
|
|
});
|
|
const result = getSupervisor(ctx).spawnBatch(accepted);
|
|
return textResult({ accepted: result.accepted, failed: [...failed, ...result.failed] });
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_list",
|
|
label: "List subagents",
|
|
description: "List active and recent subagents for this parent session",
|
|
parameters: Type.Object({}),
|
|
async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
|
|
return textResult(getSupervisor(ctx).list());
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_status",
|
|
label: "Get subagent status",
|
|
description: "Get current lifecycle status for one subagent",
|
|
parameters: Type.Object({
|
|
id: Type.String({ description: "Subagent id returned by subagent_spawn" }),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
return textResult(getSupervisor(ctx).status(String((params as { id: unknown }).id)));
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_result",
|
|
label: "Get subagent result",
|
|
description: "Return still-running before completion and the final answer after completion",
|
|
parameters: Type.Object({
|
|
id: Type.String({ description: "Subagent id returned by subagent_spawn" }),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
return textResult(getSupervisor(ctx).result(String((params as { id: unknown }).id)));
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_wait",
|
|
label: "Wait for subagents",
|
|
description: "Block until multiple subagents are terminal or a timeout expires. Prefer setting timeoutMs so the parent turn cannot hang forever",
|
|
parameters: Type.Object({
|
|
ids: Type.Array(Type.String({ description: "Subagent id returned by subagent_spawn or subagent_batch" })),
|
|
timeoutMs: Type.Optional(Type.Number({ description: "Maximum milliseconds to wait. Omit or use 0 to wait indefinitely" })),
|
|
mode: Type.Optional(Type.Union([Type.Literal("all"), Type.Literal("any")], { description: "Wait for all ids by default, or return after any id is terminal" })),
|
|
}),
|
|
async execute(_toolCallId, params, signal, _onUpdate, ctx) {
|
|
const input = params as { ids?: unknown; timeoutMs?: unknown; mode?: unknown };
|
|
const ids = Array.isArray(input.ids) ? input.ids.map(String) : [];
|
|
const timeoutMs = typeof input.timeoutMs === "number" && Number.isFinite(input.timeoutMs) ? input.timeoutMs : undefined;
|
|
const mode = input.mode === "any" ? "any" : "all";
|
|
return textResult(await getSupervisor(ctx).wait(ids, { timeoutMs, mode, signal }));
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "subagent_cancel",
|
|
label: "Cancel subagent",
|
|
description: "Cancel a running subagent",
|
|
parameters: Type.Object({
|
|
id: Type.String({ description: "Subagent id returned by subagent_spawn" }),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
return textResult(await getSupervisor(ctx).cancel(String((params as { id: unknown }).id)));
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-spawn", {
|
|
description: "Start an ad hoc independent subagent",
|
|
handler: async (args, ctx) => {
|
|
const accepted = getSupervisor(ctx).spawn(resolve(ctx, parseSpawnArgs(args)));
|
|
ctx.ui.notify(`Started subagent ${accepted.label}`, "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-batch", {
|
|
description: "Start ad hoc independent subagents split by |",
|
|
handler: async (args, ctx) => {
|
|
const requests = args
|
|
.split("|")
|
|
.map((prompt) => prompt.trim())
|
|
.filter(Boolean)
|
|
.map((prompt) => resolve(ctx, { prompt }));
|
|
ctx.ui.notify(JSON.stringify(getSupervisor(ctx).spawnBatch(requests), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-list", {
|
|
description: "Show subagent status records",
|
|
handler: async (_args, ctx) => {
|
|
ctx.ui.notify(JSON.stringify(getSupervisor(ctx).list(), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-status", {
|
|
description: "Show a subagent status by id",
|
|
handler: async (args, ctx) => {
|
|
ctx.ui.notify(JSON.stringify(getSupervisor(ctx).status(args.trim()), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-result", {
|
|
description: "Show a subagent result by id",
|
|
handler: async (args, ctx) => {
|
|
ctx.ui.notify(JSON.stringify(getSupervisor(ctx).result(args.trim()), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-wait", {
|
|
description: "Wait for subagent ids separated by spaces",
|
|
handler: async (args, ctx) => {
|
|
const { ids, timeoutMs, mode } = parseWaitArgs(args);
|
|
ctx.ui.notify(JSON.stringify(await getSupervisor(ctx).wait(ids, { timeoutMs, mode }), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-ui", {
|
|
description: "Toggle the bundled subagent status inspector",
|
|
handler: async (_args, ctx) => {
|
|
uiExpanded = !uiExpanded;
|
|
updateUi(ctx, true);
|
|
ctx.ui.notify(`Subagent inspector ${uiExpanded ? "expanded" : "collapsed"}`, "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-diagnostics", {
|
|
description: "Show subagent configuration diagnostics from the last load",
|
|
handler: async (_args, ctx) => {
|
|
ctx.ui.notify(JSON.stringify(lastDiagnostics, null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("subagent-cancel", {
|
|
description: "Cancel a running subagent by id",
|
|
handler: async (args, ctx) => {
|
|
ctx.ui.notify(JSON.stringify(await getSupervisor(ctx).cancel(args.trim()), null, 2), "info");
|
|
},
|
|
});
|
|
|
|
pi.on("session_shutdown", async () => {
|
|
await supervisor?.shutdown();
|
|
supervisor = undefined;
|
|
});
|
|
}
|
|
|
|
function updateUi(ctx: ExtensionContext, enabled: boolean) {
|
|
if (!ctx.hasUI) return;
|
|
ctx.ui.setWidget("subagents", enabled ? widget(lastStatuses, uiExpanded) : undefined);
|
|
}
|
|
|
|
function parseSpawnArgs(args: string): SpawnRequest {
|
|
const parts = args.trim().split(/\s+/u);
|
|
const request: Partial<SpawnRequest> = {};
|
|
while (parts.length >= 2 && parts[0].startsWith("--")) {
|
|
const flag = parts.shift();
|
|
const value = parts.shift();
|
|
if (flag === "--agent") request.agent = value;
|
|
else if (flag === "--label") request.label = value;
|
|
else if (flag === "--context" && (value === "independent" || value === "fork")) request.context = value;
|
|
else if (flag === "--tools") request.tools = value;
|
|
else if (flag === "--model") request.model = value;
|
|
else if (flag === "--thinking") request.thinking = value;
|
|
}
|
|
return { ...request, prompt: parts.join(" ") || args } as SpawnRequest;
|
|
}
|
|
|
|
function parseWaitArgs(args: string): { ids: string[]; timeoutMs?: number; mode?: "all" | "any" } {
|
|
const parts = args.trim().split(/\s+/u).filter(Boolean);
|
|
let timeoutMs: number | undefined;
|
|
let mode: "all" | "any" | undefined;
|
|
const ids: string[] = [];
|
|
while (parts.length > 0) {
|
|
const part = parts.shift();
|
|
if (!part) continue;
|
|
if (part === "--timeout-ms" && parts[0]) {
|
|
const parsed = Number(parts.shift());
|
|
if (Number.isFinite(parsed)) timeoutMs = parsed;
|
|
} else if (part === "--mode" && (parts[0] === "all" || parts[0] === "any")) {
|
|
mode = parts.shift() as "all" | "any";
|
|
} else {
|
|
ids.push(part);
|
|
}
|
|
}
|
|
return { ids, timeoutMs, mode };
|
|
}
|
|
|
|
function isProjectTrusted(ctx: ExtensionContext): boolean {
|
|
const value = (ctx as unknown as { isProjectTrusted?: () => boolean }).isProjectTrusted?.();
|
|
return value === true;
|
|
}
|
|
|
|
function cwdOf(ctx: ExtensionContext): string {
|
|
const sessionCwd = (ctx as unknown as { sessionManager?: { getCwd?: () => string }; cwd?: string }).sessionManager?.getCwd?.();
|
|
return sessionCwd ?? (ctx as unknown as { cwd?: string }).cwd ?? process.cwd();
|
|
}
|
|
|
|
function textResult(value: unknown) {
|
|
return {
|
|
content: [{ type: "text" as const, text: JSON.stringify(value, null, 2) }],
|
|
details: value,
|
|
};
|
|
}
|