pi-ext

import { spawnSync } from "node:child_process";
import {
	mkdirSync,
	mkdtempSync,
	readFileSync,
	rmSync,
	symlinkSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join, resolve } from "node:path";
import { fileURLToPath } from "node:url";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { buildExtension } from "./build.ts";
import { forbiddenImports } from "./verify.ts";

const root = resolve(import.meta.dirname, "../..");
const workspaceManifest = resolve(
	fileURLToPath(import.meta.resolve("@earendil-works/pi-coding-agent")),
	"../../package.json",
);
const hosts = [
	...new Set([
		workspaceManifest,
		process.env["PI_CONTEXT_PACKAGE_JSON_PATH"] ?? workspaceManifest,
	]),
].map((manifest) => {
	const metadata = JSON.parse(readFileSync(manifest, "utf8")) as {
		version: string;
		bin: { pi: string };
	};
	return { dir: resolve(manifest, ".."), ...metadata };
});

// Observe native SDK evaluation, not merely resolution or text in a bundle.
const hook = `import { registerHooks } from 'node:module';
registerHooks({ load(url, context, nextLoad) {
	const loaded = nextLoad(url, context);
	if (url.replaceAll('\\\\', '/').endsWith('/pi-coding-agent/dist/index.js'))
		process.stderr.write('PFUI_SDK_LOAD=' + url + '\\n');
	return loaded;
}});`;

function sandboxEnvironment(sandbox: string, agentDir: string) {
	const env = { ...process.env };
	for (const name of Object.keys(env)) {
		if (name.startsWith("PI_CONTEXT_") || name === "PI_MANAGED_INSTALL_ROOT")
			delete env[name];
	}
	return Object.assign(env, {
		HOME: sandbox,
		USERPROFILE: sandbox,
		APPDATA: sandbox,
		LOCALAPPDATA: sandbox,
		PI_CODING_AGENT_DIR: agentDir,
		XDG_CONFIG_HOME: sandbox,
		XDG_CACHE_HOME: sandbox,
		XDG_STATE_HOME: sandbox,
		XDG_DATA_HOME: sandbox,
		PI_OFFLINE: "1",
	});
}

function probe(host: (typeof hosts)[number], entry: string | string[]) {
	const sandbox = mkdtempSync(join(tmpdir(), "pfui-runtime-"));
	try {
		const agentDir = join(sandbox, "agent");
		mkdirSync(agentDir);
		const observer = join(sandbox, "observer.ts");
		writeFileSync(
			observer,
			`export default function(pi) {
	pi.on('session_start', () => {
		console.log(JSON.stringify({ type: 'pfui_context', version: process.env.PI_CONTEXT_VERSION, packageDir: process.env.PI_CONTEXT_PACKAGE_DIR }));
	});
}`,
		);
		const env = sandboxEnvironment(sandbox, agentDir);
		const started = performance.now();
		const child = spawnSync(
			process.execPath,
			[
				"--import",
				`data:text/javascript,${encodeURIComponent(hook)}`,
				join(host.dir, host.bin.pi),
				"--mode",
				"rpc",
				"--no-session",
				"--approve",
				"--no-extensions",
				"--no-context-files",
				"--no-skills",
				"--no-prompt-templates",
				"--no-themes",
				...[entry].flat().flatMap((path) => ["--extension", path]),
				"--extension",
				observer,
			],
			{
				cwd: sandbox,
				env,
				encoding: "utf8",
				timeout: 20_000,
				input: '{"type":"get_state"}\n',
				maxBuffer: 1024 * 1024,
			},
		);
		const milliseconds = performance.now() - started;
		expect(child.error, child.stderr).toBeUndefined();
		expect(child.status, child.stderr).toBe(0);
		const lines = child.stdout
			.trim()
			.split("\n")
			.map((line) => JSON.parse(line));
		expect(lines).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					type: "response",
					command: "get_state",
					success: true,
				}),
			]),
		);
		const context = child.stderr
			.split("\n")
			.filter((line) => line.startsWith('{"type":"pfui_context"'))
			.map((line) => JSON.parse(line))[0] as
			| { version: string; packageDir: string }
			| undefined;
		expect(context, child.stderr).toBeDefined();
		const sdkLoads = child.stderr
			.split("\n")
			.filter((line) => line.startsWith("PFUI_SDK_LOAD="));
		return { context: context!, sdkLoads, milliseconds };
	} finally {
		rmSync(sandbox, { recursive: true, force: true });
	}
}

// Drive the real public SDK synchronously to completion, not RPC startup followed
// by stdin EOF. This exercises the same resource loader and actual child sessions.
function compositionProbe(host: (typeof hosts)[number], entry: string[]) {
	const sandbox = mkdtempSync(join(tmpdir(), "sdk-composition-"));
	try {
		const agentDir = join(sandbox, "agent");
		mkdirSync(agentDir);
		const observer = join(sandbox, "observer.ts");
		writeFileSync(
			observer,
			`import { VERSION, getPackageDir } from '@earendil-works/pi-coding-agent';
export default function(pi) {
	pi.registerProvider('identity', globalThis[Symbol.for('pi-ext-sdk-probe/provider')]);
	const records = globalThis[Symbol.for('pi-ext-sdk-probe')] ??= [];
	records.push({ stage: 'factory', version: VERSION, packageDir: getPackageDir() });
	pi.on('session_start', (_event, ctx) => records.push({
		stage: 'session', id: ctx.sessionManager.getSessionId(),
		version: VERSION, packageDir: getPackageDir(),
		contextVersion: process.env.PI_CONTEXT_VERSION,
		contextPackageDir: process.env.PI_CONTEXT_PACKAGE_DIR,
		activeTools: pi.getActiveTools(),
		nuOwner: pi.getAllTools().find(tool => tool.name === 'nu')?.sourceInfo.path,
		hasNuSession: pi.getAllTools().some(tool => tool.name === 'nu_session')
	}));
	pi.on('tool_result', event => {
		if (event.toolName === 'nu') records.push({ stage: 'nu_result', isError: event.isError, content: event.content });
	});
}`,
		);
		const sources = [entry[0]!, resolve(entry[3]!, ".."), observer];
		writeFileSync(
			join(agentDir, "settings.json"),
			JSON.stringify({
				angel: {
					pairs: [{ executor: "identity/probe", advisor: "identity/probe" }],
					subagentExtensions: sources,
				},
				ultra: {
					autoEnable: true,
					subagentExtensions: sources,
					maxRetries: 0,
					actionSummarizerModel: null,
				},
			}),
		);
		const driver = join(sandbox, "driver.mjs");
		writeFileSync(
			driver,
			`import { pathToFileURL } from 'node:url';
const { hostDir, entries, observer, agentDir } = JSON.parse(process.argv[2]);
globalThis.fetch = async () => { throw new Error('Unexpected network request in SDK identity probe'); };
const sdk = await import(pathToFileURL(hostDir + '/dist/index.js').href);
const ai = await import(import.meta.resolve('@earendil-works/pi-ai', pathToFileURL(hostDir + '/package.json').href));
sdk.initTheme('dark', false);
const core = ai.createFauxCore({ api: 'openai', provider: 'identity', models: [{ id: 'probe', name: 'Identity Probe' }] });
const runtime = await sdk.ModelRuntime.create({
	credentials: new ai.InMemoryCredentialStore(), modelsStore: new ai.InMemoryModelsStore(), refreshOnCreate: false
});
const provider = { api: 'openai', apiKey: 'test-key', baseUrl: 'https://identity.invalid',
	streamSimple: core.streamSimple, models: [{ id: 'probe', name: 'Identity Probe', reasoning: false,
		input: ['text'], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 200000, maxTokens: 8192 }] };
globalThis[Symbol.for('pi-ext-sdk-probe/provider')] = provider;
runtime.registerProvider('identity', provider);
const loader = new sdk.DefaultResourceLoader({ cwd: process.cwd(), agentDir, noExtensions: true,
	noSkills: true, noPromptTemplates: true, noThemes: true, noContextFiles: true,
	additionalExtensionPaths: [...entries, observer] });
await loader.reload();
const loaded = loader.getExtensions();
if (loaded.errors.length) throw new Error(JSON.stringify(loaded.errors));
const created = await sdk.createAgentSession({ cwd: process.cwd(), agentDir, modelRuntime: runtime,
	model: runtime.getModel('identity', 'probe'), resourceLoader: loader,
	sessionManager: sdk.SessionManager.inMemory(process.cwd()) });
const session = created.session;
const errors = [];
const unsubscribe = session.extensionRunner.onError(error => errors.push(String(error.error)));
try {
	await session.bindExtensions({ mode: 'print' });
	core.setResponses([ai.fauxAssistantMessage('Parent catalog initialized.')]);
	await session.prompt('Initialize subagent flag catalog', { expandPromptTemplates: false, source: 'extension' });
	const catalogFactories = globalThis[Symbol.for('pi-ext-sdk-probe')].filter(record => record.stage === 'factory').length;
	const tool = name => {
		for (const extension of loaded.extensions) {
			const definition = extension.tools.get(name)?.definition;
			if (definition) return definition;
		}
		throw new Error('Missing tool ' + name);
	};
	const ctx = session.extensionRunner.createToolContext('identity-probe', undefined);
	core.setResponses([ai.fauxAssistantMessage('Verified Angel child.')]);
	const angel = await tool('angel').execute('angel-probe', { question: 'Verify child initialization' }, undefined, undefined, ctx);
	core.setResponses([
		ai.fauxAssistantMessage(ai.fauxToolCall('nu', { command: '1 + 1' }), { stopReason: 'toolUse' }),
		ai.fauxAssistantMessage(ai.fauxToolCall('structured_output', { verified: true }), { stopReason: 'toolUse' })
	]);
	const ultra = await tool('run_workflow').execute('ultra-probe', { spec: { name: 'identity', phases: [{
		id: 'child', kind: 'single', step: { summary: 'Verify child', prompt: 'Return verified true',
			model: 'identity/probe', tools: ['nu'], flags: { nushell: true }, schema: { type: 'object', properties: { verified: { type: 'boolean' } }, required: ['verified'], additionalProperties: false } }
	}] } }, undefined, undefined, ctx);
	console.log(JSON.stringify({ records: globalThis[Symbol.for('pi-ext-sdk-probe')], catalogFactories,
		angel: angel.details, ultra: JSON.parse(ultra.content[0].text), errors,
		warnings: [...(globalThis[Symbol.for('@bugabinga/pi-ext-pfui/state')]?.warnings ?? [])] }));
} finally {
	unsubscribe();
	await session.extensionRunner.emit({ type: 'session_shutdown', reason: 'quit' });
	session.dispose();
}
`,
		);
		const child = spawnSync(
			process.execPath,
			[
				"--import",
				`data:text/javascript,${encodeURIComponent(hook)}`,
				"--experimental-import-meta-resolve",
				driver,
				JSON.stringify({
					hostDir: host.dir,
					entries: entry,
					observer,
					agentDir,
				}),
			],
			{
				cwd: sandbox,
				env: sandboxEnvironment(sandbox, agentDir),
				encoding: "utf8",
				timeout: 30_000,
				maxBuffer: 1024 * 1024,
			},
		);
		expect(child.error, child.stderr).toBeUndefined();
		expect(child.status, child.stderr).toBe(0);
		const report = JSON.parse(child.stdout.trim()) as {
			records: {
				stage: string;
				id?: string;
				version: string;
				packageDir: string;
				contextVersion?: string;
				contextPackageDir?: string;
				activeTools?: string[];
				nuOwner?: string;
				hasNuSession?: boolean;
				isError?: boolean;
				content?: { type: string; text?: string }[];
			}[];
			catalogFactories: number;
			angel: {
				kind: string;
				result: { advice: string; metadata: { childSessionId: string } };
			};
			ultra: { executionStatus: string; result: { verified: boolean }[] };
			errors: string[];
			warnings: string[];
		};
		const sdkLoads = child.stderr
			.split("\n")
			.filter((line) => line.startsWith("PFUI_SDK_LOAD="))
			.map((line) => fileURLToPath(line.slice("PFUI_SDK_LOAD=".length)));
		return { report, sdkLoads };
	} finally {
		rmSync(sandbox, { recursive: true, force: true });
	}
}

// Build fresh artifacts; stale ignored dist must not make this test pass.
// Every relocated package gets workspace dependencies adjacent to it, so an
// accidental native .js import loads a second workspace SDK rather than hiding.
let builds: string;
let nativeControl: string;
let nativeAngelControl: string[];
const entries = [
	[
		join(root, "extensions/pfui/index.ts"),
		join(root, "extensions/angel/index.ts"),
		join(root, "extensions/ultra/index.ts"),
		join(root, "extensions/nushell/index.ts"),
	],
];
beforeAll(async () => {
	builds = mkdtempSync(join(tmpdir(), "pfui-runtime-builds-"));
	for (const release of [false, true]) {
		const packageEntries: string[] = [];
		for (const name of ["pfui", "angel", "ultra", "nushell"]) {
			const built = await buildExtension(name, {
				distRoot: join(builds, release ? "release" : "dev"),
				release,
			});
			symlinkSync(
				join(root, "node_modules"),
				join(built.dir, "node_modules"),
				"junction",
			);
			const manifest = JSON.parse(
				readFileSync(join(built.dir, "package.json"), "utf8"),
			) as {
				pi: { extensions: string[] };
			};
			expect(manifest.pi.extensions).toEqual(["index.ts"]);
			expect(forbiddenImports(built.dir)).toEqual([]);
			const entry = join(built.dir, manifest.pi.extensions[0]!);
			packageEntries.push(entry);
			if (!release && name === "pfui") {
				nativeControl = join(built.dir, "native-control.js");
				writeFileSync(nativeControl, readFileSync(entry));
			}
			if (!release && name === "angel") {
				writeFileSync(
					join(built.dir, "native-control.js"),
					readFileSync(entry),
				);
			}
		}
		if (!release) {
			nativeAngelControl = [...packageEntries];
			nativeAngelControl[1] = packageEntries[1]!.replace(
				/index\.ts$/u,
				"native-control.js",
			);
		}
		entries.push(packageEntries);
	}
}, 30_000);
afterAll(() => {
	if (builds) rmSync(builds, { recursive: true, force: true });
});

describe.each(hosts)("PFUI running Pi $version identity", (host) => {
	it("detects the duplicate SDK when the same emitted module uses a native JS suffix", () => {
		const result = probe(host, nativeControl);
		expect(result.sdkLoads, JSON.stringify(result)).toHaveLength(1);
		expect(
			fileURLToPath(result.sdkLoads[0]!.slice("PFUI_SDK_LOAD=".length)),
		).toBe(join(resolve(workspaceManifest, ".."), "dist/index.js"));
		expect(result.context.packageDir).toBe(resolve(workspaceManifest, ".."));
	});

	it.skipIf(host.dir === resolve(workspaceManifest, ".."))(
		"native generated Angel loads an adjacent SDK even when both installations have the same version",
		() => {
			const { sdkLoads } = compositionProbe(host, nativeAngelControl);
			expect(sdkLoads).toEqual([
				join(host.dir, "dist/index.js"),
				join(resolve(workspaceManifest, ".."), "dist/index.js"),
			]);
		},
	);

	it("source plus every relocated dev and release entry reuse the running CLI SDK", () => {
		for (const entry of entries) {
			const result = probe(host, entry);
			expect(result.context.version, JSON.stringify(result)).toBe(host.version);
			expect(result.context.packageDir, JSON.stringify(result)).toBe(host.dir);
			expect(result.sdkLoads, JSON.stringify(result)).toEqual([]);
		}
	});

	it("public SDK parent activates lazy runtimes, catalogs PFUI, and completes actual Angel/Ultra children", () => {
		for (const entry of entries) {
			const { report, sdkLoads } = compositionProbe(host, entry);
			expect(sdkLoads, JSON.stringify(report)).toEqual([
				join(host.dir, "dist/index.js"),
			]);
			expect(report.errors).toEqual([]);
			expect(report.warnings).toEqual([]);
			expect(report.catalogFactories).toBe(2);
			expect(report.angel.kind).toBe("complete");
			expect(report.angel.result.advice).toBe("Verified Angel child.");
			expect(report.ultra.executionStatus, JSON.stringify(report.ultra)).toBe(
				"finished",
			);
			expect(report.ultra.result).toEqual([{ verified: true }]);
			const sessions = report.records.filter(
				(record) => record.stage === "session",
			);
			expect(new Set(sessions.map((record) => record.id)).size).toBe(3);
			expect(sessions.map((record) => record.id)).toContain(
				report.angel.result.metadata.childSessionId,
			);
			const nuResult = report.records.filter(
				(record) => record.stage === "nu_result",
			);
			expect(nuResult, JSON.stringify(report)).toHaveLength(1);
			expect(nuResult[0]?.isError).toBe(false);
			expect(nuResult[0]?.content).toEqual(
				expect.arrayContaining([
					expect.objectContaining({
						type: "text",
						text: expect.stringMatching(/^2\s*$/),
					}),
				]),
			);
			const nuChild = sessions.find((record) =>
				record.activeTools?.includes("nu"),
			);
			expect(nuChild, JSON.stringify(report)).toBeDefined();
			expect(nuChild?.nuOwner).toBe(entry[3]);
			expect(nuChild?.hasNuSession).toBe(false);
			for (const record of report.records.filter(
				(record) => record.stage !== "nu_result",
			)) {
				expect(record.version).toBe(host.version);
				expect(record.packageDir).toBe(host.dir);
				if (record.stage === "session") {
					expect(record.contextVersion).toBe(host.version);
					expect(record.contextPackageDir).toBe(host.dir);
				}
			}
		}
	}, 120_000);
});