refactor(subagent): split extension into modules
This commit is contained in:
617
index.ts
617
index.ts
@@ -1,591 +1,18 @@
|
|||||||
// Subagent Extension - Registers a tool for delegating work to prompt-defined
|
// Subagent Extension - Registers a tool for delegating work to prompt-defined
|
||||||
// subagents with constrained tool permissions.
|
// subagents with constrained tool permissions.
|
||||||
import { spawn } from "node:child_process";
|
import { randomUUID } from "node:crypto";
|
||||||
import { createHash, randomUUID } from "node:crypto";
|
|
||||||
import * as fs from "node:fs";
|
import * as fs from "node:fs";
|
||||||
import * as os from "node:os";
|
|
||||||
import * as path from "node:path";
|
import * as path from "node:path";
|
||||||
import { fileURLToPath } from "node:url";
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
import type { ExtensionAPI, Theme } from "@mariozechner/pi-coding-agent";
|
|
||||||
import {
|
|
||||||
parseFrontmatter,
|
|
||||||
withFileMutationQueue,
|
|
||||||
} from "@mariozechner/pi-coding-agent";
|
|
||||||
import { Markdown, Text, type MarkdownTheme } from "@mariozechner/pi-tui";
|
|
||||||
import { Type } from "typebox";
|
import { Type } from "typebox";
|
||||||
|
import { FINALIZE_TOOL_NAME, MAX_FINALIZE_RETRIES } from "./src/constants.ts";
|
||||||
interface PromptConfig {
|
import { formatPromptList, formatSubagentContent } from "./src/format.ts";
|
||||||
name: string;
|
import { discoverPrompts } from "./src/prompts.ts";
|
||||||
description: string;
|
import { renderSubagentCall, renderSubagentResult } from "./src/render.ts";
|
||||||
approvedTools: string[];
|
import { runAgent } from "./src/runner.ts";
|
||||||
deniedTools: string[];
|
import { getSubagentSessionPath } from "./src/session.ts";
|
||||||
systemPrompt: string;
|
import { FinalizeStatus, type SubagentResult } from "./src/types.ts";
|
||||||
filePath: string;
|
import { resolveTools, SubagentParams } from "./src/tools.ts";
|
||||||
}
|
|
||||||
|
|
||||||
interface SubagentToolActivity {
|
|
||||||
id: string;
|
|
||||||
toolName: string;
|
|
||||||
summary: string;
|
|
||||||
status: "running" | "done" | "error";
|
|
||||||
}
|
|
||||||
|
|
||||||
interface SubagentStatus {
|
|
||||||
state: "starting" | "thinking" | "running" | "done";
|
|
||||||
toolCallCount: number;
|
|
||||||
activeToolCalls: number;
|
|
||||||
recentToolCalls: SubagentToolActivity[];
|
|
||||||
}
|
|
||||||
|
|
||||||
enum FinalizeStatus {
|
|
||||||
SUCCESS = "SUCCESS",
|
|
||||||
ERROR = "ERROR",
|
|
||||||
}
|
|
||||||
|
|
||||||
interface SubagentFinalizePayload {
|
|
||||||
status: FinalizeStatus;
|
|
||||||
result?: string;
|
|
||||||
error?: string;
|
|
||||||
}
|
|
||||||
|
|
||||||
interface SubagentResult {
|
|
||||||
agent: string;
|
|
||||||
task: string;
|
|
||||||
tools: string[];
|
|
||||||
sessionId: string;
|
|
||||||
exitCode: number;
|
|
||||||
output: string;
|
|
||||||
stderr: string;
|
|
||||||
status: SubagentStatus;
|
|
||||||
finalized?: SubagentFinalizePayload;
|
|
||||||
error?: string;
|
|
||||||
}
|
|
||||||
|
|
||||||
const EXTENSION_ENTRY = fileURLToPath(import.meta.url);
|
|
||||||
const EXTENSION_DIR = path.dirname(EXTENSION_ENTRY);
|
|
||||||
const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");
|
|
||||||
const FINALIZE_TOOL_NAME = "subagent_finalize";
|
|
||||||
const MAX_FINALIZE_RETRIES = 2;
|
|
||||||
|
|
||||||
// Markdown Theme - Render finalized subagent text with Pi's terminal markdown component.
|
|
||||||
function getSubagentMarkdownTheme(theme: Theme): MarkdownTheme {
|
|
||||||
return {
|
|
||||||
heading: (text) => theme.fg("mdHeading", text),
|
|
||||||
link: (text) => theme.fg("mdLink", text),
|
|
||||||
linkUrl: (text) => theme.fg("mdLinkUrl", text),
|
|
||||||
code: (text) => theme.fg("mdCode", text),
|
|
||||||
codeBlock: (text) => theme.fg("mdCodeBlock", text),
|
|
||||||
codeBlockBorder: (text) => theme.fg("mdCodeBlockBorder", text),
|
|
||||||
quote: (text) => theme.fg("mdQuote", text),
|
|
||||||
quoteBorder: (text) => theme.fg("mdQuoteBorder", text),
|
|
||||||
hr: (text) => theme.fg("mdHr", text),
|
|
||||||
listBullet: (text) => theme.fg("mdListBullet", text),
|
|
||||||
bold: (text) => theme.bold(text),
|
|
||||||
italic: (text) => theme.italic(text),
|
|
||||||
strikethrough: (text) => theme.strikethrough(text),
|
|
||||||
underline: (text) => theme.underline(text),
|
|
||||||
codeBlockIndent: " ",
|
|
||||||
};
|
|
||||||
}
|
|
||||||
|
|
||||||
// Format Tool Content - Some clients hide structured details from the model.
|
|
||||||
function formatSubagentContent(
|
|
||||||
status: FinalizeStatus,
|
|
||||||
sessionId: string,
|
|
||||||
result?: string,
|
|
||||||
error?: string,
|
|
||||||
): string {
|
|
||||||
const header = [`**Status:** ${status}`, `**Session ID:** \`${sessionId}\``];
|
|
||||||
if (error?.trim()) header.push(`**Error:** ${error.trim()}`);
|
|
||||||
|
|
||||||
const body = result?.trimEnd();
|
|
||||||
return body
|
|
||||||
? `${header.join(" \n")}\n\n---\n\n${body}`
|
|
||||||
: header.join(" \n");
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse Tool List - Frontmatter may use YAML arrays or comma-delimited strings.
|
|
||||||
function parseToolList(value: unknown): string[] {
|
|
||||||
if (!value) return [];
|
|
||||||
if (Array.isArray(value)) {
|
|
||||||
return value
|
|
||||||
.map(String)
|
|
||||||
.map((tool) => tool.trim())
|
|
||||||
.filter(Boolean);
|
|
||||||
}
|
|
||||||
if (typeof value === "string") {
|
|
||||||
return value
|
|
||||||
.split(",")
|
|
||||||
.map((tool) => tool.trim())
|
|
||||||
.filter(Boolean);
|
|
||||||
}
|
|
||||||
return [];
|
|
||||||
}
|
|
||||||
|
|
||||||
// Discover Prompts - Load prompt markdown files from this extension's prompts dir.
|
|
||||||
function discoverPrompts(): PromptConfig[] {
|
|
||||||
if (!fs.existsSync(PROMPTS_DIR)) return [];
|
|
||||||
|
|
||||||
const prompts: PromptConfig[] = [];
|
|
||||||
const entries = fs.readdirSync(PROMPTS_DIR, { withFileTypes: true });
|
|
||||||
for (const entry of entries) {
|
|
||||||
if (!entry.name.endsWith(".md")) continue;
|
|
||||||
if (!entry.isFile() && !entry.isSymbolicLink()) continue;
|
|
||||||
|
|
||||||
const filePath = path.join(PROMPTS_DIR, entry.name);
|
|
||||||
const content = fs.readFileSync(filePath, "utf-8");
|
|
||||||
const { frontmatter, body } =
|
|
||||||
parseFrontmatter<Record<string, unknown>>(content);
|
|
||||||
|
|
||||||
if (
|
|
||||||
typeof frontmatter.name !== "string" ||
|
|
||||||
typeof frontmatter.description !== "string"
|
|
||||||
) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
prompts.push({
|
|
||||||
name: frontmatter.name,
|
|
||||||
description: frontmatter.description,
|
|
||||||
approvedTools: parseToolList(
|
|
||||||
frontmatter.approved_tools ?? frontmatter.allowed_tools,
|
|
||||||
),
|
|
||||||
deniedTools: parseToolList(frontmatter.denied_tools),
|
|
||||||
systemPrompt: body.trim(),
|
|
||||||
filePath,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
return prompts.sort((a, b) => a.name.localeCompare(b.name));
|
|
||||||
}
|
|
||||||
|
|
||||||
// Resolve Tools - Use exactly one permission mode. approved_tools/allowed_tools
|
|
||||||
// is a whitelist; denied_tools is a blacklist over the currently active tools.
|
|
||||||
function resolveTools(agent: PromptConfig, activeTools: string[]): string[] {
|
|
||||||
const withoutDelegationTools = (tools: string[]) =>
|
|
||||||
tools.filter((tool) => tool !== "subagent" && tool !== FINALIZE_TOOL_NAME);
|
|
||||||
|
|
||||||
const resolved =
|
|
||||||
agent.approvedTools.length > 0
|
|
||||||
? withoutDelegationTools([...new Set(agent.approvedTools)])
|
|
||||||
: withoutDelegationTools(
|
|
||||||
[...new Set(activeTools)].filter(
|
|
||||||
(tool) => !new Set(agent.deniedTools).has(tool),
|
|
||||||
),
|
|
||||||
);
|
|
||||||
|
|
||||||
return [...new Set([...resolved, FINALIZE_TOOL_NAME])];
|
|
||||||
}
|
|
||||||
|
|
||||||
// Build Finalize Prompt - Child agents must terminate by calling this tool.
|
|
||||||
function buildSubagentPrompt(agent: PromptConfig): string {
|
|
||||||
const finalizePrompt = [
|
|
||||||
"You are running as a subagent.",
|
|
||||||
`When the task is complete, call ${FINALIZE_TOOL_NAME} as your final action.`,
|
|
||||||
"Do not provide the final answer as normal assistant text.",
|
|
||||||
`${FINALIZE_TOOL_NAME} requires status ${FinalizeStatus.SUCCESS} with result, or status ${FinalizeStatus.ERROR} with error and optional result.`,
|
|
||||||
].join("\n");
|
|
||||||
|
|
||||||
return [agent.systemPrompt, finalizePrompt].filter(Boolean).join("\n\n");
|
|
||||||
}
|
|
||||||
|
|
||||||
// Session Path - Persist child sessions as <agent>_<uuid>.jsonl under a cwd hash.
|
|
||||||
function getSubagentSessionPath(
|
|
||||||
cwd: string,
|
|
||||||
agentName: string,
|
|
||||||
sessionId: string,
|
|
||||||
): string {
|
|
||||||
const cwdHash = createHash("sha256").update(cwd).digest("hex").slice(0, 16);
|
|
||||||
const safeAgent = agentName.replace(/[^\w.-]+/g, "_");
|
|
||||||
const safeSessionId = sessionId.replace(/[^\w.-]+/g, "_");
|
|
||||||
return path.join(
|
|
||||||
os.homedir(),
|
|
||||||
".pi",
|
|
||||||
"subagent-sessions",
|
|
||||||
cwdHash,
|
|
||||||
`${safeAgent}_${safeSessionId}.jsonl`,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Validate Finalize Payload - Keep the parent contract strict and small.
|
|
||||||
function validateFinalizePayload(
|
|
||||||
value: unknown,
|
|
||||||
): SubagentFinalizePayload | null {
|
|
||||||
if (!value || typeof value !== "object") return null;
|
|
||||||
const payload = value as Record<string, unknown>;
|
|
||||||
if (payload.status === FinalizeStatus.SUCCESS) {
|
|
||||||
return typeof payload.result === "string" && payload.result.trim()
|
|
||||||
? { status: FinalizeStatus.SUCCESS, result: payload.result }
|
|
||||||
: null;
|
|
||||||
}
|
|
||||||
if (payload.status === FinalizeStatus.ERROR) {
|
|
||||||
const result =
|
|
||||||
typeof payload.result === "string" ? payload.result : undefined;
|
|
||||||
return typeof payload.error === "string" && payload.error.trim()
|
|
||||||
? { status: FinalizeStatus.ERROR, error: payload.error, result }
|
|
||||||
: null;
|
|
||||||
}
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Write Prompt - pi accepts appended system prompts via file path.
|
|
||||||
async function writePromptToTempFile(
|
|
||||||
agentName: string,
|
|
||||||
prompt: string,
|
|
||||||
): Promise<{ dir: string; filePath: string }> {
|
|
||||||
const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
|
|
||||||
const safeName = agentName.replace(/[^\w.-]+/g, "_");
|
|
||||||
const filePath = path.join(dir, `${safeName}.md`);
|
|
||||||
await withFileMutationQueue(filePath, async () => {
|
|
||||||
await fs.promises.writeFile(filePath, prompt, {
|
|
||||||
encoding: "utf-8",
|
|
||||||
mode: 0o600,
|
|
||||||
});
|
|
||||||
});
|
|
||||||
return { dir, filePath };
|
|
||||||
}
|
|
||||||
|
|
||||||
// Pi Invocation - Prefer the current pi entrypoint when running inside pi.
|
|
||||||
function getPiInvocation(args: string[]): { command: string; args: string[] } {
|
|
||||||
const currentScript = process.argv[1];
|
|
||||||
const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
|
|
||||||
if (currentScript && !isBunVirtualScript && fs.existsSync(currentScript)) {
|
|
||||||
return { command: process.execPath, args: [currentScript, ...args] };
|
|
||||||
}
|
|
||||||
|
|
||||||
const execName = path.basename(process.execPath).toLowerCase();
|
|
||||||
const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
|
|
||||||
if (!isGenericRuntime) return { command: process.execPath, args };
|
|
||||||
return { command: "pi", args };
|
|
||||||
}
|
|
||||||
|
|
||||||
// Truncate One Line - Keep live status compact and stable in the TUI.
|
|
||||||
function truncateOneLine(text: string, maxLength: number): string {
|
|
||||||
const firstLine = text.split(/\r?\n/, 1)[0]?.trim() ?? "";
|
|
||||||
return firstLine.length > maxLength
|
|
||||||
? `${firstLine.slice(0, Math.max(0, maxLength - 3))}...`
|
|
||||||
: firstLine;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Format Tool Args - Keep status updates compact enough for the tool UI.
|
|
||||||
function formatToolArgs(args: unknown): string {
|
|
||||||
if (!args || typeof args !== "object") return "";
|
|
||||||
|
|
||||||
const record = args as Record<string, unknown>;
|
|
||||||
const priorityKeys = ["path", "command", "query", "pattern", "file", "agent"];
|
|
||||||
const key = priorityKeys.find((candidate) => record[candidate] !== undefined);
|
|
||||||
const value = key ? record[key] : undefined;
|
|
||||||
if (value === undefined) return "";
|
|
||||||
|
|
||||||
const text = typeof value === "string" ? value : JSON.stringify(value);
|
|
||||||
const compact = text.replace(/\s+/g, " ").trim();
|
|
||||||
return compact.length > 80 ? `${compact.slice(0, 77)}...` : compact;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Format Status Text - Plain fallback for non-TUI consumers and JSON logs.
|
|
||||||
function formatStatusText(result: SubagentResult): string {
|
|
||||||
const statusText =
|
|
||||||
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
|
||||||
const lines = [
|
|
||||||
`Subagent: ${result.agent}`,
|
|
||||||
`Task: ${result.task}`,
|
|
||||||
`Status: ${statusText}`,
|
|
||||||
`Tool Calls: ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
|
|
||||||
];
|
|
||||||
|
|
||||||
if (result.status.recentToolCalls.length > 0) {
|
|
||||||
lines.push("", "Tool Calls - Last 3:");
|
|
||||||
for (const call of result.status.recentToolCalls.slice(0, 3)) {
|
|
||||||
const mark =
|
|
||||||
call.status === "done" ? "✓" : call.status === "error" ? "✗" : "•";
|
|
||||||
const suffix = call.summary ? ` ${call.summary}` : "";
|
|
||||||
lines.push(`- ${mark} ${call.toolName}${suffix}`);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return lines.join("\n");
|
|
||||||
}
|
|
||||||
|
|
||||||
// Render Subagent Summary - Compact collapsed view; Ctrl+O expands details.
|
|
||||||
function renderSubagentSummary(result: SubagentResult, theme: Theme): Text {
|
|
||||||
const statusText =
|
|
||||||
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
|
||||||
const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
|
|
||||||
const latest = result.status.recentToolCalls[0];
|
|
||||||
const latestText = latest
|
|
||||||
? ` · ${latest.toolName}${latest.summary ? ` ${latest.summary}` : ""}`
|
|
||||||
: "";
|
|
||||||
const text =
|
|
||||||
theme.fg(statusColor, statusText) +
|
|
||||||
theme.fg(
|
|
||||||
"muted",
|
|
||||||
` · ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active${latestText}`,
|
|
||||||
);
|
|
||||||
return new Text(text, 0, 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Render Subagent Status - Expanded TUI renderer with theme-aware colors.
|
|
||||||
function renderSubagentStatus(result: SubagentResult, theme: Theme): Text {
|
|
||||||
const statusText =
|
|
||||||
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
|
||||||
const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
|
|
||||||
|
|
||||||
const lines = [
|
|
||||||
`${theme.fg("dim", "Subagent:")} ${theme.fg("accent", result.agent)}`,
|
|
||||||
`${theme.fg("dim", "Task:")} ${theme.fg("muted", result.task)}`,
|
|
||||||
`${theme.fg("dim", "Status:")} ${theme.fg(statusColor, statusText)}`,
|
|
||||||
`${theme.fg("dim", "Tool Calls:")} ${theme.fg(
|
|
||||||
"muted",
|
|
||||||
`${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
|
|
||||||
)}`,
|
|
||||||
];
|
|
||||||
|
|
||||||
if (result.status.recentToolCalls.length > 0) {
|
|
||||||
lines.push("", theme.fg("dim", "Tool Calls - Last 3:"));
|
|
||||||
for (const call of result.status.recentToolCalls.slice(0, 3)) {
|
|
||||||
const mark =
|
|
||||||
call.status === "done"
|
|
||||||
? theme.fg("success", "✓")
|
|
||||||
: call.status === "error"
|
|
||||||
? theme.fg("error", "✗")
|
|
||||||
: theme.fg("warning", "•");
|
|
||||||
const tool = theme.fg("toolTitle", call.toolName.padEnd(6));
|
|
||||||
const summary = call.summary ? ` ${theme.fg("muted", call.summary)}` : "";
|
|
||||||
lines.push(`- ${mark} ${tool}${summary}`);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return new Text(lines.join("\n"), 0, 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Run Agent - Spawn an isolated pi process in JSON mode and collect final text.
|
|
||||||
async function runAgent(
|
|
||||||
cwd: string,
|
|
||||||
agent: PromptConfig,
|
|
||||||
task: string,
|
|
||||||
tools: string[],
|
|
||||||
sessionId: string,
|
|
||||||
sessionPath: string,
|
|
||||||
signal?: AbortSignal,
|
|
||||||
onUpdate?: (partial: {
|
|
||||||
content: { type: "text"; text: string }[];
|
|
||||||
details: SubagentResult;
|
|
||||||
}) => void,
|
|
||||||
): Promise<SubagentResult> {
|
|
||||||
let tmpDir: string | null = null;
|
|
||||||
let tmpPromptPath: string | null = null;
|
|
||||||
|
|
||||||
const result: SubagentResult = {
|
|
||||||
agent: agent.name,
|
|
||||||
task,
|
|
||||||
tools,
|
|
||||||
sessionId,
|
|
||||||
exitCode: 0,
|
|
||||||
output: "",
|
|
||||||
stderr: "",
|
|
||||||
status: {
|
|
||||||
state: "starting",
|
|
||||||
toolCallCount: 0,
|
|
||||||
activeToolCalls: 0,
|
|
||||||
recentToolCalls: [],
|
|
||||||
},
|
|
||||||
};
|
|
||||||
|
|
||||||
const activeToolIds = new Set<string>();
|
|
||||||
|
|
||||||
const emitUpdate = () => {
|
|
||||||
onUpdate?.({
|
|
||||||
content: [{ type: "text", text: formatStatusText(result) }],
|
|
||||||
details: {
|
|
||||||
...result,
|
|
||||||
status: {
|
|
||||||
...result.status,
|
|
||||||
recentToolCalls: [...result.status.recentToolCalls],
|
|
||||||
},
|
|
||||||
},
|
|
||||||
});
|
|
||||||
};
|
|
||||||
|
|
||||||
const rememberToolCall = (activity: SubagentToolActivity) => {
|
|
||||||
const existing = result.status.recentToolCalls.findIndex(
|
|
||||||
(call) => call.id === activity.id,
|
|
||||||
);
|
|
||||||
if (existing >= 0) result.status.recentToolCalls.splice(existing, 1);
|
|
||||||
result.status.recentToolCalls.unshift(activity);
|
|
||||||
result.status.recentToolCalls = result.status.recentToolCalls.slice(0, 3);
|
|
||||||
};
|
|
||||||
|
|
||||||
try {
|
|
||||||
const args = [
|
|
||||||
"--mode",
|
|
||||||
"json",
|
|
||||||
"-p",
|
|
||||||
"--session",
|
|
||||||
sessionPath,
|
|
||||||
"--extension",
|
|
||||||
EXTENSION_ENTRY,
|
|
||||||
];
|
|
||||||
args.push("--tools", tools.join(","));
|
|
||||||
|
|
||||||
const prompt = buildSubagentPrompt(agent);
|
|
||||||
if (prompt) {
|
|
||||||
const tmp = await writePromptToTempFile(agent.name, prompt);
|
|
||||||
tmpDir = tmp.dir;
|
|
||||||
tmpPromptPath = tmp.filePath;
|
|
||||||
args.push("--append-system-prompt", tmpPromptPath);
|
|
||||||
}
|
|
||||||
|
|
||||||
args.push(task);
|
|
||||||
emitUpdate();
|
|
||||||
|
|
||||||
const exitCode = await new Promise<number>((resolve) => {
|
|
||||||
const invocation = getPiInvocation(args);
|
|
||||||
const proc = spawn(invocation.command, invocation.args, {
|
|
||||||
cwd,
|
|
||||||
env: { ...process.env, PI_SUBAGENT_CHILD: "1" },
|
|
||||||
shell: false,
|
|
||||||
stdio: ["ignore", "pipe", "pipe"],
|
|
||||||
});
|
|
||||||
|
|
||||||
let buffer = "";
|
|
||||||
let aborted = false;
|
|
||||||
|
|
||||||
const processLine = (line: string) => {
|
|
||||||
if (!line.trim()) return;
|
|
||||||
try {
|
|
||||||
const event = JSON.parse(line);
|
|
||||||
if (
|
|
||||||
event.type === "message_start" &&
|
|
||||||
event.message?.role === "assistant"
|
|
||||||
) {
|
|
||||||
result.status.state = "thinking";
|
|
||||||
emitUpdate();
|
|
||||||
}
|
|
||||||
if (event.type === "tool_execution_start") {
|
|
||||||
const id = String(
|
|
||||||
event.toolCallId ?? result.status.toolCallCount + 1,
|
|
||||||
);
|
|
||||||
const toolName = String(event.toolName ?? "tool");
|
|
||||||
if (toolName === FINALIZE_TOOL_NAME) {
|
|
||||||
result.finalized =
|
|
||||||
validateFinalizePayload(event.args) ?? undefined;
|
|
||||||
}
|
|
||||||
activeToolIds.add(id);
|
|
||||||
result.status.state = "running";
|
|
||||||
result.status.toolCallCount += 1;
|
|
||||||
result.status.activeToolCalls = activeToolIds.size;
|
|
||||||
rememberToolCall({
|
|
||||||
id,
|
|
||||||
toolName,
|
|
||||||
summary: formatToolArgs(event.args),
|
|
||||||
status: "running",
|
|
||||||
});
|
|
||||||
emitUpdate();
|
|
||||||
}
|
|
||||||
if (event.type === "tool_execution_end") {
|
|
||||||
const id = String(event.toolCallId ?? "");
|
|
||||||
if (id) activeToolIds.delete(id);
|
|
||||||
result.status.activeToolCalls = activeToolIds.size;
|
|
||||||
const previous = result.status.recentToolCalls.find(
|
|
||||||
(call) => call.id === id,
|
|
||||||
);
|
|
||||||
rememberToolCall({
|
|
||||||
id,
|
|
||||||
toolName: String(event.toolName ?? previous?.toolName ?? "tool"),
|
|
||||||
summary: previous?.summary ?? "",
|
|
||||||
status: event.isError ? "error" : "done",
|
|
||||||
});
|
|
||||||
emitUpdate();
|
|
||||||
}
|
|
||||||
if (
|
|
||||||
event.type === "message_end" &&
|
|
||||||
event.message?.role === "assistant"
|
|
||||||
) {
|
|
||||||
for (const part of event.message.content ?? []) {
|
|
||||||
if (part.type === "text") result.output = part.text;
|
|
||||||
}
|
|
||||||
if (event.message.stopReason === "error") {
|
|
||||||
result.error =
|
|
||||||
event.message.errorMessage ?? "Subagent stopped with an error.";
|
|
||||||
}
|
|
||||||
emitUpdate();
|
|
||||||
}
|
|
||||||
} catch {
|
|
||||||
// Ignore non-JSON output.
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
proc.stdout.on("data", (data) => {
|
|
||||||
buffer += data.toString();
|
|
||||||
const lines = buffer.split("\n");
|
|
||||||
buffer = lines.pop() ?? "";
|
|
||||||
for (const line of lines) processLine(line);
|
|
||||||
});
|
|
||||||
|
|
||||||
proc.stderr.on("data", (data) => {
|
|
||||||
result.stderr += data.toString();
|
|
||||||
});
|
|
||||||
|
|
||||||
proc.on("close", (code) => {
|
|
||||||
if (buffer.trim()) processLine(buffer);
|
|
||||||
resolve(aborted ? 130 : (code ?? 0));
|
|
||||||
});
|
|
||||||
|
|
||||||
proc.on("error", (error) => {
|
|
||||||
result.error = error.message;
|
|
||||||
resolve(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
if (signal) {
|
|
||||||
const killProc = () => {
|
|
||||||
aborted = true;
|
|
||||||
proc.kill("SIGTERM");
|
|
||||||
setTimeout(() => {
|
|
||||||
if (!proc.killed) proc.kill("SIGKILL");
|
|
||||||
}, 5000);
|
|
||||||
};
|
|
||||||
if (signal.aborted) killProc();
|
|
||||||
else signal.addEventListener("abort", killProc, { once: true });
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
result.exitCode = exitCode;
|
|
||||||
result.status.state = "done";
|
|
||||||
result.status.activeToolCalls = 0;
|
|
||||||
if (exitCode === 130) result.error = "Subagent was aborted.";
|
|
||||||
return result;
|
|
||||||
} finally {
|
|
||||||
if (tmpPromptPath) await fs.promises.rm(tmpPromptPath, { force: true });
|
|
||||||
if (tmpDir) await fs.promises.rm(tmpDir, { force: true, recursive: true });
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Format Prompt List - Keep available prompt hints compact for the LLM.
|
|
||||||
function formatPromptList(prompts: PromptConfig[]): string {
|
|
||||||
if (prompts.length === 0) return "none";
|
|
||||||
return prompts
|
|
||||||
.map((prompt) => `${prompt.name}: ${prompt.description}`)
|
|
||||||
.join("; ");
|
|
||||||
}
|
|
||||||
|
|
||||||
const SubagentParams = Type.Object({
|
|
||||||
agent: Type.String({
|
|
||||||
description: "Name of the prompt-defined subagent to invoke",
|
|
||||||
}),
|
|
||||||
task: Type.String({ description: "Task to delegate to the subagent" }),
|
|
||||||
sessionId: Type.Optional(
|
|
||||||
Type.String({
|
|
||||||
description:
|
|
||||||
"Optional sticky subagent session id. Reuse to continue a previous subagent context.",
|
|
||||||
}),
|
|
||||||
),
|
|
||||||
cwd: Type.Optional(
|
|
||||||
Type.String({
|
|
||||||
description:
|
|
||||||
"Working directory for the subagent process. Defaults to current cwd.",
|
|
||||||
}),
|
|
||||||
),
|
|
||||||
});
|
|
||||||
|
|
||||||
export default function (pi: ExtensionAPI) {
|
export default function (pi: ExtensionAPI) {
|
||||||
if (process.env.PI_SUBAGENT_CHILD === "1") {
|
if (process.env.PI_SUBAGENT_CHILD === "1") {
|
||||||
@@ -783,29 +210,11 @@ export default function (pi: ExtensionAPI) {
|
|||||||
},
|
},
|
||||||
|
|
||||||
renderCall(args, theme) {
|
renderCall(args, theme) {
|
||||||
const task = truncateOneLine(args.task, 90);
|
return renderSubagentCall(args, theme);
|
||||||
const text =
|
|
||||||
theme.fg("toolTitle", theme.bold("subagent ")) +
|
|
||||||
theme.fg("accent", args.agent) +
|
|
||||||
(task ? ` ${theme.fg("dim", task)}` : "");
|
|
||||||
return new Text(text, 0, 0);
|
|
||||||
},
|
},
|
||||||
|
|
||||||
renderResult(result, { expanded, isPartial }, theme) {
|
renderResult(result, options, theme) {
|
||||||
const details = result.details as Partial<SubagentResult> | undefined;
|
return renderSubagentResult(result, options, theme);
|
||||||
if (isPartial && details?.status && details.agent && details.task) {
|
|
||||||
return expanded
|
|
||||||
? renderSubagentStatus(details as SubagentResult, theme)
|
|
||||||
: renderSubagentSummary(details as SubagentResult, theme);
|
|
||||||
}
|
|
||||||
|
|
||||||
const text = result.content[0];
|
|
||||||
return new Markdown(
|
|
||||||
text?.type === "text" ? text.text : "",
|
|
||||||
0,
|
|
||||||
0,
|
|
||||||
getSubagentMarkdownTheme(theme),
|
|
||||||
);
|
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|||||||
10
src/constants.ts
Normal file
10
src/constants.ts
Normal file
@@ -0,0 +1,10 @@
|
|||||||
|
import * as path from "node:path";
|
||||||
|
import { fileURLToPath } from "node:url";
|
||||||
|
|
||||||
|
const SRC_DIR = path.dirname(fileURLToPath(import.meta.url));
|
||||||
|
|
||||||
|
export const EXTENSION_DIR = path.dirname(SRC_DIR);
|
||||||
|
export const EXTENSION_ENTRY = path.join(EXTENSION_DIR, "index.ts");
|
||||||
|
export const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");
|
||||||
|
export const FINALIZE_TOOL_NAME = "subagent_finalize";
|
||||||
|
export const MAX_FINALIZE_RETRIES = 2;
|
||||||
73
src/format.ts
Normal file
73
src/format.ts
Normal file
@@ -0,0 +1,73 @@
|
|||||||
|
import type { PromptConfig, SubagentResult } from "./types.ts";
|
||||||
|
import { FinalizeStatus } from "./types.ts";
|
||||||
|
|
||||||
|
// Format Tool Content - Some clients hide structured details from the model.
|
||||||
|
export function formatSubagentContent(
|
||||||
|
status: FinalizeStatus,
|
||||||
|
sessionId: string,
|
||||||
|
result?: string,
|
||||||
|
error?: string,
|
||||||
|
): string {
|
||||||
|
const header = [`**Status:** ${status}`, `**Session ID:** \`${sessionId}\``];
|
||||||
|
if (error?.trim()) header.push(`**Error:** ${error.trim()}`);
|
||||||
|
|
||||||
|
const body = result?.trimEnd();
|
||||||
|
return body
|
||||||
|
? `${header.join(" \n")}\n\n---\n\n${body}`
|
||||||
|
: header.join(" \n");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Truncate One Line - Keep live status compact and stable in the TUI.
|
||||||
|
export function truncateOneLine(text: string, maxLength: number): string {
|
||||||
|
const firstLine = text.split(/\r?\n/, 1)[0]?.trim() ?? "";
|
||||||
|
return firstLine.length > maxLength
|
||||||
|
? `${firstLine.slice(0, Math.max(0, maxLength - 3))}...`
|
||||||
|
: firstLine;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Format Tool Args - Keep status updates compact enough for the tool UI.
|
||||||
|
export function formatToolArgs(args: unknown): string {
|
||||||
|
if (!args || typeof args !== "object") return "";
|
||||||
|
|
||||||
|
const record = args as Record<string, unknown>;
|
||||||
|
const priorityKeys = ["path", "command", "query", "pattern", "file", "agent"];
|
||||||
|
const key = priorityKeys.find((candidate) => record[candidate] !== undefined);
|
||||||
|
const value = key ? record[key] : undefined;
|
||||||
|
if (value === undefined) return "";
|
||||||
|
|
||||||
|
const text = typeof value === "string" ? value : JSON.stringify(value);
|
||||||
|
const compact = text.replace(/\s+/g, " ").trim();
|
||||||
|
return compact.length > 80 ? `${compact.slice(0, 77)}...` : compact;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Format Status Text - Plain fallback for non-TUI consumers and JSON logs.
|
||||||
|
export function formatStatusText(result: SubagentResult): string {
|
||||||
|
const statusText =
|
||||||
|
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
||||||
|
const lines = [
|
||||||
|
`Subagent: ${result.agent}`,
|
||||||
|
`Task: ${result.task}`,
|
||||||
|
`Status: ${statusText}`,
|
||||||
|
`Tool Calls: ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
|
||||||
|
];
|
||||||
|
|
||||||
|
if (result.status.recentToolCalls.length > 0) {
|
||||||
|
lines.push("", "Tool Calls - Last 3:");
|
||||||
|
for (const call of result.status.recentToolCalls.slice(0, 3)) {
|
||||||
|
const mark =
|
||||||
|
call.status === "done" ? "✓" : call.status === "error" ? "✗" : "•";
|
||||||
|
const suffix = call.summary ? ` ${call.summary}` : "";
|
||||||
|
lines.push(`- ${mark} ${call.toolName}${suffix}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return lines.join("\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Format Prompt List - Keep available prompt hints compact for the LLM.
|
||||||
|
export function formatPromptList(prompts: PromptConfig[]): string {
|
||||||
|
if (prompts.length === 0) return "none";
|
||||||
|
return prompts
|
||||||
|
.map((prompt) => `${prompt.name}: ${prompt.description}`)
|
||||||
|
.join("; ");
|
||||||
|
}
|
||||||
73
src/prompts.ts
Normal file
73
src/prompts.ts
Normal file
@@ -0,0 +1,73 @@
|
|||||||
|
import * as fs from "node:fs";
|
||||||
|
import * as path from "node:path";
|
||||||
|
import { parseFrontmatter } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { FINALIZE_TOOL_NAME, PROMPTS_DIR } from "./constants.ts";
|
||||||
|
import type { PromptConfig } from "./types.ts";
|
||||||
|
import { FinalizeStatus } from "./types.ts";
|
||||||
|
|
||||||
|
// Parse Tool List - Frontmatter may use YAML arrays or comma-delimited strings.
|
||||||
|
export function parseToolList(value: unknown): string[] {
|
||||||
|
if (!value) return [];
|
||||||
|
if (Array.isArray(value)) {
|
||||||
|
return value
|
||||||
|
.map(String)
|
||||||
|
.map((tool) => tool.trim())
|
||||||
|
.filter(Boolean);
|
||||||
|
}
|
||||||
|
if (typeof value === "string") {
|
||||||
|
return value
|
||||||
|
.split(",")
|
||||||
|
.map((tool) => tool.trim())
|
||||||
|
.filter(Boolean);
|
||||||
|
}
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
|
||||||
|
// Discover Prompts - Load prompt markdown files from this extension's prompts dir.
|
||||||
|
export function discoverPrompts(): PromptConfig[] {
|
||||||
|
if (!fs.existsSync(PROMPTS_DIR)) return [];
|
||||||
|
|
||||||
|
const prompts: PromptConfig[] = [];
|
||||||
|
const entries = fs.readdirSync(PROMPTS_DIR, { withFileTypes: true });
|
||||||
|
for (const entry of entries) {
|
||||||
|
if (!entry.name.endsWith(".md")) continue;
|
||||||
|
if (!entry.isFile() && !entry.isSymbolicLink()) continue;
|
||||||
|
|
||||||
|
const filePath = path.join(PROMPTS_DIR, entry.name);
|
||||||
|
const content = fs.readFileSync(filePath, "utf-8");
|
||||||
|
const { frontmatter, body } =
|
||||||
|
parseFrontmatter<Record<string, unknown>>(content);
|
||||||
|
|
||||||
|
if (
|
||||||
|
typeof frontmatter.name !== "string" ||
|
||||||
|
typeof frontmatter.description !== "string"
|
||||||
|
) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
prompts.push({
|
||||||
|
name: frontmatter.name,
|
||||||
|
description: frontmatter.description,
|
||||||
|
approvedTools: parseToolList(
|
||||||
|
frontmatter.approved_tools ?? frontmatter.allowed_tools,
|
||||||
|
),
|
||||||
|
deniedTools: parseToolList(frontmatter.denied_tools),
|
||||||
|
systemPrompt: body.trim(),
|
||||||
|
filePath,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
return prompts.sort((a, b) => a.name.localeCompare(b.name));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build Finalize Prompt - Child agents must terminate by calling this tool.
|
||||||
|
export function buildSubagentPrompt(agent: PromptConfig): string {
|
||||||
|
const finalizePrompt = [
|
||||||
|
"You are running as a subagent.",
|
||||||
|
`When the task is complete, call ${FINALIZE_TOOL_NAME} as your final action.`,
|
||||||
|
"Do not provide the final answer as normal assistant text.",
|
||||||
|
`${FINALIZE_TOOL_NAME} requires status ${FinalizeStatus.SUCCESS} with result, or status ${FinalizeStatus.ERROR} with error and optional result.`,
|
||||||
|
].join("\n");
|
||||||
|
|
||||||
|
return [agent.systemPrompt, finalizePrompt].filter(Boolean).join("\n\n");
|
||||||
|
}
|
||||||
121
src/render.ts
Normal file
121
src/render.ts
Normal file
@@ -0,0 +1,121 @@
|
|||||||
|
import type { Theme } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { Markdown, Text, type MarkdownTheme } from "@mariozechner/pi-tui";
|
||||||
|
import { formatSubagentContent, truncateOneLine } from "./format.ts";
|
||||||
|
import type { SubagentResult } from "./types.ts";
|
||||||
|
import { FinalizeStatus } from "./types.ts";
|
||||||
|
|
||||||
|
// Markdown Theme - Render finalized subagent text with Pi's terminal markdown component.
|
||||||
|
export function getSubagentMarkdownTheme(theme: Theme): MarkdownTheme {
|
||||||
|
return {
|
||||||
|
heading: (text) => theme.fg("mdHeading", text),
|
||||||
|
link: (text) => theme.fg("mdLink", text),
|
||||||
|
linkUrl: (text) => theme.fg("mdLinkUrl", text),
|
||||||
|
code: (text) => theme.fg("mdCode", text),
|
||||||
|
codeBlock: (text) => theme.fg("mdCodeBlock", text),
|
||||||
|
codeBlockBorder: (text) => theme.fg("mdCodeBlockBorder", text),
|
||||||
|
quote: (text) => theme.fg("mdQuote", text),
|
||||||
|
quoteBorder: (text) => theme.fg("mdQuoteBorder", text),
|
||||||
|
hr: (text) => theme.fg("mdHr", text),
|
||||||
|
listBullet: (text) => theme.fg("mdListBullet", text),
|
||||||
|
bold: (text) => theme.bold(text),
|
||||||
|
italic: (text) => theme.italic(text),
|
||||||
|
strikethrough: (text) => theme.strikethrough(text),
|
||||||
|
underline: (text) => theme.underline(text),
|
||||||
|
codeBlockIndent: " ",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Render Subagent Summary - Compact collapsed view; Ctrl+O expands details.
|
||||||
|
export function renderSubagentSummary(
|
||||||
|
result: SubagentResult,
|
||||||
|
theme: Theme,
|
||||||
|
): Text {
|
||||||
|
const statusText =
|
||||||
|
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
||||||
|
const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
|
||||||
|
const latest = result.status.recentToolCalls[0];
|
||||||
|
const latestText = latest
|
||||||
|
? ` · ${latest.toolName}${latest.summary ? ` ${latest.summary}` : ""}`
|
||||||
|
: "";
|
||||||
|
const text =
|
||||||
|
theme.fg(statusColor, statusText) +
|
||||||
|
theme.fg(
|
||||||
|
"muted",
|
||||||
|
` · ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active${latestText}`,
|
||||||
|
);
|
||||||
|
return new Text(text, 0, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Render Subagent Status - Expanded TUI renderer with theme-aware colors.
|
||||||
|
export function renderSubagentStatus(
|
||||||
|
result: SubagentResult,
|
||||||
|
theme: Theme,
|
||||||
|
): Text {
|
||||||
|
const statusText =
|
||||||
|
result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
|
||||||
|
const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
|
||||||
|
|
||||||
|
const lines = [
|
||||||
|
`${theme.fg("dim", "Subagent:")} ${theme.fg("accent", result.agent)}`,
|
||||||
|
`${theme.fg("dim", "Task:")} ${theme.fg("muted", result.task)}`,
|
||||||
|
`${theme.fg("dim", "Status:")} ${theme.fg(statusColor, statusText)}`,
|
||||||
|
`${theme.fg("dim", "Tool Calls:")} ${theme.fg(
|
||||||
|
"muted",
|
||||||
|
`${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
|
||||||
|
)}`,
|
||||||
|
];
|
||||||
|
|
||||||
|
if (result.status.recentToolCalls.length > 0) {
|
||||||
|
lines.push("", theme.fg("dim", "Tool Calls - Last 3:"));
|
||||||
|
for (const call of result.status.recentToolCalls.slice(0, 3)) {
|
||||||
|
const mark =
|
||||||
|
call.status === "done"
|
||||||
|
? theme.fg("success", "✓")
|
||||||
|
: call.status === "error"
|
||||||
|
? theme.fg("error", "✗")
|
||||||
|
: theme.fg("warning", "•");
|
||||||
|
const tool = theme.fg("toolTitle", call.toolName.padEnd(6));
|
||||||
|
const summary = call.summary ? ` ${theme.fg("muted", call.summary)}` : "";
|
||||||
|
lines.push(`- ${mark} ${tool}${summary}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return new Text(lines.join("\n"), 0, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Render Subagent Call - Keep tool invocations readable in compact displays.
|
||||||
|
export function renderSubagentCall(
|
||||||
|
args: { agent: string; task: string },
|
||||||
|
theme: Theme,
|
||||||
|
): Text {
|
||||||
|
const task = truncateOneLine(args.task, 90);
|
||||||
|
const text =
|
||||||
|
theme.fg("toolTitle", theme.bold("subagent ")) +
|
||||||
|
theme.fg("accent", args.agent) +
|
||||||
|
(task ? ` ${theme.fg("dim", task)}` : "");
|
||||||
|
return new Text(text, 0, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Render Subagent Result - Show live status while partial, markdown after completion.
|
||||||
|
export function renderSubagentResult(
|
||||||
|
result: { content: { type: string; text?: string }[]; details?: unknown },
|
||||||
|
options: { expanded: boolean; isPartial: boolean },
|
||||||
|
theme: Theme,
|
||||||
|
): Text | Markdown {
|
||||||
|
const details = result.details as Partial<SubagentResult> | undefined;
|
||||||
|
if (options.isPartial && details?.status && details.agent && details.task) {
|
||||||
|
return options.expanded
|
||||||
|
? renderSubagentStatus(details as SubagentResult, theme)
|
||||||
|
: renderSubagentSummary(details as SubagentResult, theme);
|
||||||
|
}
|
||||||
|
|
||||||
|
const text = result.content[0];
|
||||||
|
return new Markdown(
|
||||||
|
text?.type === "text" ? (text.text ?? "") : "",
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
getSubagentMarkdownTheme(theme),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export { formatSubagentContent, FinalizeStatus };
|
||||||
247
src/runner.ts
Normal file
247
src/runner.ts
Normal file
@@ -0,0 +1,247 @@
|
|||||||
|
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 { withFileMutationQueue } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { EXTENSION_ENTRY, FINALIZE_TOOL_NAME } from "./constants.ts";
|
||||||
|
import { formatStatusText, formatToolArgs } from "./format.ts";
|
||||||
|
import { buildSubagentPrompt } from "./prompts.ts";
|
||||||
|
import type {
|
||||||
|
PromptConfig,
|
||||||
|
SubagentResult,
|
||||||
|
SubagentToolActivity,
|
||||||
|
} from "./types.ts";
|
||||||
|
import { validateFinalizePayload } from "./tools.ts";
|
||||||
|
|
||||||
|
// Write Prompt - pi accepts appended system prompts via file path.
|
||||||
|
async function writePromptToTempFile(
|
||||||
|
agentName: string,
|
||||||
|
prompt: string,
|
||||||
|
): Promise<{ dir: string; filePath: string }> {
|
||||||
|
const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
|
||||||
|
const safeName = agentName.replace(/[^\w.-]+/g, "_");
|
||||||
|
const filePath = path.join(dir, `${safeName}.md`);
|
||||||
|
await withFileMutationQueue(filePath, async () => {
|
||||||
|
await fs.promises.writeFile(filePath, prompt, {
|
||||||
|
encoding: "utf-8",
|
||||||
|
mode: 0o600,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
return { dir, filePath };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pi Invocation - Prefer the current pi entrypoint when running inside pi.
|
||||||
|
function getPiInvocation(args: string[]): { command: string; args: string[] } {
|
||||||
|
const currentScript = process.argv[1];
|
||||||
|
const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
|
||||||
|
if (currentScript && !isBunVirtualScript && fs.existsSync(currentScript)) {
|
||||||
|
return { command: process.execPath, args: [currentScript, ...args] };
|
||||||
|
}
|
||||||
|
|
||||||
|
const execName = path.basename(process.execPath).toLowerCase();
|
||||||
|
const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
|
||||||
|
if (!isGenericRuntime) return { command: process.execPath, args };
|
||||||
|
return { command: "pi", args };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run Agent - Spawn an isolated pi process in JSON mode and collect final text.
|
||||||
|
export async function runAgent(
|
||||||
|
cwd: string,
|
||||||
|
agent: PromptConfig,
|
||||||
|
task: string,
|
||||||
|
tools: string[],
|
||||||
|
sessionId: string,
|
||||||
|
sessionPath: string,
|
||||||
|
signal?: AbortSignal,
|
||||||
|
onUpdate?: (partial: {
|
||||||
|
content: { type: "text"; text: string }[];
|
||||||
|
details: SubagentResult;
|
||||||
|
}) => void,
|
||||||
|
): Promise<SubagentResult> {
|
||||||
|
let tmpDir: string | null = null;
|
||||||
|
let tmpPromptPath: string | null = null;
|
||||||
|
|
||||||
|
const result: SubagentResult = {
|
||||||
|
agent: agent.name,
|
||||||
|
task,
|
||||||
|
tools,
|
||||||
|
sessionId,
|
||||||
|
exitCode: 0,
|
||||||
|
output: "",
|
||||||
|
stderr: "",
|
||||||
|
status: {
|
||||||
|
state: "starting",
|
||||||
|
toolCallCount: 0,
|
||||||
|
activeToolCalls: 0,
|
||||||
|
recentToolCalls: [],
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const activeToolIds = new Set<string>();
|
||||||
|
|
||||||
|
const emitUpdate = () => {
|
||||||
|
onUpdate?.({
|
||||||
|
content: [{ type: "text", text: formatStatusText(result) }],
|
||||||
|
details: {
|
||||||
|
...result,
|
||||||
|
status: {
|
||||||
|
...result.status,
|
||||||
|
recentToolCalls: [...result.status.recentToolCalls],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
const rememberToolCall = (activity: SubagentToolActivity) => {
|
||||||
|
const existing = result.status.recentToolCalls.findIndex(
|
||||||
|
(call) => call.id === activity.id,
|
||||||
|
);
|
||||||
|
if (existing >= 0) result.status.recentToolCalls.splice(existing, 1);
|
||||||
|
result.status.recentToolCalls.unshift(activity);
|
||||||
|
result.status.recentToolCalls = result.status.recentToolCalls.slice(0, 3);
|
||||||
|
};
|
||||||
|
|
||||||
|
try {
|
||||||
|
const args = [
|
||||||
|
"--mode",
|
||||||
|
"json",
|
||||||
|
"-p",
|
||||||
|
"--session",
|
||||||
|
sessionPath,
|
||||||
|
"--extension",
|
||||||
|
EXTENSION_ENTRY,
|
||||||
|
];
|
||||||
|
args.push("--tools", tools.join(","));
|
||||||
|
|
||||||
|
const prompt = buildSubagentPrompt(agent);
|
||||||
|
if (prompt) {
|
||||||
|
const tmp = await writePromptToTempFile(agent.name, prompt);
|
||||||
|
tmpDir = tmp.dir;
|
||||||
|
tmpPromptPath = tmp.filePath;
|
||||||
|
args.push("--append-system-prompt", tmpPromptPath);
|
||||||
|
}
|
||||||
|
|
||||||
|
args.push(task);
|
||||||
|
emitUpdate();
|
||||||
|
|
||||||
|
const exitCode = await new Promise<number>((resolve) => {
|
||||||
|
const invocation = getPiInvocation(args);
|
||||||
|
const proc = spawn(invocation.command, invocation.args, {
|
||||||
|
cwd,
|
||||||
|
env: { ...process.env, PI_SUBAGENT_CHILD: "1" },
|
||||||
|
shell: false,
|
||||||
|
stdio: ["ignore", "pipe", "pipe"],
|
||||||
|
});
|
||||||
|
|
||||||
|
let buffer = "";
|
||||||
|
let aborted = false;
|
||||||
|
|
||||||
|
const processLine = (line: string) => {
|
||||||
|
if (!line.trim()) return;
|
||||||
|
try {
|
||||||
|
const event = JSON.parse(line);
|
||||||
|
if (
|
||||||
|
event.type === "message_start" &&
|
||||||
|
event.message?.role === "assistant"
|
||||||
|
) {
|
||||||
|
result.status.state = "thinking";
|
||||||
|
emitUpdate();
|
||||||
|
}
|
||||||
|
if (event.type === "tool_execution_start") {
|
||||||
|
const id = String(
|
||||||
|
event.toolCallId ?? result.status.toolCallCount + 1,
|
||||||
|
);
|
||||||
|
const toolName = String(event.toolName ?? "tool");
|
||||||
|
if (toolName === FINALIZE_TOOL_NAME) {
|
||||||
|
result.finalized =
|
||||||
|
validateFinalizePayload(event.args) ?? undefined;
|
||||||
|
}
|
||||||
|
activeToolIds.add(id);
|
||||||
|
result.status.state = "running";
|
||||||
|
result.status.toolCallCount += 1;
|
||||||
|
result.status.activeToolCalls = activeToolIds.size;
|
||||||
|
rememberToolCall({
|
||||||
|
id,
|
||||||
|
toolName,
|
||||||
|
summary: formatToolArgs(event.args),
|
||||||
|
status: "running",
|
||||||
|
});
|
||||||
|
emitUpdate();
|
||||||
|
}
|
||||||
|
if (event.type === "tool_execution_end") {
|
||||||
|
const id = String(event.toolCallId ?? "");
|
||||||
|
if (id) activeToolIds.delete(id);
|
||||||
|
result.status.activeToolCalls = activeToolIds.size;
|
||||||
|
const previous = result.status.recentToolCalls.find(
|
||||||
|
(call) => call.id === id,
|
||||||
|
);
|
||||||
|
rememberToolCall({
|
||||||
|
id,
|
||||||
|
toolName: String(event.toolName ?? previous?.toolName ?? "tool"),
|
||||||
|
summary: previous?.summary ?? "",
|
||||||
|
status: event.isError ? "error" : "done",
|
||||||
|
});
|
||||||
|
emitUpdate();
|
||||||
|
}
|
||||||
|
if (
|
||||||
|
event.type === "message_end" &&
|
||||||
|
event.message?.role === "assistant"
|
||||||
|
) {
|
||||||
|
for (const part of event.message.content ?? []) {
|
||||||
|
if (part.type === "text") result.output = part.text;
|
||||||
|
}
|
||||||
|
if (event.message.stopReason === "error") {
|
||||||
|
result.error =
|
||||||
|
event.message.errorMessage ?? "Subagent stopped with an error.";
|
||||||
|
}
|
||||||
|
emitUpdate();
|
||||||
|
}
|
||||||
|
} catch {
|
||||||
|
// Ignore Non-JSON Output.
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
proc.stdout.on("data", (data) => {
|
||||||
|
buffer += data.toString();
|
||||||
|
const lines = buffer.split("\n");
|
||||||
|
buffer = lines.pop() ?? "";
|
||||||
|
for (const line of lines) processLine(line);
|
||||||
|
});
|
||||||
|
|
||||||
|
proc.stderr.on("data", (data) => {
|
||||||
|
result.stderr += data.toString();
|
||||||
|
});
|
||||||
|
|
||||||
|
proc.on("close", (code) => {
|
||||||
|
if (buffer.trim()) processLine(buffer);
|
||||||
|
resolve(aborted ? 130 : (code ?? 0));
|
||||||
|
});
|
||||||
|
|
||||||
|
proc.on("error", (error) => {
|
||||||
|
result.error = error.message;
|
||||||
|
resolve(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
if (signal) {
|
||||||
|
const killProc = () => {
|
||||||
|
aborted = true;
|
||||||
|
proc.kill("SIGTERM");
|
||||||
|
setTimeout(() => {
|
||||||
|
if (!proc.killed) proc.kill("SIGKILL");
|
||||||
|
}, 5000);
|
||||||
|
};
|
||||||
|
if (signal.aborted) killProc();
|
||||||
|
else signal.addEventListener("abort", killProc, { once: true });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
result.exitCode = exitCode;
|
||||||
|
result.status.state = "done";
|
||||||
|
result.status.activeToolCalls = 0;
|
||||||
|
if (exitCode === 130) result.error = "Subagent was aborted.";
|
||||||
|
return result;
|
||||||
|
} finally {
|
||||||
|
if (tmpPromptPath) await fs.promises.rm(tmpPromptPath, { force: true });
|
||||||
|
if (tmpDir) await fs.promises.rm(tmpDir, { force: true, recursive: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
21
src/session.ts
Normal file
21
src/session.ts
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
import { createHash } from "node:crypto";
|
||||||
|
import * as os from "node:os";
|
||||||
|
import * as path from "node:path";
|
||||||
|
|
||||||
|
// Session Path - Persist child sessions as <agent>_<uuid>.jsonl under a cwd hash.
|
||||||
|
export function getSubagentSessionPath(
|
||||||
|
cwd: string,
|
||||||
|
agentName: string,
|
||||||
|
sessionId: string,
|
||||||
|
): string {
|
||||||
|
const cwdHash = createHash("sha256").update(cwd).digest("hex").slice(0, 16);
|
||||||
|
const safeAgent = agentName.replace(/[^\w.-]+/g, "_");
|
||||||
|
const safeSessionId = sessionId.replace(/[^\w.-]+/g, "_");
|
||||||
|
return path.join(
|
||||||
|
os.homedir(),
|
||||||
|
".pi",
|
||||||
|
"subagent-sessions",
|
||||||
|
cwdHash,
|
||||||
|
`${safeAgent}_${safeSessionId}.jsonl`,
|
||||||
|
);
|
||||||
|
}
|
||||||
65
src/tools.ts
Normal file
65
src/tools.ts
Normal file
@@ -0,0 +1,65 @@
|
|||||||
|
import { Type } from "typebox";
|
||||||
|
import { FINALIZE_TOOL_NAME } from "./constants.ts";
|
||||||
|
import type { PromptConfig, SubagentFinalizePayload } from "./types.ts";
|
||||||
|
import { FinalizeStatus } from "./types.ts";
|
||||||
|
|
||||||
|
// Resolve Tools - Use exactly one permission mode. approved_tools/allowed_tools
|
||||||
|
// is a whitelist; denied_tools is a blacklist over the currently active tools.
|
||||||
|
export function resolveTools(
|
||||||
|
agent: PromptConfig,
|
||||||
|
activeTools: string[],
|
||||||
|
): string[] {
|
||||||
|
const withoutDelegationTools = (tools: string[]) =>
|
||||||
|
tools.filter((tool) => tool !== "subagent" && tool !== FINALIZE_TOOL_NAME);
|
||||||
|
|
||||||
|
const resolved =
|
||||||
|
agent.approvedTools.length > 0
|
||||||
|
? withoutDelegationTools([...new Set(agent.approvedTools)])
|
||||||
|
: withoutDelegationTools(
|
||||||
|
[...new Set(activeTools)].filter(
|
||||||
|
(tool) => !new Set(agent.deniedTools).has(tool),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
|
return [...new Set([...resolved, FINALIZE_TOOL_NAME])];
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate Finalize Payload - Keep the parent contract strict and small.
|
||||||
|
export function validateFinalizePayload(
|
||||||
|
value: unknown,
|
||||||
|
): SubagentFinalizePayload | null {
|
||||||
|
if (!value || typeof value !== "object") return null;
|
||||||
|
const payload = value as Record<string, unknown>;
|
||||||
|
if (payload.status === FinalizeStatus.SUCCESS) {
|
||||||
|
return typeof payload.result === "string" && payload.result.trim()
|
||||||
|
? { status: FinalizeStatus.SUCCESS, result: payload.result }
|
||||||
|
: null;
|
||||||
|
}
|
||||||
|
if (payload.status === FinalizeStatus.ERROR) {
|
||||||
|
const result =
|
||||||
|
typeof payload.result === "string" ? payload.result : undefined;
|
||||||
|
return typeof payload.error === "string" && payload.error.trim()
|
||||||
|
? { status: FinalizeStatus.ERROR, error: payload.error, result }
|
||||||
|
: null;
|
||||||
|
}
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
export const SubagentParams = Type.Object({
|
||||||
|
agent: Type.String({
|
||||||
|
description: "Name of the prompt-defined subagent to invoke",
|
||||||
|
}),
|
||||||
|
task: Type.String({ description: "Task to delegate to the subagent" }),
|
||||||
|
sessionId: Type.Optional(
|
||||||
|
Type.String({
|
||||||
|
description:
|
||||||
|
"Optional sticky subagent session id. Reuse to continue a previous subagent context.",
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
cwd: Type.Optional(
|
||||||
|
Type.String({
|
||||||
|
description:
|
||||||
|
"Working directory for the subagent process. Defaults to current cwd.",
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
});
|
||||||
46
src/types.ts
Normal file
46
src/types.ts
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
export interface PromptConfig {
|
||||||
|
name: string;
|
||||||
|
description: string;
|
||||||
|
approvedTools: string[];
|
||||||
|
deniedTools: string[];
|
||||||
|
systemPrompt: string;
|
||||||
|
filePath: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SubagentToolActivity {
|
||||||
|
id: string;
|
||||||
|
toolName: string;
|
||||||
|
summary: string;
|
||||||
|
status: "running" | "done" | "error";
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SubagentStatus {
|
||||||
|
state: "starting" | "thinking" | "running" | "done";
|
||||||
|
toolCallCount: number;
|
||||||
|
activeToolCalls: number;
|
||||||
|
recentToolCalls: SubagentToolActivity[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export enum FinalizeStatus {
|
||||||
|
SUCCESS = "SUCCESS",
|
||||||
|
ERROR = "ERROR",
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SubagentFinalizePayload {
|
||||||
|
status: FinalizeStatus;
|
||||||
|
result?: string;
|
||||||
|
error?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SubagentResult {
|
||||||
|
agent: string;
|
||||||
|
task: string;
|
||||||
|
tools: string[];
|
||||||
|
sessionId: string;
|
||||||
|
exitCode: number;
|
||||||
|
output: string;
|
||||||
|
stderr: string;
|
||||||
|
status: SubagentStatus;
|
||||||
|
finalized?: SubagentFinalizePayload;
|
||||||
|
error?: string;
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user