From a541b3d98bff5e5b94242f52b5a8d2082e9f12c5 Mon Sep 17 00:00:00 2001 From: Amos Date: Mon, 6 Jul 2026 23:47:01 +0300 Subject: [PATCH] fixed session resume launching with global extensions. resumed subagents launch with subagent-specific tools. also fixed freezing upon resume --- README.md | 6 + agents/scout.md | 2 +- pi-extension/subagents/index.ts | 207 +++++++++++++++++++++++------- pi-extension/subagents/session.ts | 64 +++++++++ test/test.ts | 102 +++++++++++++++ 5 files changed, 333 insertions(+), 48 deletions(-) diff --git a/README.md b/README.md index bd02c98..b3690c9 100644 --- a/README.md +++ b/README.md @@ -182,8 +182,12 @@ subagent_message({ sessionId: "019f05b2-f1c3", message: "Now write the tests too Resuming is fire-and-forget async: the relaunched session always runs its follow-up task autonomously and its result is delivered later as a steer message, exactly like a fresh `subagent` spawn. (There is no per-call behavior knob — resume is always autonomous, matching this result-delivery model.) +**Resume replays the original sandbox.** A resumed subagent is *not* relaunched as a bare `pi --session`. At spawn time the extension snapshots the agent's fully-resolved loadout — the `--tools` allowlist, backing extensions, model + thinking level, identity/system prompt, spawn whitelist, cwd, and config dir — into a `.loadout.json` sidecar next to the session file. Resume reads that snapshot and reconstructs the exact same `--no-extensions` + `--tools` sandbox, so the reincarnated process gets the same restricted toolset it originally ran with instead of pi's default (all global extensions + the full toolset). The snapshot is stored, not re-derived from the agent `.md` by name, so resume stays faithful even if the agent definition was later edited, moved, or deleted. + > **Guard:** a `sessionId` that maps to a still-running subagent is rejected — resuming would launch a second process mutating the same session file. Steer it by `name` instead. There is no hard-abort tool; to forcibly stop a subagent, use its pane directly. +> **Refusal:** if a session has no loadout snapshot (it predates this feature, or the sidecar was deleted), resume is **refused** with a clear error rather than relaunching unrestricted. Re-run the task as a fresh subagent instead. + **`subagent_message` parameters:** - `name` — exact display name of a running subagent to steer (mutually exclusive with `sessionId`) - `sessionId` — id (or id prefix) of a finished session to resume (mutually exclusive with `name`) @@ -339,6 +343,8 @@ name: worker Access is **whitelist-only**. Every sub-agent process is launched with `--no-extensions` (extension discovery disabled) and `--tools `. Only the tools named in the agent's `tools` frontmatter are exposed, and only the extensions that register those tools are loaded back in via explicit `--extension` flags. There is no global default toolset and no deny-list to maintain — an agent gets exactly what it asks for. +This restriction survives **resume**: the resolved sandbox is snapshotted to a `.loadout.json` sidecar at spawn and replayed verbatim when the session is later resumed via `subagent_message({ sessionId })`. A session with no snapshot is refused rather than relaunched unrestricted, so there is no path by which a resume silently escalates a sub-agent back to the full toolset + all global extensions. + ### Spawns must name a known agent The agent whitelist is enforced at **every** depth, not just for restricted sub-agents: diff --git a/agents/scout.md b/agents/scout.md index 766441c..2a1be1c 100644 --- a/agents/scout.md +++ b/agents/scout.md @@ -2,7 +2,7 @@ name: scout description: Fast codebase recon — explores files, finds patterns, maps architecture tools: read, grep, find, ls -model: openrouter/z-ai/glm-5.2 +model: openrouter/deepseek/deepseek-v4-pro thinking: low system-prompt: append auto-exit: true diff --git a/pi-extension/subagents/index.ts b/pi-extension/subagents/index.ts index 64d65c8..b36bc33 100644 --- a/pi-extension/subagents/index.ts +++ b/pi-extension/subagents/index.ts @@ -31,10 +31,13 @@ import { findLastAssistantMessage, getNewEntries, getSessionId, + readSubagentLoadout, resolveSessionFileById, seedSubagentSessionFile, summarizeSessionStats, + writeSubagentLoadout, type SessionStats, + type SubagentLoadout, } from "./session.ts"; import { type StatusSnapshot, @@ -164,6 +167,40 @@ function getAgentConfigDir(): string { return process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent"); } +// ── Runtime tool-extension registration ───────────────────────────────────── +// `getToolExtensionPath` otherwise only knows a closed set of tool names. Other +// pi extensions that bundle a tool for subagents (e.g. a project-local +// extension exposing a bespoke tool) register its name → extension-file path +// here at load/session_start time so a child process can be launched with +// `--no-extensions` + an explicit `-e ` for it. Mirrors the legacy +// `subagents` extension's `registerToolExtension` hook. +const EXTRA_TOOL_EXTENSIONS = new Map(); + +/** Register (or re-register) a custom tool's backing extension file. */ +export function registerToolExtension(name: string, extensionPath: string): void { + if (BUILTIN_TOOLS.has(name)) { + throw new Error(`Cannot register custom tool "${name}": shadows a built-in pi tool`); + } + if ((SPAWNING_TOOLS as readonly string[]).includes(name)) { + throw new Error(`Cannot register custom tool "${name}": shadows a spawning tool`); + } + const existing = EXTRA_TOOL_EXTENSIONS.get(name); + if (existing === extensionPath) return; // idempotent / reload-safe + if (existing !== undefined) { + throw new Error( + `Tool extension already registered for "${name}": ${existing} (refusing to overwrite with ${extensionPath})`, + ); + } + EXTRA_TOOL_EXTENSIONS.set(name, extensionPath); +} + +// Expose registration on a process-global so project-local extensions loaded +// via jiti (separate module instances) can reach this shared map. Set at module +// load so it's available before any `session_start` listener runs. +(globalThis as any).__pi_interactive_subagents = { + registerToolExtension, +}; + /** * Map a custom (non-built-in) tool name to the pi-extension file that * registers it. Used to build the child's `--extension` whitelist after @@ -185,7 +222,12 @@ function getToolExtensionPath(tool: string): string | undefined { google_image_search: join(extBase, "google-image-search", "index.ts"), safe_bash: join(SUBAGENTS_DIR, "tools", "safe-bash.ts"), }; - return map[tool]; + // Prefer the built-in path, but fall back to a runtime-registered extension + // when that path no longer exists on disk (e.g. a built-in tool extension + // was disabled/removed but a project-local extension re-registered it). + const builtin = map[tool]; + if (builtin && existsSync(builtin)) return builtin; + return EXTRA_TOOL_EXTENSIONS.get(tool); } /** @@ -777,6 +819,60 @@ function buildSubagentToolAllowlist( return [...allow].join(","); } +/** + * Apply a loadout snapshot's sandbox to a pi command's `parts` array: model, + * identity (system prompt), and the default-deny tool/extension restriction + * (`--no-extensions` + `--tools` + one `-e` per tool-backing extension). + * + * This is the single source of truth for reconstructing a subagent's sandbox, + * used both by the initial `launchSubagent` and by the `subagent_message` + * resume path so the two can never drift. Env vars (PI_SUBAGENT_AGENT / + * PI_SUBAGENT_ALLOWED / PI_CODING_AGENT_DIR) and cwd are the caller's + * responsibility since they differ slightly between launch and resume. + */ +function applySandboxToParts( + parts: string[], + loadout: SubagentLoadout, + opts: { artifactDir: string; name: string }, +): void { + if (loadout.model) { + const model = loadout.thinking ? `${loadout.model}:${loadout.thinking}` : loadout.model; + parts.push("--model", shellEscape(model)); + } + + if (loadout.identity) { + const flag = loadout.systemPromptMode === "replace" ? "--system-prompt" : "--append-system-prompt"; + const spTimestamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19); + const spSafeName = opts.name + .toLowerCase() + .replace(/[^a-z0-9\s-]/g, "") + .replace(/\s+/g, "-") + .replace(/-+/g, "-") + .replace(/^-|-$/g, ""); + const spPath = join(opts.artifactDir, `context/${spSafeName || "subagent"}-sysprompt-${spTimestamp}.md`); + mkdirSync(dirname(spPath), { recursive: true }); + writeFileSync(spPath, loadout.identity, "utf8"); + parts.push(flag, shellEscape(spPath)); + } + + // Default-deny: disable global extension discovery and re-enable only the + // extensions backing the whitelisted tools. A null allowlist means the spawn + // was intentionally unrestricted (e.g. a fork clone) and is replayed as-is. + if (loadout.toolAllowlist) { + parts.push("--no-extensions"); + parts.push("--tools", shellEscape(loadout.toolAllowlist)); + + const extPaths = new Set(); + for (const tool of loadout.toolAllowlist.split(",")) { + const extPath = getToolExtensionPath(tool); + if (extPath && existsSync(extPath)) extPaths.add(extPath); + } + for (const extPath of extPaths) { + parts.push("-e", shellEscape(extPath)); + } + } +} + function buildPiPromptArgs(params: { effectiveSkills?: string; taskDelivery: "direct" | "artifact"; @@ -1025,6 +1121,7 @@ export const __test__ = { resolveLaunchBehavior, resolveEffectiveInteractive, buildSubagentToolAllowlist, + applySandboxToParts, buildPiPromptArgs, formatWidgetRightLabel, observeRunningSubagent, @@ -1216,59 +1313,46 @@ async function launchSubagent( const subagentDonePath = join(SUBAGENTS_DIR, "subagent-done.ts"); parts.push("-e", shellEscape(subagentDonePath)); - if (effectiveModel) { - const model = effectiveThinking ? `${effectiveModel}:${effectiveThinking}` : effectiveModel; - parts.push("--model", shellEscape(model)); - } - - // Pass agent body as system prompt via file to avoid shell escaping issues - // with multiline content. Pi's --append-system-prompt and --system-prompt - // auto-detect file paths and read their contents. - if (identityInSystemPrompt && identity) { - const flag = systemPromptMode === "replace" ? "--system-prompt" : "--append-system-prompt"; - const spTimestamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19); - const spSafeName = params.name - .toLowerCase() - .replace(/[^a-z0-9\s-]/g, "") - .replace(/\s+/g, "-") - .replace(/-+/g, "-") - .replace(/^-|-$/g, ""); - const syspromptPath = join(artifactDir, `context/${spSafeName || "subagent"}-sysprompt-${spTimestamp}.md`); - mkdirSync(dirname(syspromptPath), { recursive: true }); - writeFileSync(syspromptPath, identity, "utf8"); - parts.push(flag, shellEscape(syspromptPath)); - } + // Resolve the config dir the child sees: a target-local .pi/agent/ wins, + // else the propagated global dir. Captured once so the launch env and the + // resume snapshot agree. + const resolvedAgentDir = + localAgentDir && existsSync(localAgentDir) + ? localAgentDir + : process.env.PI_CODING_AGENT_DIR ?? null; // Default-deny model: when an agent restricts its tools (or is granted the // spawning toolset), we disable global extension discovery and re-enable only // the extensions backing the whitelisted tools. Bare/fork spawns with no tool // restriction keep their full default toolset and all global extensions. const toolAllowlist = buildSubagentToolAllowlist(effectiveTools, { grantSpawning }); - if (toolAllowlist) { - parts.push("--no-extensions"); - parts.push("--tools", shellEscape(toolAllowlist)); - const extPaths = new Set(); - for (const tool of toolAllowlist.split(",")) { - const extPath = getToolExtensionPath(tool); - // Skip extensions that can't be resolved on disk so a missing optional - // extension never breaks the launch (the tool simply won't be available). - if (extPath && existsSync(extPath)) extPaths.add(extPath); - } - for (const extPath of extPaths) { - parts.push("-e", shellEscape(extPath)); - } - } + // Snapshot the fully-resolved sandbox beside the session file so a later + // `subagent_message({ sessionId })` resume can replay the exact same + // restriction instead of relaunching pi with all global extensions + tools. + const loadout: SubagentLoadout = { + agent: params.agent ?? null, + toolAllowlist, + model: effectiveModel ?? null, + thinking: effectiveThinking ?? null, + systemPromptMode: systemPromptMode ?? null, + identity: identityInSystemPrompt ? identity : null, + spawnable: agentDefs?.subagentAgents ?? null, + autoExit: agentDefs?.autoExit ?? false, + cwd: effectiveCwd ?? null, + agentDir: resolvedAgentDir, + }; + writeSubagentLoadout(subagentSessionFile, loadout); + + // Apply model, identity, and the default-deny tool/extension restriction via + // the shared helper (same code path resume uses — they can't drift). + applySandboxToParts(parts, loadout, { artifactDir, name: params.name }); // Build env prefix: subagent identity + config dir propagation + spawn allowlist const envParts: string[] = []; - // If the target cwd has its own .pi/agent/, use that as the config root. - // Otherwise propagate the current/global agent dir. - if (localAgentDir && existsSync(localAgentDir)) { - envParts.push(`PI_CODING_AGENT_DIR=${shellEscape(localAgentDir)}`); - } else if (process.env.PI_CODING_AGENT_DIR) { - envParts.push(`PI_CODING_AGENT_DIR=${shellEscape(process.env.PI_CODING_AGENT_DIR)}`); + if (resolvedAgentDir) { + envParts.push(`PI_CODING_AGENT_DIR=${shellEscape(resolvedAgentDir)}`); } if (grantSpawning && agentDefs?.subagentAgents) { @@ -2020,6 +2104,19 @@ export default function subagentsExtension(pi: ExtensionAPI) { } } + // Reconstruct the sandbox from the snapshot written at spawn time. + // Without it we cannot safely resume: relaunching bare would load every + // global extension + the full toolset. Refuse rather than escalate. + const loadout = readSubagentLoadout(sessionPath); + if (!loadout) { + const err = + `Cannot safely resume "${requestedId}": no sandbox snapshot found for this session ` + + `(it predates sandboxed resume, or its .loadout.json sidecar was removed). ` + + `Resuming would relaunch with all global extensions and the full toolset, so this is refused. ` + + `Re-run the task as a fresh subagent instead.`; + return { content: [{ type: "text" as const, text: err }], details: { error: err } }; + } + const resumedSessionId = getSessionId(sessionPath) ?? requestedId; // Record entry count before resuming so we can extract new messages @@ -2040,6 +2137,9 @@ export default function subagentsExtension(pi: ExtensionAPI) { const activityFile = getSubagentActivityFile(artifactDir, id); mkdirSync(dirname(activityFile), { recursive: true }); + // Replay the model, identity, and default-deny tool/extension sandbox. + applySandboxToParts(parts, loadout, { artifactDir, name }); + let resumeMsgFile: string | undefined; if (params.message) { const msgTimestamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19); @@ -2058,10 +2158,19 @@ export default function subagentsExtension(pi: ExtensionAPI) { parts.push(shellEscape(`@${resumeMsgFile}`)); } - // Build env prefix — propagate PI_CODING_AGENT_DIR for config isolation + // Build env prefix — replay the snapshot's config dir + spawn whitelist + // so the resumed process resolves the same agents/extensions and keeps + // the same nested-spawn restriction it originally ran with. const resumeEnvParts: string[] = []; - if (process.env.PI_CODING_AGENT_DIR) { - resumeEnvParts.push(`PI_CODING_AGENT_DIR=${shellEscape(process.env.PI_CODING_AGENT_DIR)}`); + const resumeAgentDir = loadout.agentDir ?? process.env.PI_CODING_AGENT_DIR ?? null; + if (resumeAgentDir) { + resumeEnvParts.push(`PI_CODING_AGENT_DIR=${shellEscape(resumeAgentDir)}`); + } + if (loadout.spawnable && loadout.spawnable.length > 0) { + resumeEnvParts.push(`PI_SUBAGENT_ALLOWED=${shellEscape(loadout.spawnable.join(","))}`); + } + if (loadout.agent) { + resumeEnvParts.push(`PI_SUBAGENT_AGENT=${shellEscape(loadout.agent)}`); } resumeEnvParts.push(`PI_SUBAGENT_NAME=${shellEscape(name)}`); resumeEnvParts.push(`PI_SUBAGENT_SESSION=${shellEscape(sessionPath)}`); @@ -2072,7 +2181,11 @@ export default function subagentsExtension(pi: ExtensionAPI) { } const resumeEnvPrefix = resumeEnvParts.join(" ") + " "; - const command = `${resumeEnvPrefix}${parts.join(" ")}; echo '__SUBAGENT_DONE_'$?'__'`; + // Resume in the subagent's original cwd so its tools (safe_bash, edits) + // operate where they did before. + const resumeCdPrefix = loadout.cwd ? `cd ${shellEscape(loadout.cwd)} && ` : ""; + + const command = `${resumeCdPrefix}${resumeEnvPrefix}${parts.join(" ")}; echo '__SUBAGENT_DONE_'$?'__'`; const launchScriptFile = join( artifactDir, "subagent-scripts", diff --git a/pi-extension/subagents/session.ts b/pi-extension/subagents/session.ts index d3937f4..25494d5 100644 --- a/pi-extension/subagents/session.ts +++ b/pi-extension/subagents/session.ts @@ -79,6 +79,70 @@ export function seedSubagentSessionFile(params: { writeFileSync(params.childSessionFile, lines.join("\n") + "\n", "utf8"); } +/** + * A snapshot of everything needed to reconstruct a subagent's sandbox when its + * session is later resumed via `subagent_message({ sessionId })`. + * + * Written next to the session file as `.loadout.json` at spawn + * time. Resume replays this exact snapshot so the reincarnated process gets the + * same `--no-extensions` + `--tools` restriction, model, identity, spawn + * whitelist, cwd, and config dir it originally ran with — instead of falling + * back to pi's default (all global extensions + full toolset). Storing the + * resolved loadout (rather than re-deriving from the agent `.md` by name) keeps + * resume faithful even if the agent definition is later edited, moved, or + * deleted. + */ +export interface SubagentLoadout { + /** Agent profile name (for PI_SUBAGENT_AGENT); null for agentless spawns. */ + agent: string | null; + /** The `--tools` allowlist string, or null when the spawn was unrestricted. */ + toolAllowlist: string | null; + /** Model id (without thinking suffix), or null to use the session default. */ + model: string | null; + /** Thinking level appended to the model as `model:level`, or null. */ + thinking: string | null; + /** How the identity text was applied: append/replace, or null. */ + systemPromptMode: "append" | "replace" | null; + /** The system-prompt/identity text, only when it lived in the system prompt. */ + identity: string | null; + /** Agents this subagent was allowed to spawn (for PI_SUBAGENT_ALLOWED). */ + spawnable: string[] | null; + /** Whether the agent auto-exits (informational; resume forces autonomous). */ + autoExit: boolean; + /** Working directory the subagent ran in, or null. */ + cwd: string | null; + /** PI_CODING_AGENT_DIR the subagent resolved config/extensions from, or null. */ + agentDir: string | null; +} + +/** Path of the loadout sidecar written next to a subagent session file. */ +export function loadoutSidecarPath(sessionFile: string): string { + return `${sessionFile}.loadout.json`; +} + +/** Persist a subagent's resolved sandbox loadout beside its session file. */ +export function writeSubagentLoadout(sessionFile: string, loadout: SubagentLoadout): void { + try { + writeFileSync(loadoutSidecarPath(sessionFile), JSON.stringify(loadout), "utf8"); + } catch { + // Best-effort: a missing snapshot only means resume will refuse, never that + // it launches unrestricted. + } +} + +/** Read a subagent's loadout snapshot, or null if absent/unparseable. */ +export function readSubagentLoadout(sessionFile: string): SubagentLoadout | null { + try { + const p = loadoutSidecarPath(sessionFile); + if (!existsSync(p)) return null; + const parsed = JSON.parse(readFileSync(p, "utf8")); + if (!parsed || typeof parsed !== "object") return null; + return parsed as SubagentLoadout; + } catch { + return null; + } +} + function readEntries(sessionFile: string): SessionEntry[] { const raw = readFileSync(sessionFile, "utf8"); return raw diff --git a/test/test.ts b/test/test.ts index f2babae..1f0f670 100644 --- a/test/test.ts +++ b/test/test.ts @@ -11,6 +11,10 @@ import { getLeafId, getNewEntries, getSessionId, + readSubagentLoadout, + writeSubagentLoadout, + loadoutSidecarPath, + type SubagentLoadout, resolveSessionFileById, findLastAssistantMessage, appendBranchSummary, @@ -297,6 +301,44 @@ describe("session.ts", () => { }); }); + describe("subagent loadout snapshot", () => { + const sample: SubagentLoadout = { + agent: "worker", + toolAllowlist: "read,write,edit,safe_bash,web_search,subagent,ask_question", + model: "openrouter/z-ai/glm-5.2", + thinking: "medium", + systemPromptMode: "append", + identity: "You are a worker agent.", + spawnable: ["scout", "researcher"], + autoExit: true, + cwd: "/work/dir", + agentDir: "/home/u/.pi/agent", + }; + + it("writes the sidecar next to the session file", () => { + const sf = join(dir, "s1.jsonl"); + writeSubagentLoadout(sf, sample); + assert.equal(loadoutSidecarPath(sf), sf + ".loadout.json"); + assert.ok(existsSync(sf + ".loadout.json")); + }); + + it("round-trips the full loadout", () => { + const sf = join(dir, "s2.jsonl"); + writeSubagentLoadout(sf, sample); + assert.deepEqual(readSubagentLoadout(sf), sample); + }); + + it("returns null when the sidecar is absent", () => { + assert.equal(readSubagentLoadout(join(dir, "missing.jsonl")), null); + }); + + it("returns null when the sidecar is corrupt", () => { + const sf = join(dir, "s3.jsonl"); + writeFileSync(sf + ".loadout.json", "not json{", "utf8"); + assert.equal(readSubagentLoadout(sf), null); + }); + }); + describe("findLastAssistantMessage", () => { it("finds last assistant text", () => { const entries = [USER_MSG, ASSISTANT_MSG, ASSISTANT_MSG_2] as any[]; @@ -1190,6 +1232,66 @@ describe("subagent discovery", () => { assert.equal(testApi.buildSubagentToolAllowlist(""), null); }); + it("applySandboxToParts replays model, identity, and default-deny tool restriction", () => { + withTempDir((d) => { + const parts: string[] = []; + testApi.applySandboxToParts( + parts, + { + agent: "worker", + toolAllowlist: "read,write,safe_bash", + model: "openrouter/z-ai/glm-5.2", + thinking: "medium", + systemPromptMode: "append", + identity: "You are a worker.", + spawnable: ["scout"], + autoExit: true, + cwd: null, + agentDir: null, + }, + { artifactDir: d, name: "worker" }, + ); + const joined = parts.join(" "); + // Model with thinking suffix. + assert.ok(joined.includes("--model"), "expected --model"); + assert.ok(joined.includes("openrouter/z-ai/glm-5.2:medium"), "expected model:thinking"); + // Identity written to a file and appended. + assert.ok(joined.includes("--append-system-prompt"), "expected --append-system-prompt"); + // Default-deny restriction. + assert.ok(parts.includes("--no-extensions"), "expected --no-extensions"); + const toolsIdx = parts.indexOf("--tools"); + assert.ok(toolsIdx >= 0, "expected --tools"); + // The value is shell-escaped (single-quoted) before joining. + assert.ok( + parts[toolsIdx + 1].includes("read,write,safe_bash"), + "expected the tool allowlist as the --tools value", + ); + }); + }); + + it("applySandboxToParts omits restriction flags when the loadout was unrestricted", () => { + withTempDir((d) => { + const parts: string[] = []; + testApi.applySandboxToParts( + parts, + { + agent: null, + toolAllowlist: null, + model: null, + thinking: null, + systemPromptMode: null, + identity: null, + spawnable: null, + autoExit: false, + cwd: null, + agentDir: null, + }, + { artifactDir: d, name: "fork" }, + ); + assert.deepEqual(parts, []); + }); + }); + it("buildPiPromptArgs inserts separator for artifact-backed launches with skills", () => { assert.deepEqual( testApi.buildPiPromptArgs({ effectiveSkills: "review,lint", taskDelivery: "artifact", taskArg: "@artifact.md" }),