336 lines
10 KiB
TypeScript
336 lines
10 KiB
TypeScript
import { spawn } from "node:child_process";
|
|
import * as fs from "node:fs";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import { fileURLToPath } from "node:url";
|
|
import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
|
|
import type { AgentConfig } from "./agents.js";
|
|
import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js";
|
|
import {
|
|
applyAssistantUsage,
|
|
createEmptyUsageStats,
|
|
getFinalOutput,
|
|
type OnUpdateCallback,
|
|
type SingleResult,
|
|
type SubagentDetails,
|
|
} from "./subagent-types.js";
|
|
|
|
interface RpcResponse {
|
|
type: "response";
|
|
id?: string;
|
|
command: string;
|
|
success: boolean;
|
|
error?: string;
|
|
}
|
|
|
|
interface PendingRequest {
|
|
resolve: (response: RpcResponse) => void;
|
|
reject: (error: Error) => void;
|
|
timeout: NodeJS.Timeout;
|
|
}
|
|
|
|
type RpcExtensionUIRequest =
|
|
| { type: "extension_ui_request"; id: string; method: "select"; title: string; options: string[]; timeout?: number }
|
|
| { type: "extension_ui_request"; id: string; method: "confirm"; title: string; message: string; timeout?: number }
|
|
| { type: "extension_ui_request"; id: string; method: "input"; title: string; placeholder?: string; timeout?: number }
|
|
| { type: "extension_ui_request"; id: string; method: "editor"; title: string; prefill?: string }
|
|
| { type: "extension_ui_request"; id: string; method: "notify"; message: string; notifyType?: "info" | "warning" | "error" }
|
|
| { type: "extension_ui_request"; id: string; method: "setStatus"; statusKey: string; statusText?: string }
|
|
| {
|
|
type: "extension_ui_request";
|
|
id: string;
|
|
method: "setWidget";
|
|
widgetKey: string;
|
|
widgetLines?: string[];
|
|
widgetPlacement?: "aboveEditor" | "belowEditor";
|
|
}
|
|
| { type: "extension_ui_request"; id: string; method: "setTitle"; title: string }
|
|
| { type: "extension_ui_request"; id: string; method: "set_editor_text"; text: string };
|
|
|
|
function writePromptToTempFile(agentName: string, prompt: string): { dir: string; filePath: string } {
|
|
const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-subagent-rpc-"));
|
|
const safeName = agentName.replace(/[^\w.-]+/g, "_");
|
|
const filePath = path.join(tmpDir, `prompt-${safeName}.md`);
|
|
fs.writeFileSync(filePath, prompt, { encoding: "utf-8", mode: 0o600 });
|
|
return { dir: tmpDir, filePath };
|
|
}
|
|
|
|
function terminateProcess(proc: ReturnType<typeof spawn>): void {
|
|
if (proc.exitCode !== null) return;
|
|
proc.kill("SIGTERM");
|
|
setTimeout(() => {
|
|
if (proc.exitCode === null) proc.kill("SIGKILL");
|
|
}, 3000);
|
|
}
|
|
|
|
function safeJsonParse(line: string): any | undefined {
|
|
if (!line.trim()) return undefined;
|
|
try {
|
|
return JSON.parse(line);
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
async function bridgeExtensionUiRequest(
|
|
request: RpcExtensionUIRequest,
|
|
ctx: ExtensionContext,
|
|
signal: AbortSignal | undefined,
|
|
sendJson: (obj: unknown) => void,
|
|
): Promise<void> {
|
|
const sendCancelled = () => sendJson({ type: "extension_ui_response", id: request.id, cancelled: true });
|
|
|
|
try {
|
|
switch (request.method) {
|
|
case "select": {
|
|
const value = await ctx.ui.select(request.title, request.options, { signal, timeout: request.timeout });
|
|
if (value === undefined) sendCancelled();
|
|
else sendJson({ type: "extension_ui_response", id: request.id, value });
|
|
return;
|
|
}
|
|
case "confirm": {
|
|
const confirmed = await ctx.ui.confirm(request.title, request.message, { signal, timeout: request.timeout });
|
|
sendJson({ type: "extension_ui_response", id: request.id, confirmed });
|
|
return;
|
|
}
|
|
case "input": {
|
|
const value = await ctx.ui.input(request.title, request.placeholder, { signal, timeout: request.timeout });
|
|
if (value === undefined) sendCancelled();
|
|
else sendJson({ type: "extension_ui_response", id: request.id, value });
|
|
return;
|
|
}
|
|
case "editor": {
|
|
const value = await ctx.ui.editor(request.title, request.prefill ?? "");
|
|
if (value === undefined) sendCancelled();
|
|
else sendJson({ type: "extension_ui_response", id: request.id, value });
|
|
return;
|
|
}
|
|
case "notify":
|
|
ctx.ui.notify(request.message, request.notifyType);
|
|
return;
|
|
case "setStatus":
|
|
ctx.ui.setStatus(request.statusKey, request.statusText);
|
|
return;
|
|
case "setWidget":
|
|
if (request.widgetLines === undefined || Array.isArray(request.widgetLines)) {
|
|
ctx.ui.setWidget(request.widgetKey, request.widgetLines, { placement: request.widgetPlacement });
|
|
}
|
|
return;
|
|
case "setTitle":
|
|
ctx.ui.setTitle(request.title);
|
|
return;
|
|
case "set_editor_text":
|
|
ctx.ui.setEditorText(request.text);
|
|
return;
|
|
}
|
|
} catch {
|
|
if (request.method === "select" || request.method === "confirm" || request.method === "input" || request.method === "editor") {
|
|
sendCancelled();
|
|
}
|
|
}
|
|
}
|
|
|
|
export async function runInteractiveAgentRpc(
|
|
defaultCwd: string,
|
|
agents: AgentConfig[],
|
|
agentName: string,
|
|
task: string,
|
|
cwd: string | undefined,
|
|
step: number | undefined,
|
|
signal: AbortSignal | undefined,
|
|
onUpdate: OnUpdateCallback | undefined,
|
|
makeDetails: (results: SingleResult[]) => SubagentDetails,
|
|
ctx: ExtensionContext,
|
|
): Promise<SingleResult> {
|
|
const agent = agents.find((entry) => entry.name === agentName);
|
|
if (!agent) {
|
|
const available = agents.map((entry) => `"${entry.name}"`).join(", ") || "none";
|
|
return {
|
|
agent: agentName,
|
|
agentSource: "unknown",
|
|
task,
|
|
exitCode: 1,
|
|
messages: [],
|
|
stderr: `Unknown agent: "${agentName}". Available agents: ${available}.`,
|
|
usage: createEmptyUsageStats(),
|
|
step,
|
|
};
|
|
}
|
|
|
|
const args: string[] = ["--mode", "rpc", "--no-session"];
|
|
if (agent.model) args.push("--model", agent.model);
|
|
|
|
let tmpPromptDir: string | null = null;
|
|
let tmpPromptPath: string | null = null;
|
|
if (agent.systemPrompt.trim()) {
|
|
const tmp = writePromptToTempFile(agent.name, agent.systemPrompt);
|
|
tmpPromptDir = tmp.dir;
|
|
tmpPromptPath = tmp.filePath;
|
|
args.push("--append-system-prompt", tmpPromptPath);
|
|
}
|
|
|
|
const extensionDir = path.dirname(fileURLToPath(import.meta.url));
|
|
const helperExtensionPath = path.join(extensionDir, "interactive-tools.ts");
|
|
args.push("--extension", helperExtensionPath);
|
|
|
|
const env = { ...process.env };
|
|
if (agent.tools && agent.tools.length > 0) {
|
|
env.PI_SUBAGENT_TOOLS = JSON.stringify(agent.tools);
|
|
}
|
|
|
|
const currentResult: SingleResult = {
|
|
agent: agentName,
|
|
agentSource: agent.source,
|
|
task,
|
|
exitCode: 0,
|
|
messages: [],
|
|
stderr: "",
|
|
usage: createEmptyUsageStats(),
|
|
model: agent.model,
|
|
step,
|
|
};
|
|
|
|
const emitUpdate = () => {
|
|
if (!onUpdate) return;
|
|
onUpdate({
|
|
content: [{ type: "text", text: getFinalOutput(currentResult.messages) || "(running...)" }],
|
|
details: makeDetails([currentResult]),
|
|
});
|
|
};
|
|
|
|
const proc = spawn("pi", args, { cwd: cwd ?? defaultCwd, shell: false, stdio: ["pipe", "pipe", "pipe"], env });
|
|
const pendingRequests = new Map<string, PendingRequest>();
|
|
let requestCount = 0;
|
|
let agentEnded = false;
|
|
let resolveAgentEnd: (() => void) | undefined;
|
|
const agentEndPromise = new Promise<void>((resolve) => {
|
|
resolveAgentEnd = resolve;
|
|
});
|
|
const processExitPromise = new Promise<number>((resolve) => {
|
|
proc.on("close", (code) => {
|
|
if (!agentEnded) resolveAgentEnd?.();
|
|
for (const [id, pending] of pendingRequests.entries()) {
|
|
clearTimeout(pending.timeout);
|
|
pending.reject(new Error(`RPC process exited before response: ${id}`));
|
|
}
|
|
pendingRequests.clear();
|
|
resolve(code ?? 0);
|
|
});
|
|
});
|
|
let wasAborted = false;
|
|
|
|
const sendJson = (obj: unknown) => {
|
|
if (proc.stdin.destroyed || proc.stdin.writableEnded) return;
|
|
proc.stdin.write(serializeJsonLine(obj));
|
|
};
|
|
|
|
const sendCommand = (command: Record<string, unknown>): Promise<RpcResponse> => {
|
|
const id = `subagent_${++requestCount}`;
|
|
return new Promise<RpcResponse>((resolve, reject) => {
|
|
const timeout = setTimeout(() => {
|
|
pendingRequests.delete(id);
|
|
reject(new Error(`RPC command timed out: ${command.type}`));
|
|
}, 30000);
|
|
pendingRequests.set(id, { resolve, reject, timeout });
|
|
sendJson({ ...command, id });
|
|
});
|
|
};
|
|
|
|
const detachStdout = attachJsonlLineReader(proc.stdout, (line) => {
|
|
const event = safeJsonParse(line);
|
|
if (!event) {
|
|
currentResult.stderr += `[rpc] Ignored non-JSON output: ${line}\n`;
|
|
return;
|
|
}
|
|
|
|
if (event.type === "response" && event.id && pendingRequests.has(event.id)) {
|
|
const pending = pendingRequests.get(event.id)!;
|
|
pendingRequests.delete(event.id);
|
|
clearTimeout(pending.timeout);
|
|
pending.resolve(event as RpcResponse);
|
|
return;
|
|
}
|
|
|
|
if (event.type === "extension_ui_request") {
|
|
void bridgeExtensionUiRequest(event as RpcExtensionUIRequest, ctx, signal, sendJson);
|
|
return;
|
|
}
|
|
|
|
if (event.type === "message_end" && event.message) {
|
|
const msg = event.message;
|
|
currentResult.messages.push(msg);
|
|
applyAssistantUsage(currentResult, msg);
|
|
emitUpdate();
|
|
return;
|
|
}
|
|
|
|
if (event.type === "agent_end") {
|
|
agentEnded = true;
|
|
resolveAgentEnd?.();
|
|
return;
|
|
}
|
|
|
|
if (event.type === "extension_error") {
|
|
const extensionPath = typeof event.extensionPath === "string" ? event.extensionPath : "unknown";
|
|
const error = typeof event.error === "string" ? event.error : "unknown extension error";
|
|
currentResult.stderr += `[extension:${extensionPath}] ${error}\n`;
|
|
}
|
|
});
|
|
|
|
proc.stderr.on("data", (data) => {
|
|
currentResult.stderr += data.toString();
|
|
});
|
|
|
|
proc.on("error", (error) => {
|
|
currentResult.stderr += `${error.message}\n`;
|
|
});
|
|
|
|
const abortHandler = () => {
|
|
wasAborted = true;
|
|
try {
|
|
sendJson({ type: "abort" });
|
|
} catch {}
|
|
terminateProcess(proc);
|
|
};
|
|
if (signal) {
|
|
if (signal.aborted) abortHandler();
|
|
else signal.addEventListener("abort", abortHandler, { once: true });
|
|
}
|
|
|
|
try {
|
|
const promptResponse = await sendCommand({ type: "prompt", message: `Task: ${task}` });
|
|
if (!promptResponse.success) {
|
|
currentResult.exitCode = 1;
|
|
currentResult.stderr += `Prompt failed: ${promptResponse.error || "unknown error"}`;
|
|
return currentResult;
|
|
}
|
|
|
|
await agentEndPromise;
|
|
if (currentResult.stopReason === "error") currentResult.exitCode = 1;
|
|
if (wasAborted) throw new Error("Subagent was aborted");
|
|
return currentResult;
|
|
} finally {
|
|
detachStdout();
|
|
for (const [id, pending] of pendingRequests.entries()) {
|
|
clearTimeout(pending.timeout);
|
|
pending.reject(new Error(`RPC command interrupted: ${id}`));
|
|
}
|
|
pendingRequests.clear();
|
|
if (signal) signal.removeEventListener("abort", abortHandler);
|
|
terminateProcess(proc);
|
|
const processCode = await processExitPromise;
|
|
if (!wasAborted && currentResult.exitCode === 0 && processCode !== 0 && !agentEnded) {
|
|
currentResult.exitCode = processCode;
|
|
}
|
|
if (tmpPromptPath) {
|
|
try {
|
|
fs.unlinkSync(tmpPromptPath);
|
|
} catch {}
|
|
}
|
|
if (tmpPromptDir) {
|
|
try {
|
|
fs.rmdirSync(tmpPromptDir);
|
|
} catch {}
|
|
}
|
|
}
|
|
}
|