diff --git a/extensions/shared/child-session.ts b/extensions/shared/child-session.ts index 68c2acee..115e2e03 100644 --- a/extensions/shared/child-session.ts +++ b/extensions/shared/child-session.ts @@ -13,6 +13,11 @@ import { type SessionShutdownEvent, SettingsManager, } from "@earendil-works/pi-coding-agent"; +import { + OPENPI_OWNER_SOURCE_PATHS, + OPENPI_TOOL_SURFACE, + type OpenPiToolOwner, +} from "./tool-surface.ts"; export const CHILD_SHUTDOWN_TIMEOUT_MS = 5_000; @@ -471,6 +476,42 @@ export const CHILD_EXCLUDED_TOOL_NAMES = [ "context_pivot", ] as const; +const PARENT_ONLY_OPENPI_EXTENSION_PATHS = new Set( + (Object.keys(OPENPI_TOOL_SURFACE) as OpenPiToolOwner[]) + .filter((owner) => { + const { entry, deferred } = OPENPI_TOOL_SURFACE[owner]; + const toolNames = [...entry, ...deferred]; + return ( + toolNames.length > 0 && + toolNames.every((name) => + CHILD_EXCLUDED_TOOL_NAMES.includes(name as never), + ) + ); + }) + .map((owner) => canonicalExistingPath(OPENPI_OWNER_SOURCE_PATHS[owner])) + .filter( + (extensionPath): extensionPath is string => extensionPath !== undefined, + ), +); + +function isVerifiedParentOnlyOpenPiExtension(extension: { + path: string; + resolvedPath: string; + sourceInfo: { path: string }; +}) { + return [ + extension.path, + extension.resolvedPath, + extension.sourceInfo.path, + ].some((candidate) => { + const canonicalPath = canonicalExistingPath(candidate); + return ( + canonicalPath !== undefined && + PARENT_ONLY_OPENPI_EXTENSION_PATHS.has(canonicalPath) + ); + }); +} + /** * Fresh SDK options avoid turning the denylist into an accidental allowlist. * @@ -517,7 +558,15 @@ export async function createChildResources(options: ChildResourceOptions) { cwd: options.cwd, agentDir, settingsManager, - extensionsOverride: excludeOpenPiGitInfoExtension, + extensionsOverride(base) { + const withoutGitInfo = excludeOpenPiGitInfoExtension(base); + return { + ...withoutGitInfo, + extensions: withoutGitInfo.extensions.filter( + (extension) => !isVerifiedParentOnlyOpenPiExtension(extension), + ), + }; + }, ...(options.appendSystemPrompt ? { appendSystemPrompt: options.appendSystemPrompt } : {}), diff --git a/extensions/shared/tool-surface.ts b/extensions/shared/tool-surface.ts index 86c8973f..c3e0f443 100644 --- a/extensions/shared/tool-surface.ts +++ b/extensions/shared/tool-surface.ts @@ -135,7 +135,7 @@ interface ToolSurfaceState { subscribed: boolean; } -const OWNER_SOURCE_PATHS = { +export const OPENPI_OWNER_SOURCE_PATHS = { capabilities: fileURLToPath( new URL("../capabilities/index.ts", import.meta.url), ), @@ -268,7 +268,7 @@ function availableOwnedToolNames( } const expectedSource = - state.sourceByOwner.get(owner) ?? OWNER_SOURCE_PATHS[owner]; + state.sourceByOwner.get(owner) ?? OPENPI_OWNER_SOURCE_PATHS[owner]; return new Set( reportedTools .filter( diff --git a/package.json b/package.json index a6b2787e..4f02ddbb 100644 --- a/package.json +++ b/package.json @@ -85,6 +85,7 @@ "format:check": "biome format .", "lint": "biome lint . --error-on-warnings", "typecheck": "tsc --noEmit", + "benchmark:workflow-child-startup": "node --experimental-strip-types scripts/benchmark-workflow-child-startup.mjs", "test": "node scripts/run-tests.mjs", "provenance": "node scripts/provenance.mjs" }, diff --git a/scripts/benchmark-workflow-child-startup.mjs b/scripts/benchmark-workflow-child-startup.mjs new file mode 100644 index 00000000..0a4b09ed --- /dev/null +++ b/scripts/benchmark-workflow-child-startup.mjs @@ -0,0 +1,218 @@ +import fs from "node:fs"; +import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { syncBuiltinESMExports } from "node:module"; +import { tmpdir } from "node:os"; +import * as path from "node:path"; +import { performance } from "node:perf_hooks"; +import { fileURLToPath } from "node:url"; +import { + createAgentSession, + SessionManager, +} from "@earendil-works/pi-coding-agent"; +import { + bindChildSessionExtensions, + childToolPolicy, + createChildResources, + shutdownAndDisposeChildSession, +} from "../extensions/shared/child-session.ts"; + +const foreignRunCounts = [0, 100, 1_000]; +const childCounts = [1, 4, 8]; +const packageRoot = fileURLToPath(new URL("../", import.meta.url)); + +async function writeForeignRuns(agentDir, count) { + const artifacts = new Map(); + for (let index = 0; index < count; index++) { + const runId = `wf_${index.toString(16).padStart(8, "0")}`; + const runDir = path.join(agentDir, "workflows", runId); + const workflow = JSON.stringify({ + runId, + sessionId: `foreign-session-${index}`, + status: "completed", + startedAt: index, + finishedAt: index + 1, + agents: [], + phases: [], + resultArtifact: "result.json", + transcriptArtifact: "transcripts.json", + }); + const result = JSON.stringify({ result: `foreign result ${index}` }); + const transcripts = JSON.stringify({}); + await mkdir(runDir, { recursive: true }); + await Promise.all( + [ + ["workflow.json", workflow], + ["result.json", result], + ["transcripts.json", transcripts], + ].map(async ([name, content]) => { + const artifactPath = path.join(runDir, name); + artifacts.set(artifactPath, content); + await writeFile(artifactPath, content); + }), + ); + } + return artifacts; +} + +function instrumentWorkflowArtifacts(artifacts) { + const originalReadFileSync = fs.readFileSync; + const originalJsonParse = JSON.parse; + const workflowJson = new Set( + [...artifacts.entries()] + .filter( + ([artifactPath]) => path.basename(artifactPath) === "workflow.json", + ) + .map(([, content]) => content), + ); + const metrics = { + workflowSyncReadCalls: 0, + workflowBytesRead: 0, + workflowJsonParses: 0, + workflowSyncReadMilliseconds: 0, + }; + + fs.readFileSync = (...args) => { + const startedAt = performance.now(); + const content = originalReadFileSync(...args); + const elapsed = performance.now() - startedAt; + const artifactPath = args[0]; + if (typeof artifactPath === "string" && artifacts.has(artifactPath)) { + metrics.workflowSyncReadCalls++; + metrics.workflowSyncReadMilliseconds += elapsed; + metrics.workflowBytesRead += Buffer.byteLength( + typeof content === "string" ? content : content.toString(), + ); + } + return content; + }; + syncBuiltinESMExports(); + JSON.parse = (text, reviver) => { + if (typeof text === "string" && workflowJson.has(text)) { + metrics.workflowJsonParses++; + } + return originalJsonParse(text, reviver); + }; + + return { + metrics, + restore() { + fs.readFileSync = originalReadFileSync; + syncBuiltinESMExports(); + JSON.parse = originalJsonParse; + }, + }; +} + +async function startChild({ cwd, agentDir, index }) { + let session; + try { + const { loader, settingsManager } = await createChildResources({ + cwd, + agentDir, + projectTrusted: true, + }); + ({ session } = await createAgentSession({ + cwd, + agentDir, + resourceLoader: loader, + settingsManager, + sessionManager: SessionManager.inMemory(path.join(cwd, String(index))), + ...childToolPolicy(), + })); + await bindChildSessionExtensions(session); + return session; + } catch (error) { + if (session) await shutdownAndDisposeChildSession(session); + throw error; + } +} + +async function benchmarkScenario({ foreignRunCount, childCount }) { + const directory = await mkdtemp( + path.join(tmpdir(), "openpi-workflow-child-startup-"), + ); + const agentDir = path.join(directory, "agent"); + const cwd = path.join(directory, "project"); + const originalAgentDir = process.env.PI_CODING_AGENT_DIR; + let instrumentation; + let sessions = []; + + try { + await Promise.all([ + mkdir(cwd, { recursive: true }), + mkdir(agentDir, { recursive: true }), + ]); + await writeFile( + path.join(agentDir, "settings.json"), + JSON.stringify({ packages: [packageRoot] }), + ); + const artifacts = await writeForeignRuns(agentDir, foreignRunCount); + process.env.PI_CODING_AGENT_DIR = agentDir; + instrumentation = instrumentWorkflowArtifacts(artifacts); + + const startedAt = performance.now(); + const starts = await Promise.allSettled( + Array.from({ length: childCount }, (_value, index) => + startChild({ cwd, agentDir, index }), + ), + ); + sessions = starts + .filter((result) => result.status === "fulfilled") + .map((result) => result.value); + const failedStart = starts.find((result) => result.status === "rejected"); + if (failedStart) throw failedStart.reason; + const batchStartupMilliseconds = performance.now() - startedAt; + const result = { + childCount, + foreignRunCount, + ...instrumentation.metrics, + batchStartupMilliseconds, + }; + if ( + result.workflowSyncReadCalls !== 0 || + result.workflowBytesRead !== 0 || + result.workflowJsonParses !== 0 + ) { + throw new Error( + `Workflow artifact gate failed for ${childCount} children and ${foreignRunCount} foreign runs: ` + + `${result.workflowSyncReadCalls} reads, ${result.workflowBytesRead} bytes, ` + + `${result.workflowJsonParses} parses`, + ); + } + return result; + } finally { + await Promise.all( + sessions.map((session) => shutdownAndDisposeChildSession(session)), + ); + instrumentation?.restore(); + if (originalAgentDir === undefined) { + delete process.env.PI_CODING_AGENT_DIR; + } else { + process.env.PI_CODING_AGENT_DIR = originalAgentDir; + } + await rm(directory, { recursive: true, force: true }); + } +} + +const scenarios = []; +for (const foreignRunCount of foreignRunCounts) { + for (const childCount of childCounts) { + scenarios.push(await benchmarkScenario({ foreignRunCount, childCount })); + } +} + +console.log( + "children | foreign runs | reads | bytes | parses | read ms | startup ms", +); +for (const scenario of scenarios) { + console.log( + `${String(scenario.childCount).padStart(8)} | ` + + `${String(scenario.foreignRunCount).padStart(12)} | ` + + `${String(scenario.workflowSyncReadCalls).padStart(5)} | ` + + `${String(scenario.workflowBytesRead).padStart(5)} | ` + + `${String(scenario.workflowJsonParses).padStart(6)} | ` + + `${scenario.workflowSyncReadMilliseconds.toFixed(3).padStart(7)} | ` + + scenario.batchStartupMilliseconds.toFixed(3), + ); +} +console.log(JSON.stringify({ scenarios })); diff --git a/tests/extensions/shared/child-session.test.ts b/tests/extensions/shared/child-session.test.ts index 730336f3..4ef5daeb 100644 --- a/tests/extensions/shared/child-session.test.ts +++ b/tests/extensions/shared/child-session.test.ts @@ -8,6 +8,8 @@ import { rm, writeFile, } from "node:fs/promises"; +import fs from "node:fs"; +import { syncBuiltinESMExports } from "node:module"; import { tmpdir } from "node:os"; import * as path from "node:path"; import test from "node:test"; @@ -201,6 +203,184 @@ test("child binding restores only requested child-safe package tools after paren }); }); +test("child resources remove only verified parent-only OpenPI extensions", async () => { + await withTempDir(async (directory) => { + const cwd = path.join(directory, "project"); + const agentDir = path.join(directory, "agent"); + const extensionsDir = path.join(agentDir, "extensions"); + await mkdir(extensionsDir, { recursive: true }); + await writeFile( + path.join(agentDir, "settings.json"), + JSON.stringify({ + packages: [fileURLToPath(new URL("../../../", import.meta.url))], + }), + ); + await writeFile( + path.join(extensionsDir, "third-party.ts"), + `export default function (pi) { + pi.registerTool({ + name: "subagent_spawn", + label: "Third-party subagent spawn", + description: "fixture", + parameters: { type: "object", properties: {} }, + async execute() { return { content: [{ type: "text", text: "ok" }] }; }, + }); + }`, + ); + + const { loader } = await createChildResources({ + cwd, + agentDir, + projectTrusted: true, + }); + const extensions = loader.getExtensions().extensions; + + assert.equal( + extensions.some((extension) => extension.tools.has("openpi_load_tools")), + false, + "parent-only OpenPI extension should not reach the child runtime", + ); + assert.equal( + extensions.some((extension) => extension.tools.has("fd")), + true, + "child-safe OpenPI file-search extension should remain", + ); + assert.equal( + extensions.some((extension) => extension.tools.has("git_show")), + true, + "child-safe OpenPI git-read extension should remain", + ); + assert.equal( + extensions.some((extension) => extension.tools.has("subagent_spawn")), + true, + "ordinary third-party extensions must survive tool-name collisions", + ); + }); +}); + +test("production child binding skips foreign Workflow artifacts", async () => { + await withTempDir(async (directory) => { + const cwd = path.join(directory, "project"); + const agentDir = path.join(directory, "agent"); + const runDir = path.join(agentDir, "workflows", "wf_f0e1"); + const artifactContents = [ + JSON.stringify({ + runId: "wf_f0e1", + sessionId: "foreign-session", + status: "completed", + startedAt: 1, + finishedAt: 2, + agents: [], + phases: [], + resultArtifact: "result.json", + transcriptArtifact: "transcripts.json", + }), + JSON.stringify({ result: "foreign result" }), + JSON.stringify({}), + ]; + const artifactPaths = new Set([ + path.join(runDir, "workflow.json"), + path.join(runDir, "result.json"), + path.join(runDir, "transcripts.json"), + ]); + let readCalls = 0; + let readBytes = 0; + let workflowParses = 0; + const originalReadFileSync = fs.readFileSync; + const originalJsonParse = JSON.parse; + const originalAgentDir = process.env.PI_CODING_AGENT_DIR; + let session: + | Awaited>["session"] + | undefined; + + await mkdir(runDir, { recursive: true }); + await writeFile( + path.join(agentDir, "settings.json"), + JSON.stringify({ + packages: [fileURLToPath(new URL("../../../", import.meta.url))], + }), + ); + await Promise.all( + ["workflow.json", "result.json", "transcripts.json"].map((name, index) => + writeFile(path.join(runDir, name), artifactContents[index]!), + ), + ); + + Object.defineProperty(fs, "readFileSync", { + value: (...args: Parameters) => { + const content = originalReadFileSync(...args); + const filePath = args[0]; + if (typeof filePath === "string" && artifactPaths.has(filePath)) { + readCalls++; + readBytes += Buffer.byteLength( + typeof content === "string" ? content : content.toString(), + ); + } + return content; + }, + }); + syncBuiltinESMExports(); + JSON.parse = (text, reviver) => { + if (text === artifactContents[0]) { + workflowParses++; + } + return originalJsonParse(text, reviver); + }; + process.env.PI_CODING_AGENT_DIR = agentDir; + + try { + const { loader, settingsManager } = await createChildResources({ + cwd, + agentDir, + projectTrusted: true, + }); + const structuredOutput = defineTool({ + name: "structured_output", + label: "Structured Output", + description: "fixture structured result", + parameters: Type.Object({ value: Type.String() }), + async execute(_id, params) { + return { + content: [{ type: "text", text: params.value }], + details: {}, + }; + }, + }); + ({ session } = await createAgentSession({ + cwd, + agentDir, + resourceLoader: loader, + settingsManager, + sessionManager: SessionManager.inMemory(cwd), + customTools: [structuredOutput], + ...childToolPolicy(), + })); + await bindChildSessionExtensions(session); + + assert.equal( + session.getActiveToolNames().includes("structured_output"), + true, + "dynamically registered workflow output tool should remain available", + ); + assert.equal(readCalls, 0); + assert.equal(readBytes, 0); + assert.equal(workflowParses, 0); + } finally { + if (session) await shutdownAndDisposeChildSession(session); + Object.defineProperty(fs, "readFileSync", { + value: originalReadFileSync, + }); + syncBuiltinESMExports(); + JSON.parse = originalJsonParse; + if (originalAgentDir === undefined) { + delete process.env.PI_CODING_AGENT_DIR; + } else { + process.env.PI_CODING_AGENT_DIR = originalAgentDir; + } + } + }); +}); + test("child denylist keeps extension and workflow structured tools available", async () => { await withTempDir(async (directory) => { let starts = 0;