From 821df5c59c73aaf81cd0c6ba0d1004ac43c0beca Mon Sep 17 00:00:00 2001 From: Krishna Penukonda Date: Mon, 3 Aug 2026 21:19:25 +0530 Subject: [PATCH] feat(pi-subagents): add inspection tools --- .../pi-subagents/skills/pi-subagents/SKILL.md | 13 +- .../pi-subagents/src/extension/index.ts | 203 ++++++++++++++++++ .../pi-subagents/src/extension/schemas.ts | 20 ++ .../test/unit/minimal-subagents.test.ts | 153 +++++++++++++ 4 files changed, 386 insertions(+), 3 deletions(-) diff --git a/extensions/pi-subagents/skills/pi-subagents/SKILL.md b/extensions/pi-subagents/skills/pi-subagents/SKILL.md index b5ff766..0595f80 100644 --- a/extensions/pi-subagents/skills/pi-subagents/SKILL.md +++ b/extensions/pi-subagents/skills/pi-subagents/SKILL.md @@ -15,6 +15,13 @@ Use this tool to launch unrestricted child Pi sessions. The caller must include - The returned subagent id is also the child Pi session id. - Child output is written to `result.log` under the child subagent directory. - When `model` is omitted, the child inherits the parent session's active model (e.g., `openai-codex/gpt-5.6-sol`). Pass an explicit `model` to override (e.g., `"anthropic/claude-sonnet-4-5"`). +- `list_subagents({})` + - Returns persisted subagent ids and latest running state for the current parent session. +- `get_subagent_status({ id })` + - Returns one status snapshot and result path for a spawned subagent. +- `tail_subagent({ id, lines? })` + - Returns recent complete NDJSON lines from child stdout (`lines` defaults to 20; maximum 200). + - Omits a trailing partial line that the child may still be writing. ## Usage @@ -31,7 +38,7 @@ Calls return immediately; the parent will be notified when each subagent complet ## Rules -- Available subagent tool is exactly `spawn_subagent`. +- Inspection tools are read-only snapshots. Do not build polling or sleep loops around them; completion still sends a notification. - No wait-for-completion mode exists. - No subagent types exist. - No chain or parallel-list mode exists. @@ -47,8 +54,8 @@ Calls return immediately; the parent will be notified when each subagent complet When changing the pi-subagents extension contract, update every surface together: -1. Spawn tool schema in `extensions/pi-subagents/src/extension/schemas.ts`. -2. Spawn runtime validation and user-facing messages in `extensions/pi-subagents/src/extension/index.ts`. +1. Tool schemas in `extensions/pi-subagents/src/extension/schemas.ts`. +2. Tool runtime validation and user-facing messages in `extensions/pi-subagents/src/extension/index.ts`. 3. Skill docs in `extensions/pi-subagents/skills/pi-subagents/SKILL.md`. 4. GitHub issues/PR text exactly as requested by the user; do not fabricate details. diff --git a/extensions/pi-subagents/src/extension/index.ts b/extensions/pi-subagents/src/extension/index.ts index 8d3c92e..9e8f567 100644 --- a/extensions/pi-subagents/src/extension/index.ts +++ b/extensions/pi-subagents/src/extension/index.ts @@ -31,15 +31,24 @@ import { buildPiArgs, cleanupTempDir } from "../runs/shared/pi-args.ts"; // here would execute its child-only fd watcher in the parent extension process. const PI_SUBAGENT_LIFELINE_FD = "PI_SUBAGENT_LIFELINE_FD"; import { + GetSubagentStatusParams, + type GetSubagentStatusParamsLike, + ListSubagentsParams, SpawnSubagentParams, type SpawnSubagentParamsLike, + TailSubagentParams, + type TailSubagentParamsLike, } from "./schemas.ts"; interface ToolDetails { id?: string; + sessionId?: string; running?: boolean; resultPath?: string; model?: string; + error?: string; + subagents?: Array<{ id: string; running: boolean }>; + lines?: string[]; } interface PersistedSubagentRecord { @@ -186,6 +195,13 @@ function upsertRecord(record: PersistedSubagentRecord): void { writeStore(record.parentSessionId, store); } +function findRecord( + parentId: string, + id: string, +): PersistedSubagentRecord | undefined { + return readStore(parentId).records.find((record) => record.id === id); +} + function updateRecordFields( parentId: string, id: string, @@ -455,6 +471,92 @@ function refreshRecordFromDisk( return refreshed; } +function resultPathForRecord(record: PersistedSubagentRecord): string { + return ( + record.outputFile ?? + path.join(record.parentSessionId, "subagents", record.id, "result.log") + ); +} + +function formatStatus( + record: PersistedSubagentRecord, +): AgentToolResult { + const refreshed = refreshRecordFromDisk(record); + const details: ToolDetails = { + id: refreshed.id, + sessionId: refreshed.id, + running: refreshed.running, + resultPath: resultPathForRecord(refreshed), + ...(refreshed.error ? { error: refreshed.error } : {}), + }; + return { + content: [ + { type: "text", text: JSON.stringify(details, null, 2) }, + ...(refreshed.running + ? [{ + type: "text" as const, + text: "This is a snapshot. Do not poll or sleep for the result; you will be notified when the subagent completes.", + }] + : []), + ], + details, + }; +} + +const TAIL_SUBAGENT_DEFAULT_LINES = 20; +const TAIL_SUBAGENT_MAX_READ_BYTES = 1024 * 1024; + +function readRecentCompleteLines(filePath: string, lines: number): string[] { + let fd: number | undefined; + try { + fd = fs.openSync(filePath, "r"); + const snapshotSize = fs.fstatSync(fd).size; + if (snapshotSize === 0) return []; + + const windowStart = Math.max(0, snapshotSize - TAIL_SUBAGENT_MAX_READ_BYTES); + const readStart = windowStart > 0 ? windowStart - 1 : 0; + const buffer = Buffer.allocUnsafe(snapshotSize - readStart); + let bytesRead = 0; + while (bytesRead < buffer.length) { + const count = fs.readSync( + fd, + buffer, + bytesRead, + buffer.length - bytesRead, + readStart + bytesRead, + ); + if (count === 0) break; + bytesRead += count; + } + const snapshot = buffer.subarray(0, bytesRead); + const completeEnd = snapshot.lastIndexOf(0x0a); + if (completeEnd < 0) return []; + + let completeStart = 0; + if (readStart > 0) { + if (snapshot[0] === 0x0a) completeStart = 1; + else { + const firstNewline = snapshot.indexOf(0x0a); + if (firstNewline < 0) return []; + completeStart = firstNewline + 1; + } + } + if (completeEnd < completeStart) return []; + return snapshot + .subarray(completeStart, completeEnd) + .toString("utf-8") + .split("\n") + .map((line) => line.endsWith("\r") ? line.slice(0, -1) : line) + .filter((line) => line.length > 0) + .slice(-lines); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return []; + throw error; + } finally { + if (fd !== undefined) fs.closeSync(fd); + } +} + function reconcileStore(parentId: string): ReconcileResult { const store = readStore(parentId); const refreshedResults = store.records.map((r) => refreshRecord(r)); @@ -1211,7 +1313,108 @@ export default function registerSubagentExtension(pi: ExtensionAPI): void { }, }; + const listTool: ToolDefinition = { + name: "list_subagents", + label: "List subagents", + description: + "List persisted subagents for the current parent session and their latest status.", + parameters: ListSubagentsParams, + async execute( + _toolCallId, + _params: Record, + _signal, + _onUpdate, + ctx, + ) { + const { records } = reconcileStore(parentSessionId(ctx)); + const subagents = records.map((record) => ({ + id: record.id, + running: record.running, + })); + return { + content: [{ type: "text", text: JSON.stringify(subagents, null, 2) }], + details: { subagents }, + }; + }, + renderCall(_args, theme) { + return new Text( + theme.fg("toolTitle", theme.bold("list_subagents")), + 0, + 0, + ); + }, + }; + + const statusTool: ToolDefinition< + typeof GetSubagentStatusParams, + ToolDetails + > = { + name: "get_subagent_status", + label: "Get subagent status", + description: + "Get one snapshot of a subagent's status and result path. Running subagents still notify you on completion; do not poll this tool.", + parameters: GetSubagentStatusParams, + async execute( + _toolCallId, + params: GetSubagentStatusParamsLike, + _signal, + _onUpdate, + ctx, + ) { + const record = findRecord(parentSessionId(ctx), params.id); + if (!record) throw new Error(`Unknown subagent id: ${params.id}`); + return formatStatus(record); + }, + renderCall(args, theme) { + return new Text( + `${theme.fg("toolTitle", theme.bold("get_subagent_status "))}${theme.fg("accent", args.id)}`, + 0, + 0, + ); + }, + }; + + const tailTool: ToolDefinition = { + name: "tail_subagent", + label: "Tail subagent", + description: + "Read one snapshot of recent complete NDJSON lines from a subagent's stdout log. A trailing line still being written is omitted.", + parameters: TailSubagentParams, + async execute( + _toolCallId, + params: TailSubagentParamsLike, + _signal, + _onUpdate, + ctx, + ) { + const record = findRecord(parentSessionId(ctx), params.id); + if (!record) throw new Error(`Unknown subagent id: ${params.id}`); + const refreshed = refreshRecordFromDisk(record); + const lines = readRecentCompleteLines( + refreshed.stdoutFile, + params.lines ?? TAIL_SUBAGENT_DEFAULT_LINES, + ); + return { + content: [{ + type: "text", + text: lines.join("\n") || "No complete stdout lines.", + }], + details: { id: refreshed.id, running: refreshed.running, lines }, + }; + }, + renderCall(args, theme) { + return new Text( + `${theme.fg("toolTitle", theme.bold("tail_subagent "))}${theme.fg("accent", args.id)}`, + 0, + 0, + ); + }, + }; + pi.registerTool(spawnTool); + pi.registerTool(listTool); + pi.registerTool(statusTool); + pi.registerTool(tailTool); if (typeof pi.on === "function") { pi.on("session_start", (_event, ctx) => { diff --git a/extensions/pi-subagents/src/extension/schemas.ts b/extensions/pi-subagents/src/extension/schemas.ts index be2cdec..92af23c 100644 --- a/extensions/pi-subagents/src/extension/schemas.ts +++ b/extensions/pi-subagents/src/extension/schemas.ts @@ -33,6 +33,21 @@ export const ListSubagentsParams = Type.Object( { additionalProperties: false }, ); +export const TailSubagentParams = Type.Object( + { + id: Type.String({ description: "Subagent id returned by spawn_subagent." }), + lines: Type.Optional( + Type.Integer({ + minimum: 1, + maximum: 200, + default: 20, + description: "Number of recent complete NDJSON lines to return.", + }), + ), + }, + { additionalProperties: false }, +); + export interface SpawnSubagentParamsLike { task: string; cwd?: string; @@ -42,3 +57,8 @@ export interface SpawnSubagentParamsLike { export interface GetSubagentStatusParamsLike { id: string; } + +export interface TailSubagentParamsLike { + id: string; + lines?: number; +} diff --git a/extensions/pi-subagents/test/unit/minimal-subagents.test.ts b/extensions/pi-subagents/test/unit/minimal-subagents.test.ts index dccfa96..fc6e5f7 100644 --- a/extensions/pi-subagents/test/unit/minimal-subagents.test.ts +++ b/extensions/pi-subagents/test/unit/minimal-subagents.test.ts @@ -146,6 +146,7 @@ function registerTestTools(sendMessage: (...args: unknown[]) => void = () => {}) const rawSpawnTool = registered.get("spawn_subagent"); return { handlers, + tools: registered, spawnTool: { ...rawSpawnTool, async execute(callId: string, ...args: any[]) { @@ -177,6 +178,158 @@ test("spawn schema accepts task only and rejects removed properties", () => { ); }); +test("extension registers read-only subagent inspection tools with bounded tail schema", () => { + const { tools } = registerTestTools(); + + assert.deepEqual([...tools.keys()], [ + "spawn_subagent", + "list_subagents", + "get_subagent_status", + "tail_subagent", + ]); + assert.equal(Value.Check(tools.get("list_subagents").parameters, {}), true); + assert.equal( + Value.Check(tools.get("get_subagent_status").parameters, { id: "child-1" }), + true, + ); + + const tailSchema = tools.get("tail_subagent").parameters; + assert.equal(Value.Check(tailSchema, { id: "child-1" }), true); + assert.equal(Value.Check(tailSchema, { id: "child-1", lines: 1 }), true); + assert.equal(Value.Check(tailSchema, { id: "child-1", lines: 200 }), true); + assert.equal(Value.Check(tailSchema, { id: "child-1", lines: 0 }), false); + assert.equal(Value.Check(tailSchema, { id: "child-1", lines: 201 }), false); + assert.equal(Value.Check(tailSchema, { id: "child-1", lines: 1.5 }), false); + assert.equal(Value.Check(tailSchema, { id: "child-1", extra: true }), false); + assert.equal(tailSchema.properties.lines.default, 20); +}); + +test("list and status tools inspect records persisted for current parent", async () => { + const { sessionId, ctx } = makeTestCtx("pi-subagents-inspection"); + const childDir = path.join(sessionId, "subagents", "child-1"); + const outputFile = path.join(childDir, "result.log"); + const stdoutFile = path.join(childDir, "stdout.log"); + const stderrFile = path.join(childDir, "stderr.log"); + fs.mkdirSync(childDir, { recursive: true }); + fs.writeFileSync(outputFile, "done\n"); + fs.writeFileSync(stdoutFile, "{\"type\":\"done\"}\n"); + fs.writeFileSync(stderrFile, ""); + fs.writeFileSync( + storeFile(sessionId), + JSON.stringify({ + records: [{ + id: "child-1", + parentSessionId: sessionId, + cwd: ctx.cwd, + taskPreview: "inspect me", + model: "mock/model", + running: false, + outputFile, + stdoutFile, + stderrFile, + createdAt: 10, + updatedAt: 20, + completedAt: 20, + }], + }, null, 2), + ); + + try { + const { tools } = registerTestTools(); + const signal = new AbortController().signal; + const listed = await tools.get("list_subagents").execute( + "list-call", {}, signal, undefined, ctx, + ); + assert.deepEqual(listed.details.subagents, [{ + id: "child-1", + running: false, + }]); + + const status = await tools.get("get_subagent_status").execute( + "status-call", { id: "child-1" }, signal, undefined, ctx, + ); + assert.deepEqual(status.details, { + id: "child-1", + sessionId: "child-1", + running: false, + resultPath: outputFile, + }); + assert.deepEqual(JSON.parse(status.content[0].text), status.details); + } finally { + cleanupTestCtx(ctx, sessionId); + } +}); + +test("tail_subagent returns recent complete NDJSON lines and drops a trailing partial line", async () => { + const { sessionId, ctx } = makeTestCtx("pi-subagents-tail"); + const childDir = path.join(sessionId, "subagents", "child-tail"); + const stdoutFile = path.join(childDir, "stdout.log"); + fs.mkdirSync(childDir, { recursive: true }); + const completeLines = Array.from( + { length: 25 }, + (_, index) => JSON.stringify({ type: "event", index }), + ); + fs.writeFileSync(stdoutFile, `${completeLines.join("\n")}\n{\"partial\":`); + fs.writeFileSync( + storeFile(sessionId), + JSON.stringify({ + records: [{ + id: "child-tail", + parentSessionId: sessionId, + cwd: ctx.cwd, + taskPreview: "tail me", + running: true, + pid: process.pid, + outputFile: path.join(childDir, "result.log"), + stdoutFile, + stderrFile: path.join(childDir, "stderr.log"), + createdAt: 10, + updatedAt: 20, + }], + }, null, 2), + ); + + try { + const { tools } = registerTestTools(); + const result = await tools.get("tail_subagent").execute( + "tail-call", + { id: "child-tail" }, + new AbortController().signal, + undefined, + ctx, + ); + assert.equal(result.details.id, "child-tail"); + assert.equal(result.details.running, true); + assert.deepEqual(result.details.lines, completeLines.slice(-20)); + assert.equal(result.content[0].text, completeLines.slice(-20).join("\n")); + assert.doesNotMatch(result.content[0].text, /partial/); + } finally { + cleanupTestCtx(ctx, sessionId); + } +}); + +test("inspection tools reject ids outside current parent store", async () => { + const { sessionId, ctx } = makeTestCtx("pi-subagents-unknown-inspection"); + try { + const { tools } = registerTestTools(); + const signal = new AbortController().signal; + await assert.rejects( + () => tools.get("get_subagent_status").execute( + "status-call", { id: "missing" }, signal, undefined, ctx, + ), + /Unknown subagent id: missing/, + ); + await assert.rejects( + () => tools.get("tail_subagent").execute( + "tail-call", { id: "missing" }, signal, undefined, ctx, + ), + /Unknown subagent id: missing/, + ); + } finally { + cleanupTestCtx(ctx, sessionId); + } +}); + test("model override contract is synchronized across schema, tool, README, and skill", () => { const schemaDescription = (SpawnSubagentParams.properties.model as { description?: string }).description ?? "";