Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion extensions/pi-subagents/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,31 @@ 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;
return `${CHILD_SUBAGENT_SYSTEM_LINE}\n\n${prompt}`;
}

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();
Expand All @@ -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;
Expand Down
177 changes: 145 additions & 32 deletions extensions/pi-subagents/test/unit/minimal-subagents.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<string, (event: any) => 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<void>,
) {
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: [],
Expand Down Expand Up @@ -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)
Expand Down
Loading