diff --git a/index.ts b/index.ts index 3f1fb34..c994c2b 100644 --- a/index.ts +++ b/index.ts @@ -1,591 +1,18 @@ // Subagent Extension - Registers a tool for delegating work to prompt-defined // subagents with constrained tool permissions. -import { spawn } from "node:child_process"; -import { createHash, randomUUID } from "node:crypto"; +import { randomUUID } from "node:crypto"; 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 { 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 { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { Type } from "typebox"; - -interface PromptConfig { - name: string; - description: string; - approvedTools: string[]; - deniedTools: string[]; - systemPrompt: string; - filePath: string; -} - -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>(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 _.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; - 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; - 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 { - 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(); - - 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((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.", - }), - ), -}); +import { FINALIZE_TOOL_NAME, MAX_FINALIZE_RETRIES } from "./src/constants.ts"; +import { formatPromptList, formatSubagentContent } from "./src/format.ts"; +import { discoverPrompts } from "./src/prompts.ts"; +import { renderSubagentCall, renderSubagentResult } from "./src/render.ts"; +import { runAgent } from "./src/runner.ts"; +import { getSubagentSessionPath } from "./src/session.ts"; +import { FinalizeStatus, type SubagentResult } from "./src/types.ts"; +import { resolveTools, SubagentParams } from "./src/tools.ts"; export default function (pi: ExtensionAPI) { if (process.env.PI_SUBAGENT_CHILD === "1") { @@ -783,29 +210,11 @@ export default function (pi: ExtensionAPI) { }, renderCall(args, theme) { - 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); + return renderSubagentCall(args, theme); }, - renderResult(result, { expanded, isPartial }, theme) { - const details = result.details as Partial | undefined; - 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), - ); + renderResult(result, options, theme) { + return renderSubagentResult(result, options, theme); }, }); } diff --git a/src/constants.ts b/src/constants.ts new file mode 100644 index 0000000..b3170e5 --- /dev/null +++ b/src/constants.ts @@ -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; diff --git a/src/format.ts b/src/format.ts new file mode 100644 index 0000000..a4802a9 --- /dev/null +++ b/src/format.ts @@ -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; + 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("; "); +} diff --git a/src/prompts.ts b/src/prompts.ts new file mode 100644 index 0000000..5d9cae7 --- /dev/null +++ b/src/prompts.ts @@ -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>(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"); +} diff --git a/src/render.ts b/src/render.ts new file mode 100644 index 0000000..bf022b6 --- /dev/null +++ b/src/render.ts @@ -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 | 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 }; diff --git a/src/runner.ts b/src/runner.ts new file mode 100644 index 0000000..8a55174 --- /dev/null +++ b/src/runner.ts @@ -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 { + 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(); + + 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((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 }); + } +} diff --git a/src/session.ts b/src/session.ts new file mode 100644 index 0000000..b999a03 --- /dev/null +++ b/src/session.ts @@ -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 _.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`, + ); +} diff --git a/src/tools.ts b/src/tools.ts new file mode 100644 index 0000000..8562955 --- /dev/null +++ b/src/tools.ts @@ -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; + 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.", + }), + ), +}); diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..491f3a0 --- /dev/null +++ b/src/types.ts @@ -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; +}