From 9488aab237a54e86cd05677f507469fa2fcb6434 Mon Sep 17 00:00:00 2001 From: Evan Reichard Date: Tue, 12 May 2026 12:48:26 -0400 Subject: [PATCH] feat(subagent): add sticky child finalization --- AGENTS.md | 33 +++++++ README.md | 75 +++++++++++++++- index.ts | 261 ++++++++++++++++++++++++++++++++++++++++++++++++------ 3 files changed, 340 insertions(+), 29 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..530c783 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,33 @@ +# Repo Guidelines + +## Scope + +This repo implements the `subagent` pi extension in `index.ts`. + +## Key Behavior + +- Parent-facing tool: `subagent`. +- Internal child-only tool: `subagent_finalize`. +- Never register `subagent_finalize` in the parent context; only when `PI_SUBAGENT_CHILD=1`. +- Subagent sessions are sticky and persisted at: + `~/.pi/subagent-sessions//_.jsonl`. +- Omitting `sessionId` creates a new UUID-backed session. +- Passing `sessionId` resumes the same agent/cwd child session. + +## Validation + +- Child agents must finish by calling `subagent_finalize`. +- Valid success: `status=SUCCESS` with non-empty `result`. +- Valid error: `status=ERROR` with non-empty `error`; `result` is optional partial findings. +- If finalization is missing/invalid, retry by continuing the same sticky session, capped by `MAX_FINALIZE_RETRIES`. + +## Commands + +- Typecheck: `npm run typecheck` +- Lint: `npm run lint` + +## Style + +- Keep the extension simple; avoid debug metadata unless needed. +- Do not expose retry internals to the parent by default. +- Prefer precise tool/context isolation over prompt-only enforcement. diff --git a/README.md b/README.md index 4cfa1f9..dd3cf11 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ System prompt for the subagent goes here. ``` - `approved_tools` / `allowed_tools` is whitelist mode: the subagent receives only those tools. -- `denied_tools` is blacklist mode: the subagent inherits currently active tools except those tools and `subagent` itself. +- `denied_tools` is blacklist mode: the subagent inherits currently active tools except `subagent` and internal tools. - Do not define whitelist and blacklist fields together; the extension rejects ambiguous configs. ## Tool usage @@ -28,4 +28,75 @@ System prompt for the subagent goes here. } ``` -The extension spawns an isolated `pi --mode json -p --no-session` subprocess with the prompt appended as system instructions and `--tools` set to the resolved tool list. +The tool returns the subagent's finalized result as normal text with a compact status/session header for clients that hide structured tool metadata: + +```md +**Status:** SUCCESS +**Session ID:** `550e8400-e29b-41d4-a716-446655440000` + +--- + +... +``` + +Metadata also includes the sticky session id: + +```json +{ + "sessionId": "550e8400-e29b-41d4-a716-446655440000", + "finalized": { + "status": "SUCCESS", + "result": "..." + } +} +``` + +To continue the same child subagent context, pass the returned `sessionId`: + +```json +{ + "agent": "scout", + "sessionId": "550e8400-e29b-41d4-a716-446655440000", + "task": "Keep going and inspect the registration path" +} +``` + +If `sessionId` is omitted, a new child session is created. + +## Sticky sessions + +Subagent sessions persist under: + +```txt +~/.pi/subagent-sessions//_.jsonl +``` + +Session identity is scoped by working directory, agent name, and session id. Reusing the same `sessionId` with the same agent and cwd resumes that child context. + +## Finalization + +Subagents do not return final answers as free-form assistant text. Child runs get an internal `subagent_finalize` tool and must call it as their final action: + +```ts +subagent_finalize({ + status: "SUCCESS" | "ERROR", + result?: string, + error?: string, +}) +``` + +Rules: + +- `SUCCESS` requires a non-empty `result`. +- `ERROR` requires a non-empty `error`. +- `ERROR.result` may contain partial findings. +- If the child exits without valid finalization, the extension continues the same session with corrective feedback, up to a small retry limit. +- `subagent_finalize` is registered only in child subagent processes, never in the parent context. + +## How it runs + +The extension spawns an isolated child pi process in JSON mode with the resolved tools and a persistent session file: + +```txt +pi --mode json -p --session --tools +``` diff --git a/index.ts b/index.ts index a52b485..c36bb32 100644 --- a/index.ts +++ b/index.ts @@ -1,6 +1,7 @@ // 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 * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; @@ -36,19 +37,44 @@ interface SubagentStatus { recentToolCalls: SubagentToolActivity[]; } +interface SubagentFinalizePayload { + status: "SUCCESS" | "ERROR"; + 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_DIR = path.dirname(fileURLToPath(import.meta.url)); +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; + +// Format Tool Content - Some clients hide structured details from the model. +function formatSubagentContent( + status: "SUCCESS" | "ERROR", + 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[] { @@ -108,18 +134,70 @@ function discoverPrompts(): PromptConfig[] { // 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[] { - if (agent.approvedTools.length > 0) { - return [...new Set(agent.approvedTools)].filter( - (tool) => tool !== "subagent", - ); - } + const withoutDelegationTools = (tools: string[]) => + tools.filter((tool) => tool !== "subagent" && tool !== FINALIZE_TOOL_NAME); - const denied = new Set(agent.deniedTools); - return [...new Set(activeTools)].filter( - (tool) => tool !== "subagent" && !denied.has(tool), + 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 SUCCESS with result, or status 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 === "SUCCESS") { + return typeof payload.result === "string" && payload.result.trim() + ? { status: "SUCCESS", result: payload.result } + : null; + } + if (payload.status === "ERROR") { + const result = + typeof payload.result === "string" ? payload.result : undefined; + return typeof payload.error === "string" && payload.error.trim() + ? { status: "ERROR", error: payload.error, result } + : null; + } + return null; +} + // Write Prompt - pi accepts appended system prompts via file path. async function writePromptToTempFile( agentName: string, @@ -256,6 +334,8 @@ async function runAgent( agent: PromptConfig, task: string, tools: string[], + sessionId: string, + sessionPath: string, signal?: AbortSignal, onUpdate?: (partial: { content: { type: "text"; text: string }[]; @@ -269,6 +349,7 @@ async function runAgent( agent: agent.name, task, tools, + sessionId, exitCode: 0, output: "", stderr: "", @@ -305,23 +386,33 @@ async function runAgent( }; try { - const args = ["--mode", "json", "-p", "--no-session"]; + const args = [ + "--mode", + "json", + "-p", + "--session", + sessionPath, + "--extension", + EXTENSION_ENTRY, + ]; args.push("--tools", tools.join(",")); - if (agent.systemPrompt) { - const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt); + 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: ${task}`); + 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"], }); @@ -344,13 +435,17 @@ async function runAgent( 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: String(event.toolName ?? "tool"), + toolName, summary: formatToolArgs(event.args), status: "running", }); @@ -447,6 +542,12 @@ const SubagentParams = Type.Object({ 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: @@ -456,6 +557,30 @@ const SubagentParams = Type.Object({ }); export default function (pi: ExtensionAPI) { + if (process.env.PI_SUBAGENT_CHILD === "1") { + pi.registerTool({ + name: FINALIZE_TOOL_NAME, + label: "Subagent Finalize", + description: + "Internal subagent-only tool. Call this as your final action when delegated work is complete.", + promptSnippet: + "Call subagent_finalize as your final action when subagent work is complete.", + parameters: Type.Object({ + status: Type.Union([Type.Literal("SUCCESS"), Type.Literal("ERROR")]), + result: Type.Optional(Type.String()), + error: Type.Optional(Type.String()), + }), + + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: "Subagent finalized." }], + details: params, + terminate: true, + }; + }, + }); + } + pi.registerTool({ name: "subagent", label: "Subagent", @@ -517,21 +642,103 @@ export default function (pi: ExtensionAPI) { } const cwd = path.resolve(ctx.cwd, params.cwd ?? "."); - const result = await runAgent( - cwd, - agent, - params.task, - tools, - signal, - _onUpdate, - ); - const failed = result.exitCode !== 0 || Boolean(result.error); - const fallback = result.error || result.stderr || "(no output)"; + const sessionId = params.sessionId ?? randomUUID(); + const sessionPath = getSubagentSessionPath(cwd, agent.name, sessionId); + await fs.promises.mkdir(path.dirname(sessionPath), { recursive: true }); + + let result: SubagentResult | null = null; + for (let retryCount = 0; retryCount <= MAX_FINALIZE_RETRIES; retryCount += 1) { + const task = + retryCount === 0 + ? `Task: ${params.task}` + : [ + "Your previous response did not finalize correctly.", + `If you are finished, call ${FINALIZE_TOOL_NAME}.`, + "If you are not finished, continue the original task using available tools as needed.", + `Original task: ${params.task}`, + ].join("\n\n"); + + result = await runAgent( + cwd, + agent, + task, + tools, + sessionId, + sessionPath, + signal, + _onUpdate, + ); + + if (result.finalized) break; + if (result.exitCode !== 0 || result.error) break; + } + + if (!result) { + return { + content: [ + { + type: "text", + text: formatSubagentContent( + "ERROR", + sessionId, + undefined, + "Subagent did not run.", + ), + }, + ], + details: { sessionId }, + isError: true, + }; + } + + if (!result.finalized) { + const fallback = + result.error || + result.stderr || + `Subagent did not call ${FINALIZE_TOOL_NAME}.`; + return { + content: [ + { + type: "text", + text: formatSubagentContent("ERROR", sessionId, undefined, fallback), + }, + ], + details: { sessionId }, + isError: true, + }; + } + + if (result.finalized.status === "ERROR") { + return { + content: [ + { + type: "text", + text: formatSubagentContent( + "ERROR", + sessionId, + result.finalized.result, + result.finalized.error ?? "Subagent failed.", + ), + }, + ], + details: { sessionId, finalized: result.finalized }, + isError: true, + }; + } return { - content: [{ type: "text", text: result.output || fallback }], - details: result, - isError: failed, + content: [ + { + type: "text", + text: formatSubagentContent( + "SUCCESS", + sessionId, + result.finalized.result, + ), + }, + ], + details: { sessionId, finalized: result.finalized }, + isError: false, }; },