From 705c5946586aa5f6d5f9927e76c0629523cfcbcf Mon Sep 17 00:00:00 2001 From: Alex Rosenzweig Date: Wed, 28 Oct 2026 21:10:16 -0500 Subject: [PATCH] Enforce protection checks in the Pi harness plugin Co-authored-by: Codex Signed-off-by: Alex Rosenzweig --- .github/workflows/ci.yml | 9 + src/bundled/pi/host.live.test.ts | 13 + src/bundled/pi/host.mjs | 474 ++++++++++++++++++++++--- src/bundled/pi/index.test.ts | 1 + src/bundled/pi/index.ts | 24 +- src/bundled/pi/pi-testing.ts | 8 +- src/bundled/pi/pi.test.ts | 166 +++++++++ src/bundled/pi/pi.ts | 196 +++++++++- src/bundled/pi/protection.live.test.ts | 200 +++++++++++ src/bundled/pi/runtime.test.ts | 49 +++ src/bundled/pi/runtime.ts | 83 +++-- src/bundled/pi/setup.ts | 23 ++ 12 files changed, 1160 insertions(+), 86 deletions(-) create mode 100644 src/bundled/pi/protection.live.test.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 20634de1aa..e301f1e1a5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,6 +62,15 @@ jobs: with: persist-credentials: false - uses: ./.github/actions/setup + - name: Install pinned Pi test dependencies + run: | + source bin/activate-hermit + cd src/bundled/pi/deps + npm ci --ignore-scripts --no-audit --no-fund + - name: Real Pi tools and recovery + env: + PI_DEPS: ${{ github.workspace }}/src/bundled/pi/deps + run: ./bin/node scripts/ci-test-report.mjs kind=vitest report=test-results/protection/vitest.json evidence=test-results/protection/timing.json title="Protected Pi" -- ./bin/pnpm exec vitest run src/bundled/pi/protection.live.test.ts --reporter=default --reporter=json --outputFile=test-results/protection/vitest.json - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: key: agent-protection diff --git a/src/bundled/pi/host.live.test.ts b/src/bundled/pi/host.live.test.ts index eef00bf5b3..7f899443ad 100644 --- a/src/bundled/pi/host.live.test.ts +++ b/src/bundled/pi/host.live.test.ts @@ -110,6 +110,11 @@ try { out.cleanupFailure = await pi.rpc("wait", { agent, submission: b.submissio finally { openedHarness.commit = commit; } // What a transcript and usage read: the places, and the entries without schemas. out.places = (await pi.rpc("places", { agent })).places; +// Another membership of the same identity has its own durable storage. +const secondary = agent + "@secondary.example%3A7447"; +await pi.rpc("open", { agent: secondary, tools }); +out.secondaryPlaces = (await pi.rpc("places", { agent: secondary })).places; +await pi.rpc("close", { agent: secondary }); out.recorded = Object.keys((await pi.rpc("recorded", { agent, conversation: a.conversation })).models); out.entries = (await pi.rpc("entries", { agent, conversation: a.conversation })).entries.map((entry) => [entry.kind, ...(entry.model ?? []).map((message) => message.role + (message.usage ? "+usage" : ""))]); // A policy enabled after the process started blocks the same live tool path. @@ -562,6 +567,14 @@ describe.skipIf(!DEPS)("the Pi host on Pi durable", () => { text: "Plain answer", }); expect(out.recorded).toEqual(["faux/faux-1"]); + expect(out.secondaryPlaces).toEqual([]); + expect( + ( + await stat( + join(agents, `${"b".repeat(64)}@secondary.example%3A7447.sqlite`), + ) + ).isFile(), + ).toBe(true); expect(out.places).toEqual([ { key: "chan/root", conversation: out.a.conversation, seen: 100 }, ]); diff --git a/src/bundled/pi/host.mjs b/src/bundled/pi/host.mjs index 33ccaa49df..11b047d464 100644 --- a/src/bundled/pi/host.mjs +++ b/src/bundled/pi/host.mjs @@ -9,6 +9,7 @@ import { readFileSync } from "node:fs"; import { chmod, mkdir, readFile, writeFile } from "node:fs/promises"; import { registerHooks, stripTypeScriptTypes } from "node:module"; import { homedir } from "node:os"; +import { realpathSync } from "node:fs"; import { dirname, join, resolve as resolvePath } from "node:path"; import { createInterface } from "node:readline"; import { fileURLToPath, pathToFileURL } from "node:url"; @@ -19,6 +20,7 @@ import { import { clampThinkingLevel, getSupportedThinkingLevels, + getSystemMessageText, } from "@earendil-works/pi-ai"; import { builtinModels } from "@earendil-works/pi-ai/providers/all"; import { @@ -30,6 +32,7 @@ import { Harness, hook, ToolTask, + ProviderDoc, SessionFailed, section, UsageDoc, @@ -41,6 +44,188 @@ import { createToolSearchExtension } from "@earendil-works/pi-durable/tool-searc import { CodingTools } from "@earendil-works/pi-durable/tools"; import lockfile from "proper-lockfile"; +/** Protected Pi uses the same snapshots for review and execution. The model + * barrier is inside Models.applyAuth, before every provider dispatch, including + * compaction and retry/recovery paths. Unsupported payload hooks fail closed. */ +export function protectModels(models, inspect, conversationFor) { + const requests = new Set([ + "stream", + "complete", + "streamSimple", + "completeSimple", + "generateImages", + "classify", + ]); + const deferred = new Set([ + "streamDeferred", + "fetchDeferred", + "cancelDeferred", + ]); + return new Proxy(models, { + get(target, name) { + if (deferred.has(name)) + return () => { + throw new Error( + "Protected Pi does not support deferred provider requests", + ); + }; + const method = Reflect.get(target, name); + if (!requests.has(name)) + return typeof method === "function" ? method.bind(target) : method; + return (model, input, options = {}) => { + if ( + options.onPayload || + options.transformHeaders || + options.samplingParams || + options.deferred || + model.samplingParams || + model.samplingParamsByThinkingLevel + ) + throw new Error( + "Protected Pi cannot review this provider payload override", + ); + const selected = structuredClone(model), + content = structuredClone(input); + const { signal, apiKey, headers, env, ...data } = options; + const settings = structuredClone(data); + return method.call(target, selected, content, { + ...settings, + signal, + apiKey, + headers: structuredClone(headers), + env: structuredClone(env), + transformHeaders: async (assembled) => { + signal?.throwIfAborted(); + const session = await conversationFor(settings.sessionId); + if (!Array.isArray(content.messages)) + throw new Error( + "Protected Pi requires explicit conversation messages", + ); + // Provider settings and tool schemas are not conversation text. + // Keep dynamic instructions and every native message as distinct sources. + for (const [source, value] of [ + ["instructions", content.systemPrompt ?? ""], + ...content.messages.map((message) => [ + "message", + message.role === "system" + ? { role: "system", content: getSystemMessageText(message) } + : message, + ]), + ]) { + await inspect( + { + kind: "context", + session: String(session), + content: { source, value }, + }, + signal, + ); + } + signal?.throwIfAborted(); + return assembled; + }, + }); + }; + }, + }); +} + +/** Hold all native tool output until reviewed, including streaming progress, + * diagnostics and thrown errors. Wrapping execute also covers safe replay. */ +export function protectTool(tool, inspect, preceding = async () => []) { + return { + ...tool, + async execute(original, api, signalled) { + const args = structuredClone(original); + const session = String(api.conversationId); + const signal = signalled.abortSignal; + const conversation = structuredClone(await preceding(api, signalled)); + await inspect( + { + kind: "tool", + session, + content: { title: tool.name, arguments: args, conversation }, + ...(typeof args.command === "string" + ? { command: args.command } + : {}), + }, + signal, + ); + const decoder = new TextDecoder(); + let text = "", + bytes = 0, + details; + const diagnostics = []; + const bound = (value) => { + bytes += Buffer.byteLength( + typeof value === "string" ? value : JSON.stringify(value), + ); + if (bytes > 256 * 1024) + throw new Error("Protected Pi tool output exceeds the review limit"); + }; + let result; + try { + result = await tool.execute( + args, + { + ...api, + outputWindow: undefined, + output(chunk, skipped) { + if (skipped) + throw new Error("Protected Pi cannot review omitted output"); + const value = + typeof chunk === "string" + ? decoder.decode() + chunk + : decoder.decode(chunk, { stream: true }); + bound(value); + text += value; + }, + retainedOutput: () => ({ text, truncated: false }), + diagnostic(value) { + const copy = structuredClone(value); + bound(copy); + diagnostics.push(copy); + }, + async details(value) { + const copy = structuredClone(value); + bound(copy); + details = copy; + }, + }, + signalled, + ); + text += decoder.decode(); + result = { + ...result, + output: result.output ?? (text ? [{ type: "text", text }] : []), + diagnostics: [...diagnostics, ...(result.diagnostics ?? [])], + ...(result.details !== undefined + ? { details: result.details } + : details !== undefined + ? { details } + : {}), + }; + } catch (error) { + if (signal?.aborted) + throw new Error("Protected Pi action was cancelled"); + result = { + isError: true, + output: [ + { + type: "text", + text: error instanceof Error ? error.message : String(error), + }, + ], + }; + } + result = structuredClone(result); + await inspect({ kind: "tool-result", session, content: result }, signal); + signal?.throwIfAborted(); + return result; + }, + }; +} + const context = BACKGROUND_CONTEXT; /** Each Buzz conversation key's Pi conversation, the tools' default * destination for it, and the newest message it has been shown. */ @@ -383,6 +568,7 @@ export function serve({ data, models, packages = { extensions: [], errors: [] }, + protectedAgent, dir = piDir(), input = process.stdin, output = process.stdout, @@ -393,13 +579,31 @@ export function serve({ /** Agents whose storage failed and that nothing has opened again. */ const failed = new Set(); const calls = new Map(); + const reviews = new Map(); let called = 0; const write = (message) => output.write(`${JSON.stringify(message)}\n`); const notify = (method, params) => write({ method, params }); - const call = (method, params) => + const call = (method, params, signal) => new Promise((resolve, reject) => { + signal?.throwIfAborted(); const id = `h${++called}`; - calls.set(id, { resolve, reject }); + const cancel = () => { + calls.delete(id); + write({ method: "security_cancel", params: { id } }); + reject(new Error("Protected Pi action was cancelled")); + }; + const clear = () => signal?.removeEventListener("abort", cancel); + calls.set(id, { + resolve: (value) => { + clear(); + resolve(value); + }, + reject: (error) => { + clear(); + reject(error); + }, + }); + signal?.addEventListener("abort", cancel, { once: true }); write({ id, method, params }); }); @@ -432,6 +636,8 @@ export function serve({ async function open(agent, lost) { if (!AGENT_KEY.test(agent)) throw new Error("Not an agent key"); + if (protectedAgent && agent !== protectedAgent) + throw new Error("This protected host belongs to another agent"); const path = join(data, "agents", `${agent}.sqlite`); // An agent's storage holds its conversations: only its owner may read // it, however an earlier build made it. @@ -461,15 +667,91 @@ export function serve({ await chmod(file, 0o600).catch((error) => { if (error.code !== "ENOENT") throw error; }); + if ( + protectedAgent && + (packages.extensions.length || packages.providers?.length) + ) + throw new Error("Protected Pi does not load user packages"); const registry = createRegistry(); - registry.install(CodingTools); + const inspect = async (action, signal) => { + if ( + process.env[ + { + context: "BUZZ_SECURITY_CONTEXT", + tool: "BUZZ_SECURITY_TOOL", + "tool-result": "BUZZ_SECURITY_TOOL_RESULT", + }[action.kind] + ] === "0" + ) + action = { + kind: action.kind, + session: action.session, + content: null, + }; + const result = await call("security", { agent, action }, signal); + if (result !== true) + throw new Error("Pi action denied by its protection policy"); + }; + const install = (extension) => { + if ( + protectedAgent && + (extension.tasks?.length || + extension.hooks?.length || + extension.wraps?.length) + ) + throw new Error("Unsupported protected Pi extension"); + registry.install( + protectedAgent + ? { + ...extension, + tools: extension.tools?.map((tool) => + protectTool(tool, inspect, async (api, context) => { + const conversation = await harness.conversation( + api.conversationId, + context, + ); + if (!conversation) + throw new Error( + "The tool's preceding conversation is unavailable", + ); + return (await conversation.context(context)).messages; + }), + ), + } + : extension, + ); + }; + const conversations = new Set(); + const sessions = new Map(); + let harness; + const conversationFor = async (sessionId) => { + if (typeof sessionId !== "string" || !sessionId) + throw new Error("Protected model request has no session"); + if (!sessions.has(sessionId)) { + for (const conversation of conversations) { + const document = await harness.snapshot( + ProviderDoc, + conversation, + context, + ); + if (document?.sessionId) + sessions.set(document.sessionId, conversation); + } + } + if (!sessions.has(sessionId)) + throw new Error("Protected model request has an unknown session"); + return sessions.get(sessionId); + }; + install(CodingTools); // Loads the Buzz tools that are deferred when the model asks for them. - registry.install(createToolSearchExtension()); - for (const extension of packages.extensions) registry.install(extension); - const harness = await Harness.open( + install(createToolSearchExtension()); + for (const extension of packages.extensions) install(extension); + harness = await Harness.open( await openNodeSqliteStorage(path), { - models, + models: protectedAgent + ? protectModels(models, inspect, conversationFor) + : models, registry, settings: SETTINGS, env: ({ cwd }) => new NodeExecutionEnv({ cwd: cwd ?? homedir() }), @@ -481,6 +763,8 @@ export function serve({ agent, harness, registry, + install, + conversations, release, tools: "", /** Each conversation's `core` memory as its latest turn read it, @@ -498,8 +782,12 @@ export function serve({ closing: false, }; const places = (await harness.snapshot(Places, context)) ?? {}; - for (const place of Object.values(places)) + for (const place of Object.values(places)) { + conversations.add(place.conversation); opened.contexts.set(place.conversation, place.context); + } + for (const submission of (await harness.inspect(context)).submissions) + conversations.add(submission.conversationId); return opened; } catch (error) { await release().catch(() => undefined); @@ -515,7 +803,14 @@ export function serve({ // leaves its runs to it. Opening again says the agent is open there. const opened = open(agent, () => { forget(); - void opened.then(shut, () => undefined); + void opened + .then(shut, () => undefined) + .catch((error) => + console.error( + "Pi could not close after losing storage custody", + error, + ), + ); }); started = opened; failed.delete(agent); @@ -532,7 +827,9 @@ export function serve({ console.error("Pi's storage failed for an agent", end.error); if (agents.get(agent) === opened) failed.add(agent); forget(); - void shut(ready); + void shut(ready).catch((error) => + console.error("Pi could not close an agent", error), + ); }), forget, ); @@ -573,30 +870,32 @@ export function serve({ { tag: false }, ), ], - hooks: [ - hook(ToolTask, { - beforeTool: (input, api, context) => - awaitWithContext( - call("hook", { - agent: opened.agent, - conversation: api.conversationId, - event: "beforeTool", - input, - }), - context, - ).then((value) => value ?? undefined), - afterTool: (input, result, api, context) => - awaitWithContext( - call("hook", { - agent: opened.agent, - conversation: api.conversationId, - event: "afterTool", - input: { call: input, result }, - }), - context, - ).then((value) => value ?? undefined), - }), - ], + hooks: protectedAgent + ? [] + : [ + hook(ToolTask, { + beforeTool: (input, api, context) => + awaitWithContext( + call("hook", { + agent: opened.agent, + conversation: api.conversationId, + event: "beforeTool", + input, + }), + context, + ).then((value) => value ?? undefined), + afterTool: (input, result, api, context) => + awaitWithContext( + call("hook", { + agent: opened.agent, + conversation: api.conversationId, + event: "afterTool", + input: { call: input, result }, + }), + context, + ).then((value) => value ?? undefined), + }), + ], tools: tools.map((tool) => defineTool({ name: `${PREFIX}${tool.name}`, @@ -773,6 +1072,75 @@ export function serve({ } const methods = { + async review(turn) { + if (protectedAgent) + throw new Error("Security inference runs outside requesting workers"); + if ( + typeof turn.id !== "string" || + reviews.has(turn.id) || + reviews.size >= 16 || + !turn.prompt || + typeof turn.prompt.system !== "string" || + typeof turn.prompt.input !== "string" || + Buffer.byteLength(JSON.stringify(turn.prompt)) > 768 * 1024 + ) + throw new Error("Invalid security review request"); + const controller = new AbortController(); + const signal = AbortSignal.any([ + controller.signal, + AbortSignal.timeout(TEST_TIMEOUT_MS), + ]); + reviews.set(turn.id, controller); + try { + const { found, model } = await choose(turn); + signal.throwIfAborted(); + const reply = await models.completeSimple( + found, + { + systemPrompt: turn.prompt.system, + messages: [ + { + role: "user", + content: turn.prompt.input, + timestamp: Date.now(), + }, + ], + }, + { maxTokens: 1024, maxRetries: 0, signal }, + ); + signal.throwIfAborted(); + if (reply.stopReason === "error" || reply.stopReason === "aborted") + throw new Error( + reply.errorMessage || "Security model did not answer", + ); + return { + text: reply.content + .filter((part) => part.type === "text") + .map((part) => part.text) + .join(""), + model: { provider: model.provider, id: model.modelId }, + }; + } finally { + reviews.delete(turn.id); + } + }, + "cancel-review"({ id }) { + reviews.get(id)?.abort(); + }, + runtime() { + if (protectedAgent) + throw new Error("Protected hosts cannot bootstrap another host"); + const file = fileURLToPath(import.meta.url), + node = realpathSync(process.execPath); + return { + file, + node, + codeRoot: dirname(dirname(file)), + nodeRoot: dirname(dirname(node)), + data, + auth: dir, + }; + }, async providers() { const providers = []; for (const provider of models.getProviders()) { @@ -821,6 +1189,10 @@ export function serve({ }, /** Asks the model for one short reply. */ async test(turn) { + if (protectedAgent) + throw new Error( + "Test providers in Pi setup, outside the protected worker", + ); const { found, model } = await choose(turn); const reply = await models.completeSimple( found, @@ -847,7 +1219,7 @@ export function serve({ const opened = await opening(agent); const listed = JSON.stringify([tools, hint]); if (opened.tools !== listed) { - opened.registry.install(buzz(opened, tools, hint)); + opened.install(buzz(opened, tools, hint)); opened.tools = listed; } }, @@ -886,6 +1258,13 @@ export function serve({ ...(page.next ? { next: page.next } : {}), }; }, + /** The native conversation snapshot used for a final publication review. */ + async protectionContext({ agent, conversation }) { + const { harness } = await agentOf(agent); + const found = await harness.conversation(conversation, context); + if (!found) throw new Error("The protected conversation is unavailable"); + return structuredClone((await found.context(context)).messages); + }, /** The usage Pi recorded for a conversation, by `/`: * its answers', and the requests that left no answer, such as * compaction's. */ @@ -934,6 +1313,7 @@ export function serve({ }; }, context); } + opened.conversations.add(conversation.id); opened.contexts.set(conversation.id, tools); const { busy } = await watch(opened, conversation.id); if (reply) @@ -978,6 +1358,7 @@ export function serve({ { ownership: { kind: "ownerless" } }, context, ); + opened.conversations.add(conversation.id); return submit(opened, conversation, turn, false); }, /** The requests of runs Pi has not finished, with where each one's tools @@ -1091,21 +1472,23 @@ export function serve({ * opens again. */ async close({ agent }) { const opening = agents.get(agent); + await shut(await opening?.catch(() => undefined)); agents.delete(agent); failed.delete(agent); - await shut(await opening?.catch(() => undefined)); }, }; async function shut(opened) { - if (!opened || opened.closing) return; - opened.closing = true; - for (const { stream } of opened.watched.values()) - await stream.stop().catch(() => undefined); - await opened.harness.close(context).catch((error) => { - console.error("Pi could not close an agent", error); - }); - await opened.release().catch(() => undefined); + if (!opened) return; + if (!opened.closing) + opened.closing = (async () => { + for (const { stream } of opened.watched.values()) await stream.stop(); + await opened.harness.close(context); + // A compromised lock may already be released. Native execution has + // ended above; releasing ownership is cleanup, not its acknowledgement. + await opened.release().catch(() => undefined); + })(); + await opened.closing; } async function handle({ id, method, params }) { @@ -1123,7 +1506,8 @@ export function serve({ let stopping; const stop = () => { - stopping ??= Promise.all( + for (const review of reviews.values()) review.abort(); + stopping ??= Promise.allSettled( [...agents.values()].map((opening) => opening.then(shut, () => undefined), ), diff --git a/src/bundled/pi/index.test.ts b/src/bundled/pi/index.test.ts index 461e976028..f4b283bead 100644 --- a/src/bundled/pi/index.test.ts +++ b/src/bundled/pi/index.test.ts @@ -8,6 +8,7 @@ vi.mock("@tauri-apps/api/core", () => ({ isTauri: () => true })); it("shares a Pi agent's provider and model as the pi runtime and imports them back", () => { let type: AgentType | undefined; apply({ + agentProtection: { registerModels: vi.fn() }, pluginOwner: { id: "buzz.pi", revision: "r1" }, host: { request: vi.fn() }, agents: { diff --git a/src/bundled/pi/index.ts b/src/bundled/pi/index.ts index b81c930a7e..2af038923a 100644 --- a/src/bundled/pi/index.ts +++ b/src/bundled/pi/index.ts @@ -16,7 +16,7 @@ export const TYPE_ID = "pi"; * which includes the first install of Pi's packages. */ const HANDOVER_MS = 2 * 60_000; -export const inject = ["agents", "host"]; +export const inject = ["agents", "host", "agentProtection"]; export const apply: PluginModule["apply"] = (ctx) => { // Only the desktop app can start the Pi host; no unusable type elsewhere. if (!isTauri()) return; @@ -41,6 +41,26 @@ export const apply: PluginModule["apply"] = (ctx) => { ); }, ); + ctx.agentProtection.registerModels({ + label: "Pi", + models: async () => { + const { providers } = await runtime.setup.providers(); + const choices = await Promise.all( + providers + .filter((provider) => provider.configured) + .map(async (provider) => + (await runtime.setup.models(provider.id)).map((model) => ({ + provider: provider.id, + id: model.id, + label: `${provider.name} ยท ${model.label}`, + })), + ), + ); + return choices.flat(); + }, + complete: (model, prompt, signal) => + runtime.setup.review(model, prompt, signal), + }); const { ActivityTab, Create, CreateModel, SettingsTab } = createTabs({ setup: runtime.setup, runtime, @@ -62,6 +82,8 @@ export const apply: PluginModule["apply"] = (ctx) => { agents.register({ id: TYPE_ID, hooks: { events: ["beforeTool", "afterTool"] }, + retire: (pubkey) => runtime.retire(pubkey), + protection: true, title: "Pi", icon: PiLogoIcon, description: diff --git a/src/bundled/pi/pi-testing.ts b/src/bundled/pi/pi-testing.ts index e816017657..f1616e8a9a 100644 --- a/src/bundled/pi/pi-testing.ts +++ b/src/bundled/pi/pi-testing.ts @@ -29,7 +29,7 @@ export type FakeProcess = { /** Calls one of the plugin's tools and resolves with its answer. */ tool( params: object, - method?: "tool" | "prompt", + method?: "tool" | "prompt" | "security", ): Promise<{ result?: unknown; error?: string }>; /** Exits, saying `stderr` last. */ exit(code: number | null, stderr?: string): void; @@ -37,7 +37,7 @@ export type FakeProcess = { export function fakeHost( handlers: Readonly> = {}, - options: Readonly<{ ready?: boolean }> = {}, + options: Readonly<{ ready?: boolean; close?: boolean }> = {}, ) { const processes: FakeProcess[] = []; const spawn = async ( @@ -102,7 +102,8 @@ export function fakeHost( buffer = buffer.slice(newline + 1); newline = buffer.indexOf("\n"); // The first line is the launcher's: what it installs and runs. - if (!process.launched) process.launched = message; + if (!process.launched && !spawnOptions.protection) + process.launched = message; else if (message.method) void handle(message); else { answers.get(message.id)?.(message); @@ -116,6 +117,7 @@ export function fakeHost( write, end: async () => { process.ended = true; + if (options.close) process.exit(0); }, kill: async () => { process.killed = true; diff --git a/src/bundled/pi/pi.test.ts b/src/bundled/pi/pi.test.ts index db8fd00181..229fba11dc 100644 --- a/src/bundled/pi/pi.test.ts +++ b/src/bundled/pi/pi.test.ts @@ -1,6 +1,7 @@ import { expect, it, vi } from "vitest"; import { AGENT_KEY } from "../../features/agents/membership"; import hostSource from "./host.mjs?raw"; +import type { AgentProtection } from "../../features/agent-protection/service"; import { Pi, PiStopped } from "./pi"; import { deferred, fakeHost } from "./pi-testing"; @@ -162,3 +163,168 @@ it("accepts in the host the membership keys the app names agents with", () => { const copy = /^const AGENT_KEY = (\/.*\/);$/m.exec(hostSource)?.[1]; expect(copy).toBe(`/${AGENT_KEY.source}/`); }); + +it.each(["dispose", "exit", "retire", "policy"])( + "retires completed protection checks with their Pi connection (%s)", + async (retirement) => { + const agent = "a".repeat(64); + const check = vi.fn().mockResolvedValue(); + let gate: AgentProtection = { + key: "a:1", + stateKey: "a:1", + check, + assertAvailable() {}, + processProtection: { readOnlyPaths: [], hiddenPaths: [] }, + spawn: async () => { + throw new Error("Unused protected runtime"); + }, + }; + const host = fakeHost( + { + runtime: () => ({ + file: "/runtime/deps/host.mjs", + node: "/runtime/node/bin/node", + codeRoot: "/runtime/deps", + nodeRoot: "/runtime/node", + data: "/runtime/state", + auth: "/runtime/auth", + }), + }, + { close: true }, + ); + const pi = new Pi( + () => host.spawn, + handlers(), + () => gate, + ); + const action = { + kind: "context", + session: "native-session", + content: "reviewed input", + }; + try { + await pi.request("open", { agent }); + const worker = host.processes.find( + (process) => process.options.protection, + ); + if (!worker) throw new Error("Missing protected worker"); + await expect( + worker.tool({ agent, action }, "security"), + ).resolves.toMatchObject({ result: true }); + const signal = check.mock.calls[0]?.[1]; + expect(signal?.aborted).toBe(false); + if (retirement === "dispose") await pi.dispose(); + else { + if (retirement === "exit") worker.exit(1); + else if (retirement === "retire") await pi.retire(agent); + else gate = { ...gate, key: "a:2", stateKey: "a:2" }; + if (retirement === "exit") + await vi.waitFor(() => expect(signal?.aborted).toBe(true)); + await pi.request("open", { agent }); + const next = host.processes + .filter((process) => process.options.protection) + .at(-1); + expect(next).not.toBe(worker); + await expect( + next?.tool({ agent, action }, "security"), + ).resolves.toMatchObject({ result: true }); + expect(check.mock.calls[1]?.[1].aborted).toBe(false); + } + expect(signal?.aborted).toBe(true); + } finally { + await pi.dispose(); + } + }, +); + +it("isolates protected agents, contains their runtime code and rejects crossed native checks", async () => { + const entered = deferred(); + let waiting: AbortSignal | undefined; + const gates = new Map(); + const make = (key: string): AgentProtection => ({ + key, + stateKey: key, + assertAvailable: vi.fn(), + processProtection: { + readOnlyPaths: ["/security/provider"], + hiddenPaths: ["/security/approvals"], + }, + spawn: async () => { + throw new Error("Unused protected runtime"); + }, + check: vi.fn(async (_action, signal) => { + waiting = signal; + entered.resolve(); + await new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(new Error("cancelled")), { + once: true, + }); + }); + }), + }); + const a = "a".repeat(64), + b = "b".repeat(64); + gates.set(a, make("a:1")); + gates.set(b, make("b:1")); + const host = fakeHost( + { + runtime: () => ({ + file: "/runtime/deps/version/host.mjs", + node: "/runtime/node/bin/node", + codeRoot: "/runtime/deps", + nodeRoot: "/runtime/node", + data: "/runtime/state", + auth: "/runtime/auth", + }), + }, + { close: true }, + ); + const pi = new Pi( + () => host.spawn, + handlers(), + (agent) => gates.get(agent), + ); + try { + await Promise.all([ + pi.request("open", { agent: a }), + pi.request("open", { agent: b }), + ]); + const protectedHosts = host.processes.filter( + (process) => process.options.protection, + ); + expect(protectedHosts).toHaveLength(2); + const worker = protectedHosts.find((process) => + process.options.args?.at(-1)?.includes(a), + ); + if (!worker) throw new Error("Missing protected worker"); + expect(worker.options.args?.[0]).toContain("--disable-sigusr1"); + expect(worker.options.env).toMatchObject({ + BASH_ENV: null, + ENV: null, + NODE_OPTIONS: null, + NODE_PATH: null, + DYLD_INSERT_LIBRARIES: null, + }); + expect(worker.options.protection).toEqual({ + readOnlyPaths: ["/security/provider", "/runtime/deps", "/runtime/node"], + hiddenPaths: ["/security/approvals"], + }); + const action = { kind: "tool", session: "1", content: { title: "write" } }; + await expect( + worker.tool({ agent: b, action }, "security"), + ).resolves.toMatchObject({ result: false }); + const pending = worker.tool({ agent: a, action }, "security"); + await entered.promise; + worker.send({ method: "security_cancel", params: { id: "h2" } }); + await expect(pending).resolves.toMatchObject({ result: false }); + expect(waiting?.aborted).toBe(true); + gates.set(a, make("a:2")); + await pi.request("open", { agent: a }); + expect(worker.ended).toBe(true); + expect( + host.processes.filter((process) => process.options.protection), + ).toHaveLength(3); + } finally { + await pi.dispose(); + } +}); diff --git a/src/bundled/pi/pi.ts b/src/bundled/pi/pi.ts index 0e9c10178e..7528e30973 100644 --- a/src/bundled/pi/pi.ts +++ b/src/bundled/pi/pi.ts @@ -2,6 +2,8 @@ // something needs it, spoken to over stdio, one JSON message a line. Requests // go both ways: the plugin asks the host to run turns, and the host asks the // plugin to run the Buzz tools, so the process holds no Buzz key. +import type { AgentProtection } from "../../features/agent-protection/service"; +import { protectionEnvironment } from "../../features/agent-protection/native-bridge"; import type { Context } from "../../buzz-mcp/client"; import type { HostProcess, @@ -78,7 +80,22 @@ type Message = { result?: unknown; error?: string; }; +type RuntimeInfo = { + file: string; + node: string; + codeRoot: string; + nodeRoot: string; + data: string; + auth: string; +}; +type ProtectedLaunch = { + agent: string; + gate: AgentProtection; + runtime: () => Promise; +}; type Connection = { + lifetime: AbortController; + checks: Map; process: HostProcess; calls: Map< number, @@ -92,14 +109,61 @@ export class Pi { private connection: Promise | undefined; private requests = 0; private stopped = false; + private hosts = new Map< + string, + Promise<{ key: string; pi: Pi | undefined }> + >(); constructor( private readonly spawn: () => Spawn, private readonly handlers: Handlers, + private readonly protection?: ( + agent: string, + ) => AgentProtection | undefined, + private readonly protectedLaunch?: ProtectedLaunch, ) {} /** Asks the host, starting it first if it is not running. */ async request(method: string, params: object = {}): Promise { + if (this.stopped) throw new Error("Pi has stopped"); + const agent = (params as { agent?: string }).agent; + if (agent && this.protection) { + const gate = this.protection(agent); + gate?.assertAvailable(); + const prior = this.hosts.get(agent); + const next = (async () => { + const previous = await prior?.catch(() => undefined); + if (previous?.key === (gate?.key ?? "")) return previous; + if (previous?.pi) await previous.pi.dispose(); + else if (this.connection) await this.rawRequest("close", { agent }); + const pi = gate + ? new Pi(this.spawn, this.handlers, undefined, { + agent, + gate, + runtime: () => this.rawRequest("runtime"), + }) + : undefined; + return { key: gate?.key ?? "", pi }; + })(); + this.hosts.set(agent, next); + const child = (await next).pi; + if (child) return child.request(method, params); + } + return this.rawRequest(method, params); + } + /** Ends the old worker without resolving a possibly newer policy first. */ + async retire(agent: string) { + const prior = this.hosts.get(agent); + const retired = (async () => { + const previous = await prior?.catch(() => undefined); + if (previous?.pi) await previous.pi.dispose(); + else if (this.connection) await this.rawRequest("close", { agent }); + return { key: "retired", pi: undefined }; + })(); + this.hosts.set(agent, retired); + await retired; + } + private async rawRequest(method: string, params: object = {}): Promise { const connection = await this.connect(); if (connection.gone) throw connection.gone; const id = ++this.requests; @@ -120,8 +184,16 @@ export class Pi { /** Stops the host: it closes each agent's storage, then it is killed. */ async dispose() { this.stopped = true; + await Promise.all( + [...this.hosts.values()].map(async (pending) => + (await pending.catch(() => undefined))?.pi?.dispose(), + ), + ); + this.hosts.clear(); const connection = await this.connection?.catch(() => undefined); if (!connection || connection.gone) return; + connection.lifetime.abort(); + for (const check of connection.checks.values()) check.abort(); await connection.process.end().catch(() => undefined); const timer = setTimeout(() => void connection.process.kill(), STOP_MS); await connection.process.exited; @@ -153,11 +225,17 @@ export class Pi { let buffer = ""; let stderr = ""; const connection: Connection = { + lifetime: new AbortController(), + checks: new Map(), calls: new Map(), process: undefined as unknown as HostProcess, }; const handle = (message: Message) => { - if (message.method === "ready") ready(); + if (message.method === "security_cancel") { + connection.checks.get((message.params as { id: string }).id)?.abort(); + } else if (message.method === "security" && message.id !== undefined) { + void this.security(connection, message.id, message.params); + } else if (message.method === "ready") ready(); else if (message.method === "event") this.handlers.event(message.params as PiEvent); else if ( @@ -183,9 +261,69 @@ export class Pi { // its packages go on stdin: together they are past Windows' 32 KB // command line. const windows = /^Win/i.test(globalThis.navigator?.platform ?? ""); + const protectedLaunch = this.protectedLaunch; + protectedLaunch?.gate.assertAvailable(); + const runtime = protectedLaunch && (await protectedLaunch.runtime()); + if ( + runtime && + (!runtime.node.startsWith("/") || + !runtime.file.startsWith("/") || + !runtime.codeRoot.startsWith("/") || + !runtime.nodeRoot.startsWith("/") || + !runtime.data.startsWith("/") || + !runtime.auth.startsWith("/")) + ) + throw new Error("Invalid protected Pi runtime"); + const boot = + runtime && protectedLaunch + ? ` + import {createHash} from 'node:crypto'; + import {join} from 'node:path'; + import {pathToFileURL} from 'node:url'; + import {builtinModels} from '@earendil-works/pi-ai/providers/all'; + const {AuthFile,serve}=await import(pathToFileURL(${JSON.stringify(runtime.file)}).href); + const key=createHash('sha256').update(${JSON.stringify(protectedLaunch.gate.stateKey)}).digest('hex'); + const {stop}=serve({data:join(${JSON.stringify(runtime.data)},'protected-v1',key),dir:${JSON.stringify(runtime.auth)},protectedAgent:${JSON.stringify(protectedLaunch.agent)},models:builtinModels({credentials:new AuthFile(join(${JSON.stringify(runtime.auth)},'auth.json'))})}); + process.on('SIGTERM',stop); + ` + : undefined; const starting = this.spawn()(windows ? "node" : "pi-sh", { - args: windows ? [launch] : [script, "pi.sh", launch], - cwd: "~", + args: + boot && runtime + ? [ + 'exec "$1" --disable-sigusr1 --input-type=module -e "$2"', + "buzz-pi-protected", + runtime.node, + boot, + ] + : windows + ? [launch] + : [script, "pi.sh", launch], + cwd: runtime ? runtime.file.slice(0, runtime.file.lastIndexOf("/")) : "~", + ...(runtime && protectedLaunch + ? { + env: { + ...protectionEnvironment(protectedLaunch.gate), + BASH_ENV: null, + ENV: null, + NODE_OPTIONS: null, + NODE_PATH: null, + DYLD_INSERT_LIBRARIES: null, + DYLD_LIBRARY_PATH: null, + DYLD_FRAMEWORK_PATH: null, + LD_PRELOAD: null, + LD_LIBRARY_PATH: null, + }, + protection: { + readOnlyPaths: [ + ...protectedLaunch.gate.processProtection.readOnlyPaths, + runtime.codeRoot, + runtime.nodeRoot, + ], + hiddenPaths: protectedLaunch.gate.processProtection.hiddenPaths, + }, + } + : {}), onStdout: (data) => { buffer += data; let newline = buffer.indexOf("\n"); @@ -213,13 +351,17 @@ export class Pi { ) : error; }); - void connection.process - .write( - `${JSON.stringify({ host, package: packageJson, lock: packageLock })}\n`, - ) - // A host that never has its source waits for it; it stops instead. - .catch(() => void connection.process.kill().catch(() => undefined)); + if (!protectedLaunch) + void connection.process + .write( + `${JSON.stringify({ host, package: packageJson, lock: packageLock })}\n`, + ) + // A host that never has its source waits for it; it stops instead. + .catch(() => void connection.process.kill().catch(() => undefined)); void connection.process.exited.then((code) => { + connection.lifetime.abort(); + for (const check of connection.checks.values()) check.abort(); + connection.checks.clear(); // Before it is ready, its last words are why it could not start, such // as Node missing. const why = cause(stderr); @@ -239,6 +381,42 @@ export class Pi { return connection; } + private async security( + connection: Connection, + id: number | string, + raw: unknown, + ) { + const controller = new AbortController(); + const signal = AbortSignal.any([ + controller.signal, + connection.lifetime.signal, + ]); + if (connection.checks.has(id)) { + await connection.process.kill(); + return; + } + connection.checks.set(id, controller); + let allowed = false; + try { + const request = raw as { + agent: string; + action: Parameters[0]; + }; + if (!this.protectedLaunch || request.agent !== this.protectedLaunch.agent) + throw new Error("No matching protected agent"); + await this.protectedLaunch.gate.check(request.action, signal); + allowed = !signal.aborted; + } catch { + /* Native checks fail closed. */ + } finally { + connection.checks.delete(id); + } + if (!connection.gone) + await connection.process + .write(`${JSON.stringify({ id, result: allowed })}\n`) + .catch(() => undefined); + } + /** Runs a Buzz tool the host called, and sends it the result. */ private async answer( connection: Connection, diff --git a/src/bundled/pi/protection.live.test.ts b/src/bundled/pi/protection.live.test.ts new file mode 100644 index 0000000000..758890d585 --- /dev/null +++ b/src/bundled/pi/protection.live.test.ts @@ -0,0 +1,200 @@ +// @vitest-environment node +// Real Pi durable, SQLite, built-in tools and faux provider. CI installs exactly +// deps/package-lock.json; no credentials, model service or network is used here. +import { execFile } from "node:child_process"; +import { mkdtemp, rm, symlink, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { promisify } from "node:util"; +import { it, expect } from "vitest"; +import host from "./host.mjs?raw"; +const deps = process.env.PI_DEPS; +const driver = String.raw`import assert from 'node:assert/strict'; +import {readFile,writeFile,stat} from 'node:fs/promises'; +import {join} from 'node:path'; +import {PassThrough} from 'node:stream'; +import {createInterface} from 'node:readline'; +import {createModels} from '@earendil-works/pi-ai/models'; +import {fauxProvider,fauxAssistantMessage,fauxToolCall} from '@earendil-works/pi-ai/providers/faux'; +import {serve,protectTool,protectModels} from './host.mjs'; +const data=process.argv[2], agent='b'.repeat(64), faux=fauxProvider(),models=createModels();models.setProvider(faux.provider); +const result={}; +const deferred=()=>{let resolve;const promise=new Promise(r=>{resolve=r;});return {promise,resolve};}; +const until=async test=>{const end=Date.now()+5000;while(!await test()){if(Date.now()>end)throw new Error('Missing lifecycle boundary');await new Promise(r=>setTimeout(r,5));}}; +const exists=path=>stat(path).then(()=>true,()=>false); +function start(check,protectedAgent=agent) { + const input=new PassThrough(),output=new PassThrough(),pending=new Map(),checks=new Map(),actions=[],nativeHooks=[];let id=0; + createInterface({input:output}).on('line',line=>{ + const m=JSON.parse(line); + if(m.method==='security') { + const stop=new AbortController();checks.set(m.id,stop);actions.push(m.params); + Promise.resolve().then(()=>check(m.params,stop.signal)).then(allowed=>input.write(JSON.stringify({id:m.id,result:allowed===true})+'\n'),()=>input.write(JSON.stringify({id:m.id,result:false})+'\n')).finally(()=>checks.delete(m.id)); + } else if(m.method==='hook'){nativeHooks.push(m.params);input.write(JSON.stringify({id:m.id,result:undefined})+'\n');} + else if(m.method==='security_cancel')checks.get(m.params.id)?.abort(); + else if(m.method==='tool')input.write(JSON.stringify({id:m.id,result:{content:[{type:'text',text:'sent'}]}})+'\n'); + else if(m.id!==undefined){const p=pending.get(m.id);pending.delete(m.id);m.error?p?.reject(new Error(m.error)):p?.resolve(m.result);} + }); + const rpc=(method,params)=>new Promise((resolve,reject)=>{pending.set(++id,{resolve,reject});input.write(JSON.stringify({id,method,params})+'\n');}); + const server=serve({data,models,dir:data,protectedAgent,input,output,exit:()=>{}}); + return {rpc,actions,stop:async()=>{await server.stop();assert.deepEqual(nativeHooks,[]);}}; +} +const turn=(request,text='Do the task',key='chan/root')=>({agent,key,at:100,context:{channel:'chan',root:'root'},text,steer:text,request,provider:'faux',model:'faux-1',instructions:'Test.',cwd:data}); +let inspect=async()=>true; +let pi=start((...args)=>inspect(...args)); +try { + await pi.rpc('open',{agent,tools:[]}); + // Exact approval gates a real filesystem side effect. + const entered=deferred(),allow=deferred(),file=join(data,'reviewed.txt'); + inspect=async({action},signal)=>{if(action.kind==='tool'&&action.content.title==='write'){entered.resolve();await Promise.race([allow.promise,new Promise((_,reject)=>signal.addEventListener('abort',()=>reject(new Error('cancelled')),{once:true}))]);}return true;}; + faux.setResponses([fauxAssistantMessage(fauxToolCall('write',{path:file,content:'reviewed'}),{stopReason:'toolUse'}),fauxAssistantMessage('done')]); + let submission=await pi.rpc('deliver',turn('write'));let waiting=pi.rpc('wait',{agent,submission:submission.submission}); + await entered.promise;assert.equal(await exists(file),false);allow.resolve();await waiting;assert.equal(await readFile(file,'utf8'),'reviewed');result.exactWrite=true; + // A denied write never executes. + const denied=join(data,'denied.txt');inspect=async({action})=>action.kind!=='tool'; + faux.setResponses([fauxAssistantMessage(fauxToolCall('write',{path:denied,content:'denied'}),{stopReason:'toolUse'}),fauxAssistantMessage('stopped')]); + submission=await pi.rpc('deliver',turn('deny'));await pi.rpc('wait',{agent,submission:submission.submission});assert.equal(await exists(denied),false);result.deniedWrite=true; + // Both successful and nonzero shell output survive buffering when allowed. + const secret=join(data,'source.txt'),marker='RAW_RESULT_ONLY_7301';await writeFile(secret,marker); + for(const suffix of ['','; exit 3']) { + inspect=async()=>true; + faux.setResponses([fauxAssistantMessage(fauxToolCall('bash',{command:'cat '+JSON.stringify(secret)+suffix}),{stopReason:'toolUse'}),(context)=>{assert.ok(JSON.stringify(context).includes(marker));return fauxAssistantMessage('received');}]); + submission=await pi.rpc('deliver',turn('bash'+suffix,'run it','shell/'+suffix));await pi.rpc('wait',{agent,submission:submission.submission}); + } + result.shellOutput=true; + // Denied output is absent from the next provider request and durable entries. + inspect=async({action})=>action.kind!=='tool-result'||!JSON.stringify(action.content).includes(marker); + faux.setResponses([fauxAssistantMessage(fauxToolCall('read',{path:secret}),{stopReason:'toolUse'}),(context)=>{assert.ok(!JSON.stringify(context).includes(marker));return fauxAssistantMessage('withheld');}]); + submission=await pi.rpc('deliver',turn('output','read it','output/root'));await pi.rpc('wait',{agent,submission:submission.submission}); + assert.ok(!JSON.stringify(await pi.rpc('entries',{agent,conversation:submission.conversation})).includes(marker));result.withheldOutput=true; + // Denial occurs before the provider is called. + inspect=async({action})=>action.kind!=='context';const before=faux.state.callCount; + submission=await pi.rpc('deliver',turn('context','blocked','context/root'));await pi.rpc('wait',{agent,submission:submission.submission});assert.equal(faux.state.callCount,before);result.deniedContext=true; + // Cancellation while native Bash has emitted output never commits its partial text. + const barrier=join(data,'shell-started');inspect=async()=>true; + faux.setResponses([fauxAssistantMessage(fauxToolCall('bash',{command:'cat '+JSON.stringify(secret)+'; touch '+JSON.stringify(barrier)+'; sleep 30'}),{stopReason:'toolUse'})]); + submission=await pi.rpc('deliver',turn('cancel','run','cancel/root'));waiting=pi.rpc('wait',{agent,submission:submission.submission}); + await until(()=>exists(barrier));await pi.rpc('abort',{agent,conversation:submission.conversation});await waiting; + assert.ok(!JSON.stringify(await pi.rpc('entries',{agent,conversation:submission.conversation})).includes(marker));result.cancelledPartial=true; + // Safe tool replay after closing storage is checked again. + const replayEntered=deferred();inspect=async({action},signal)=>{if(action.kind==='tool'&&action.content.title==='tool_search'){replayEntered.resolve();await new Promise((_,reject)=>signal.addEventListener('abort',()=>reject(new Error('cancelled')),{once:true}));}return true;}; + faux.setResponses([fauxAssistantMessage(fauxToolCall('tool_search',{query:'read files'}),{stopReason:'toolUse'})]); + submission=await pi.rpc('deliver',turn('replay','read','replay/root'));waiting=pi.rpc('wait',{agent,submission:submission.submission});await replayEntered.promise;await pi.stop();await waiting; + let replayChecks=0;inspect=async({action})=>{if(action.kind==='tool'&&action.content.title==='tool_search')replayChecks++;return true;};pi=start((...args)=>inspect(...args));await pi.rpc('open',{agent,tools:[]}); + faux.setResponses([(context)=>{return fauxAssistantMessage('recovered');}]);await pi.rpc('wait',{agent,submission:submission.submission});assert.ok(replayChecks>0);result.safeReplay=true; + // Old history is checked before any resumed/new provider request. + inspect=async({action})=>action.kind!=='context';const restored=faux.state.callCount; + submission=await pi.rpc('deliver',turn('restored','continue','replay/root'));await pi.rpc('wait',{agent,submission:submission.submission});assert.equal(faux.state.callCount,restored);result.restoredHistory=true; + // Concurrent conversations retain distinct, native-mapped review sessions. + inspect=async()=>true; + faux.setResponses([fauxAssistantMessage('first'),fauxAssistantMessage('second')]); + const pair=await Promise.all([pi.rpc('deliver',turn('parallel1','one','parallel/one')),pi.rpc('deliver',turn('parallel2','two','parallel/two'))]); + await Promise.all(pair.map(item=>pi.rpc('wait',{agent,submission:item.submission}))); + assert.notEqual(pair[0].conversation,pair[1].conversation); + for(const item of pair)assert.ok(pi.actions.some(row=>row.action.kind==='context'&&row.action.session===String(item.conversation))); + result.concurrentSessions=true; + // Compaction's provider call is also held, including recovery after a close. + faux.setResponses([fauxAssistantMessage('read it'),fauxAssistantMessage('read it too')]); + for(const id of ['long1','long2']) {const item=await pi.rpc('deliver',turn(id,'word '.repeat(25000),'compact/root'));await pi.rpc('wait',{agent,submission:item.submission});} + const compactEntered=deferred();inspect=async({action},signal)=>{if(action.kind==='context'){compactEntered.resolve();await new Promise((_,reject)=>signal.addEventListener('abort',()=>reject(new Error('cancelled')),{once:true}));}return true;}; + const beforeCompact=faux.state.callCount; + const compact=pi.rpc('session',{agent,key:'compact/root',action:'compact'}).catch(error=>error.message); + await compactEntered.promise;assert.equal(faux.state.callCount,beforeCompact);await pi.stop();await compact; + let compactChecks=0;inspect=async({action})=>{if(action.kind==='context')compactChecks++;return true;}; + pi=start((...args)=>inspect(...args));await pi.rpc('open',{agent,tools:[]});faux.setResponses([fauxAssistantMessage('safe summary')]); + await pi.rpc('session',{agent,key:'compact/root',action:'compact'});assert.ok(compactChecks>0);assert.ok(faux.state.callCount>beforeCompact);result.compactionRecovery=true; + assert.ok(pi.actions.every(row=>row.agent===agent&&/^\d+$/.test(row.action.session))); + await assert.rejects(pi.rpc('test',{provider:'faux',model:'faux-1'}),/outside the protected/);result.closedSetupBypass=true; +} finally {await pi.stop();} +// Protected workers accept their exact secondary membership, with the same +// identity's home and other community memberships still rejected. +const secondary=agent+'@secondary.example%3A7447'; +const secondaryHost=start(()=>true,secondary); +try { + await secondaryHost.rpc('open',{agent:secondary,tools:[]}); + for(const other of [agent,agent+'@another.example'])await assert.rejects(secondaryHost.rpc('open',{agent:other,tools:[]}),/belongs to another agent/); + result.secondaryMembership=true; +} finally {await secondaryHost.stop();} +// A one-shot reviewer has no opened agent and separates trusted instructions +// from adversarial evidence without offering tools. +const reviewer=start(()=>true,null); +try { + faux.setResponses([(context)=>{ + assert.equal(context.messages[0].role,'system'); + assert.equal(context.messages[0].content,'TRUSTED_REVIEW_INSTRUCTIONS'); + assert.equal(context.messages[1].role,'user'); + assert.equal(context.messages[1].content,'UNTRUSTED_ACTION_EVIDENCE'); + assert.ok(context.messages.every(message=>!message.toolsAdded?.length)); + return fauxAssistantMessage('{"blocked":true,"reason":"Review this action"}'); + }]); + const reviewed=await reviewer.rpc('review',{id:'isolated-review',provider:'faux',model:'faux-1',prompt:{system:'TRUSTED_REVIEW_INSTRUCTIONS',input:'UNTRUSTED_ACTION_EVIDENCE'}}); + assert.equal(JSON.parse(reviewed.text).blocked,true); + assert.equal(reviewer.actions.length,0); + result.toolFreeReview=true; +} finally {await reviewer.stop();} +// Wrappers also handle errors and exact snapshots independently. +const emitted=[];let seen; +const api={conversationId:1,output:value=>emitted.push(value),diagnostic:value=>emitted.push(value),details:async value=>emitted.push(value)}; +const wrapped=protectTool({name:'throw',execute:async(_args,api)=>{api.output('partial');api.diagnostic({message:'diagnostic'});await api.details({secret:'details'});throw new Error('private error');}},async action=>{seen=action;if(action.kind==='tool-result')throw new Error('denied');}); +await assert.rejects(wrapped.execute({},api,{abortSignal:new AbortController().signal}),/denied/);assert.deepEqual(emitted,[]);assert.ok(JSON.stringify(seen).includes('private error'));result.thrownError=true; +const gate=deferred(),entered=deferred();let captured; +const sources=[];const checked=protectModels(models,async action=>{sources.push(action.content);entered.resolve();await gate.promise;},async id=>{if(id!=='known')throw new Error('unknown session');return 1;}); +const input={messages:[{role:'system',content:'dynamic instructions',sections:{memory:'untrusted memory text'},toolsAdded:[{name:'schema_fixture',description:'tool metadata '.repeat(25000),parameters:{type:'object',properties:{}}}],timestamp:0},{role:'user',content:'reviewed',timestamp:1}]};const model=structuredClone(faux.getModel()); +faux.setResponses([(context,_opts,_state,actual)=>{captured={context,actual};return fauxAssistantMessage('ok');}]); +const reply=checked.completeSimple(model,input,{sessionId:'known'});await entered.promise;input.messages[1].content='mutated';model.id='mutated';gate.resolve();await reply; +assert.equal(captured.context.messages[1].content,'reviewed');assert.ok(JSON.stringify(sources).includes('untrusted memory text'));assert.ok(!JSON.stringify(sources).includes('tool metadata '));assert.ok(sources.every(source=>Buffer.byteLength(JSON.stringify(source))<65536));assert.notEqual(captured.actual.id,'mutated');result.exactModel=true; +const count=faux.state.callCount;await checked.completeSimple(faux.getModel(),input,{sessionId:'unknown'});assert.equal(faux.state.callCount,count);result.unknownSession=true; +for(const method of ['streamDeferred','fetchDeferred','cancelDeferred'])assert.throws(()=>checked[method](faux.getModel(),{}),/does not support deferred/); +assert.throws(()=>checked.streamSimple(faux.getModel(),input,{onPayload:()=>({})}),/payload override/);result.closedOverrides=true; +console.log(JSON.stringify(result)); +`; +it.skipIf(!deps)( + "enforces protected Pi against real side effects, provider calls and durable recovery", + async () => { + const folder = await mkdtemp(join(tmpdir(), "buzz-pi-protection-")); + try { + await symlink( + join(deps ?? "", "node_modules"), + join(folder, "node_modules"), + ); + await writeFile(join(folder, "host.mjs"), host); + await writeFile(join(folder, "driver.mjs"), driver); + const { stdout } = await promisify(execFile)( + process.execPath, + [ + "--disable-sigusr1", + "--disable-warning=ExperimentalWarning", + join(folder, "driver.mjs"), + folder, + ], + { cwd: folder, timeout: 30000 }, + ); + const result = JSON.parse(stdout.trim().split("\n").at(-1) ?? "{}"); + expect(result).toEqual( + Object.fromEntries( + [ + "exactWrite", + "deniedWrite", + "shellOutput", + "withheldOutput", + "deniedContext", + "cancelledPartial", + "safeReplay", + "restoredHistory", + "closedSetupBypass", + "thrownError", + "exactModel", + "unknownSession", + "closedOverrides", + "concurrentSessions", + "compactionRecovery", + "toolFreeReview", + "secondaryMembership", + ].map((name) => [name, true]), + ), + ); + } finally { + await rm(folder, { recursive: true, force: true }); + } + }, + 35000, +); diff --git a/src/bundled/pi/runtime.test.ts b/src/bundled/pi/runtime.test.ts index 0a89ae32db..fb16d1479e 100644 --- a/src/bundled/pi/runtime.test.ts +++ b/src/bundled/pi/runtime.test.ts @@ -2213,3 +2213,52 @@ it("keeps foreign-channel roots out of conversation context", async () => { "Private root", ); }); + +it.each(["stop", "retire"])( + "keeps a protected final answer local after %s", + async (mode) => { + const answered = deferred(); + const check = vi.fn(async () => {}); + const f = setup({ + wait: () => answered.promise, + abort: () => ({}), + runtime: () => ({ + file: "/runtime/host.mjs", + node: "/node/bin/node", + codeRoot: "/runtime", + nodeRoot: "/node", + data: "/data", + auth: "/auth", + }), + }); + let stop!: () => Promise; + const turn = f.deliver( + { type: "mention", event: reply("2", "@Pi look") }, + { + interruptible: (fn) => { + stop = fn; + }, + }, + ); + await turn.run; + Object.assign(f.agent, { + protection: { + key: "test", + stateKey: "test", + processProtection: { readOnlyPaths: [], hiddenPaths: [] }, + assertAvailable() {}, + check, + }, + }); + if (mode === "stop") await stop(); + else await f.runtime.retire(self); + answered.resolve({ + status: "done", + answer: 5, + text: "The build is green.", + }); + await f.settled(turn); + expect(f.published).toEqual([]); + expect(check).not.toHaveBeenCalled(); + }, +); diff --git a/src/bundled/pi/runtime.ts b/src/bundled/pi/runtime.ts index 207c6b261c..60bb50bd7c 100644 --- a/src/bundled/pi/runtime.ts +++ b/src/bundled/pi/runtime.ts @@ -173,6 +173,7 @@ type Entry = { export class PiRuntime { private readonly agents = new Map(); + private readonly revisions = new Map(); private readonly memoryReads = new MemoryReads(); private readonly controls = new SessionControls(() => this.notify()); private readonly listeners = new Set<() => void>(); @@ -212,42 +213,58 @@ export class PiRuntime { input: unknown, ) => Promise, ) { - this.pi = new Pi(() => this.spawn, { - tool: (call) => this.tool(call), - hook: (call) => { - if (!this.agents.has(call.agent)) - throw new Error("Pi hook agent is no longer active"); - return ( - this.hook?.( - call.agent, - String(call.conversation), - call.event, - call.input, - ) ?? Promise.resolve(null) - ); + this.pi = new Pi( + () => this.spawn, + { + tool: (call) => this.tool(call), + hook: (call) => { + if (!this.agents.has(call.agent)) + throw new Error("Pi hook agent is no longer active"); + return ( + this.hook?.( + call.agent, + String(call.conversation), + call.event, + call.input, + ) ?? Promise.resolve(null) + ); + }, + event: (event) => this.event(event), + prompt: ({ agent, request }) => + this.agents.get(agent)?.receipts.get(request)?.prompt?.() ?? + Promise.resolve(undefined), }, - event: (event) => this.event(event), - prompt: ({ agent, request }) => - this.agents.get(agent)?.receipts.get(request)?.prompt?.() ?? - Promise.resolve(undefined), - }); + (pubkey) => + (this.agents.get(pubkey)?.handle ?? this.handle(pubkey))?.protection, + ); this.setup = new PiSetup(this.pi); } + async retire(pubkey: string) { + const entry = this.agents.get(pubkey); + if (entry) { + entry.closed = true; + } + this.agents.delete(pubkey); + await this.pi.retire(pubkey); + } + /** Brings the runtime in line with this type's agents. A removed one's * storage closes; its unfinished runs continue when it is back. */ sync(agents: readonly Agent[]) { const current = new Set(agents.map((agent) => agent.key)); // Its storage may be open for a turn or only for listing its sessions. for (const agent of new Set([...this.running, ...this.agents.keys()])) - if (!current.has(agent)) { - const entry = this.agents.get(agent); - if (entry) entry.closed = true; - this.agents.delete(agent); - void this.pi.request("close", { agent }).catch(() => undefined); - } + if (!current.has(agent)) void this.retire(agent).catch(() => undefined); this.running = current; for (const agent of agents) { + const revision = agent.protectionRevision ?? 0; + const previous = this.revisions.get(agent.key); + this.revisions.set(agent.key, revision); + if (previous !== undefined && previous !== revision) { + void this.retire(agent.key).catch(() => undefined); + continue; + } const entry = this.agents.get(agent.key); if (!entry) continue; entry.config = config(agent.config); @@ -407,7 +424,8 @@ export class PiRuntime { }; if (action === "interrupt") { this.controls.interrupt(control); - this.agents.get(agent.key)?.policies.get(key)?.cancel(); + const entry = this.agents.get(agent.key); + entry?.policies.get(key)?.cancel(); return act(); } if (this.sessionState(agent, place) !== "idle") @@ -419,7 +437,9 @@ export class PiRuntime { } dispose() { - for (const entry of this.agents.values()) entry.closed = true; + for (const entry of this.agents.values()) { + entry.closed = true; + } this.agents.clear(); void this.pi.dispose(); } @@ -624,7 +644,7 @@ export class PiRuntime { }); const policy = { topLevel, - cancel() { + cancel: () => { interrupted = true; prior?.cancel(); }, @@ -644,6 +664,7 @@ export class PiRuntime { workingUntil(settled); // Stop on its activity stops the run the request is in. interruptible?.(async () => { + interrupted = true; await handed; if (!submitted) return; await this.pi.request("abort", { @@ -652,7 +673,8 @@ export class PiRuntime { }); }); void settled.then(async (result) => { - if (!result.ok) await this.report(agent, where, result.error ?? ""); + if (!result.ok && !interrupted && !entry.closed) + await this.report(agent, where, result.error ?? ""); }); } @@ -719,6 +741,11 @@ export class PiRuntime { const client = entry?.handle && appClient(entry.handle, { + protectionContext: () => + this.pi.request("protectionContext", { + agent: call.agent, + conversation: call.conversation, + }), relayAuthor: this.reads(call.agent)?.relayAuthor?.(), signingAuthority: this.reads(call.agent)?.workflowAuthority(), read: (path) => this.readFile(path, { cwd: entry.config.workspace }), diff --git a/src/bundled/pi/setup.ts b/src/bundled/pi/setup.ts index a515550038..09688892fa 100644 --- a/src/bundled/pi/setup.ts +++ b/src/bundled/pi/setup.ts @@ -3,6 +3,7 @@ // own `~/.pi/agent/auth.json`, shared with the `pi` command, never into Buzz. import type { ModelChoice } from "../../features/agents/ModelPicker"; import type { Pi } from "./pi"; +import type { ReviewPrompt } from "../../features/agent-protection/models"; export type Provider = Readonly<{ id: string; @@ -92,4 +93,26 @@ export class PiSetup { test(choice: Readonly<{ provider: string; model: string }>): Promise { return this.pi.request("test", choice); } + async review( + choice: { provider?: string; id: string }, + prompt: ReviewPrompt, + signal: AbortSignal, + ) { + signal.throwIfAborted(); + const id = crypto.randomUUID(); + const cancel = () => { + void this.pi.request("cancel-review", { id }).catch(() => undefined); + }; + signal.addEventListener("abort", cancel, { once: true }); + try { + const result = await this.pi.request<{ + text: string; + model: { provider: string; id: string }; + }>("review", { ...choice, model: choice.id, prompt, id }); + signal.throwIfAborted(); + return result; + } finally { + signal.removeEventListener("abort", cancel); + } + } }