From cfb19c05d92bfe5b47ea7d2cb61cd83bfd4980c9 Mon Sep 17 00:00:00 2001 From: Krishna Penukonda Date: Mon, 20 Jul 2026 10:50:22 +0000 Subject: [PATCH] fix(pi-subagents): resolve result env aliases in file tools --- extensions/pi-subagents/README.md | 4 +- .../runs/shared/subagent-prompt-runtime.ts | 20 +- .../test/unit/minimal-subagents.test.ts | 177 ++++++++++++++---- 3 files changed, 167 insertions(+), 34 deletions(-) diff --git a/extensions/pi-subagents/README.md b/extensions/pi-subagents/README.md index dc16fc2..f465c11 100644 --- a/extensions/pi-subagents/README.md +++ b/extensions/pi-subagents/README.md @@ -31,12 +31,14 @@ Completed children are final. Running children continue their original task unti ## Child environment -Child Pi sessions keep normal Pi capabilities: tools, skills, extensions, and project context are not hidden or restricted by this extension. The only automatic child-system-prompt addition is: +Child Pi sessions keep normal Pi capabilities: tools, skills, extensions, and project context are not hidden or restricted by this extension. Every child system prompt receives this identity line: ```text You are a Pi subagent controlled by another Pi agent. ``` +Each child receives a resolved absolute `result.log` path. Pass that literal path to Pi file tools (`write`, `edit`, and `read`), which do not expand shell environment variables. Shell commands and programs may use `$PI_SUBAGENT_RESULT_PATH`; the child runtime also narrowly corrects either exact env-var token when it is accidentally passed as a file-tool path. If the child leaves `result.log` empty, its final assistant message is saved there automatically. + Children may spawn further subagents recursively until the configured recursion-depth limit is reached. ## Removed features diff --git a/extensions/pi-subagents/src/runs/shared/subagent-prompt-runtime.ts b/extensions/pi-subagents/src/runs/shared/subagent-prompt-runtime.ts index 75170c7..d04b261 100644 --- a/extensions/pi-subagents/src/runs/shared/subagent-prompt-runtime.ts +++ b/extensions/pi-subagents/src/runs/shared/subagent-prompt-runtime.ts @@ -7,6 +7,11 @@ export const CHILD_SUBAGENT_SYSTEM_LINE = "You are a Pi subagent controlled by another Pi agent."; const RESULT_PATH_MARKER = "Your result file:"; +const RESULT_PATH_ALIASES = new Set([ + "$PI_SUBAGENT_RESULT_PATH", + "${PI_SUBAGENT_RESULT_PATH}", +]); +const FILE_TOOL_NAMES = new Set(["write", "edit", "read"]); export function rewriteSubagentPrompt(prompt: string): string { if (prompt.includes(CHILD_SUBAGENT_SYSTEM_LINE)) return prompt; @@ -14,6 +19,19 @@ export function rewriteSubagentPrompt(prompt: string): string { } export default function registerSubagentPromptRuntime(pi: ExtensionAPI): void { + pi.on("tool_call", (event) => { + if (!FILE_TOOL_NAMES.has(event.toolName)) return; + const resultPath = process.env[SUBAGENT_RESULT_PATH_ENV]?.trim(); + const input = event.input as { path?: unknown }; + if ( + resultPath && + typeof input.path === "string" && + RESULT_PATH_ALIASES.has(input.path) + ) { + input.path = resultPath; + } + }); + pi.on("before_agent_start", async (event) => { const intercomSessionName = process.env[SUBAGENT_INTERCOM_SESSION_NAME_ENV]?.trim(); @@ -25,7 +43,7 @@ export default function registerSubagentPromptRuntime(pi: ExtensionAPI): void { const resultPath = process.env[SUBAGENT_RESULT_PATH_ENV]?.trim(); if (resultPath && !rewritten.includes(RESULT_PATH_MARKER)) { - rewritten = `${rewritten}\n\nYour result file: ${resultPath}\nYou may write your final output to this file at any time using any tool (e.g., write, bash). If you leave the file empty, your final assistant message will be automatically saved there on exit. The environment variable "$PI_SUBAGENT_RESULT_PATH" is aliased to ${resultPath}; you can pipe your answer there. Particularly for very large outputs, or for programmatic outputs, use tools to write the result directly to "$PI_SUBAGENT_RESULT_PATH".`; + rewritten = `${rewritten}\n\nYour result file: ${resultPath} (resolved absolute result path)\nPi file tools (\`write\`, \`edit\`, and \`read\`) must receive this literal absolute path as \`path\`; they do not expand shell environment variables. \`PI_SUBAGENT_RESULT_PATH\` contains the same path for shell commands and programs. Use \`$PI_SUBAGENT_RESULT_PATH\` only inside bash/shell commands. If you leave the result file empty, your final assistant message will be automatically saved there on exit.`; } if (rewritten === event.systemPrompt) return; diff --git a/extensions/pi-subagents/test/unit/minimal-subagents.test.ts b/extensions/pi-subagents/test/unit/minimal-subagents.test.ts index a09f201..f5cda0f 100644 --- a/extensions/pi-subagents/test/unit/minimal-subagents.test.ts +++ b/extensions/pi-subagents/test/unit/minimal-subagents.test.ts @@ -9,7 +9,7 @@ import { Value } from "typebox/value"; import registerSubagentExtension from "../../src/extension/index.ts"; import { SpawnSubagentParams } from "../../src/extension/schemas.ts"; import { createMockPi } from "../support/mock-pi.ts"; -import { +import registerSubagentPromptRuntime, { CHILD_SUBAGENT_SYSTEM_LINE, rewriteSubagentPrompt, SUBAGENT_RESULT_PATH_ENV, @@ -252,6 +252,118 @@ test("prompt runtime prepends exactly one line and preserves content", () => { assert(!rewritten.includes("Do not propose or run subagents")); }); +function registerPromptRuntimeHandlers() { + const handlers = new Map unknown>(); + registerSubagentPromptRuntime({ + on(event: string, handler: (event: any) => unknown) { + handlers.set(event, handler); + }, + } as never); + return handlers; +} + +async function withResultPath( + resultPath: string | undefined, + callback: () => Promise, +) { + const previousResultPath = process.env[SUBAGENT_RESULT_PATH_ENV]; + if (resultPath === undefined) delete process.env[SUBAGENT_RESULT_PATH_ENV]; + else process.env[SUBAGENT_RESULT_PATH_ENV] = resultPath; + try { + await callback(); + } finally { + if (previousResultPath === undefined) { + delete process.env[SUBAGENT_RESULT_PATH_ENV]; + } else { + process.env[SUBAGENT_RESULT_PATH_ENV] = previousResultPath; + } + } +} + +test("prompt runtime rewrites a literal result-path alias for file tools", async () => { + const toolCall = registerPromptRuntimeHandlers().get("tool_call"); + assert.ok(toolCall, "prompt runtime must register a tool_call guard"); + + await withResultPath("/tmp/subagents/abc/result.log", async () => { + const event = { + toolName: "write", + input: { path: "$PI_SUBAGENT_RESULT_PATH", content: "done" }, + }; + await toolCall(event); + assert.equal(event.input.path, "/tmp/subagents/abc/result.log"); + }); +}); + +test("prompt runtime rewrites both exact aliases for write edit and read", async () => { + const toolCall = registerPromptRuntimeHandlers().get("tool_call"); + assert.ok(toolCall); + + await withResultPath("/tmp/subagents/abc/result.log", async () => { + for (const [toolName, alias] of [ + ["write", "${PI_SUBAGENT_RESULT_PATH}"], + ["edit", "$PI_SUBAGENT_RESULT_PATH"], + ["read", "${PI_SUBAGENT_RESULT_PATH}"], + ] as const) { + const event = { toolName, input: { path: alias } }; + await toolCall(event); + assert.equal(event.input.path, "/tmp/subagents/abc/result.log"); + } + }); +}); + +test("prompt runtime leaves non-exact paths and unrelated variables unchanged", async () => { + const toolCall = registerPromptRuntimeHandlers().get("tool_call"); + assert.ok(toolCall); + + await withResultPath("/tmp/subagents/abc/result.log", async () => { + for (const candidate of [ + "prefix/$PI_SUBAGENT_RESULT_PATH", + "${PI_SUBAGENT_RESULT_PATH}/suffix", + "$HOME/result.log", + "/tmp/ordinary.log", + ]) { + const event = { toolName: "write", input: { path: candidate } }; + await toolCall(event); + assert.equal(event.input.path, candidate); + } + }); +}); + +test("prompt runtime leaves aliases unchanged when result env is missing", async () => { + const toolCall = registerPromptRuntimeHandlers().get("tool_call"); + assert.ok(toolCall); + + await withResultPath(undefined, async () => { + const event = { + toolName: "read", + input: { path: "$PI_SUBAGENT_RESULT_PATH" }, + }; + await toolCall(event); + assert.equal(event.input.path, "$PI_SUBAGENT_RESULT_PATH"); + }); +}); + +test("prompt runtime never rewrites bash commands", async () => { + const toolCall = registerPromptRuntimeHandlers().get("tool_call"); + assert.ok(toolCall); + + await withResultPath("/tmp/subagents/abc/result.log", async () => { + const event = { + toolName: "bash", + input: { + path: "$PI_SUBAGENT_RESULT_PATH", + command: 'printf done > "$PI_SUBAGENT_RESULT_PATH"', + }, + }; + await toolCall(event); + assert.equal(event.input.path, "$PI_SUBAGENT_RESULT_PATH"); + assert.equal( + event.input.command, + 'printf done > "$PI_SUBAGENT_RESULT_PATH"', + ); + }); +}); + test("child pi args do not restrict tools skills extensions or MCP", () => { const built = buildPiArgs({ baseArgs: [], @@ -547,39 +659,40 @@ test("subagent-written result file content is preserved, not overwritten", async }); // Test 8: Prompt injection includes result path (Requirement 4) -test("subagent system prompt includes result file path when env var set", () => { - const prompt = "Original system prompt."; - const resultPath = "/tmp/subagents/abc/result.log"; - process.env[SUBAGENT_RESULT_PATH_ENV] = resultPath; - - // Simulate the handler's logic (mirrors registerSubagentPromptRuntime): - const RESULT_PATH_MARKER = "Your result file:"; - let rewritten = rewriteSubagentPrompt(prompt); - if (resultPath && !rewritten.includes(RESULT_PATH_MARKER)) { - rewritten = `${rewritten}\n\nYour result file: ${resultPath}\nYou may write your final output to this file at any time using any tool (e.g., write, bash). If you leave the file empty, your final assistant message will be automatically saved there on exit. The environment variable "$PI_SUBAGENT_RESULT_PATH" is aliased to ${resultPath}; you can pipe your answer there. Particularly for very large outputs, or for programmatic outputs, use tools to write the result directly to "$PI_SUBAGENT_RESULT_PATH".`; - } +test("prompt runtime explains literal file paths shell aliases and fallback", async () => { + const beforeAgentStart = + registerPromptRuntimeHandlers().get("before_agent_start"); + assert.ok(beforeAgentStart); - assert.ok(rewritten.includes(resultPath)); - assert.ok(rewritten.includes("Your result file:")); - assert.ok(rewritten.includes("write")); - assert.ok(rewritten.includes("automatically saved")); - assert.ok(rewritten.includes('"$PI_SUBAGENT_RESULT_PATH"')); - assert.ok(rewritten.includes("you can pipe your answer there")); - assert.ok(rewritten.includes("very large outputs")); - assert.ok(rewritten.includes("programmatic outputs")); - - // Idempotency: second injection must not append again - let rewrittenAgain = rewritten; - if (resultPath && !rewrittenAgain.includes(RESULT_PATH_MARKER)) { - rewrittenAgain = `${rewrittenAgain}\n\nYour result file: ${resultPath}\nYou may write...`; - } - assert.equal( - rewrittenAgain, - rewritten, - "prompt injection must be idempotent", - ); + await withResultPath("/tmp/subagents/abc/result.log", async () => { + const first = (await beforeAgentStart({ + systemPrompt: "Original system prompt.", + })) as { systemPrompt: string }; + const rewritten = first.systemPrompt; + + assert.ok( + rewritten.includes( + "Your result file: /tmp/subagents/abc/result.log", + ), + ); + assert.match(rewritten, /resolved absolute (result )?path/i); + assert.match(rewritten, /file tools.*write.*edit.*read/i); + assert.match(rewritten, /literal absolute path/i); + assert.match(rewritten, /do not expand (shell )?environment variables/i); + assert.match(rewritten, /PI_SUBAGENT_RESULT_PATH.*same path/i); + assert.match(rewritten, /only inside (bash|shell)/i); + assert.match(rewritten, /automatically saved there on exit/i); + assert.doesNotMatch(rewritten, /using any tool/i); + assert.doesNotMatch( + rewritten, + /write directly to ["']?\$PI_SUBAGENT_RESULT_PATH/i, + ); - delete process.env[SUBAGENT_RESULT_PATH_ENV]; + const second = (await beforeAgentStart({ + systemPrompt: rewritten, + })) as { systemPrompt?: string } | undefined; + assert.equal(second, undefined, "prompt injection must be idempotent"); + }); }); // Test 9: pi-args passes resultPath env var (Requirement 4)