pi-ext

import { fileURLToPath } from "node:url";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import {
	createCodemodeExtension,
	createToolSearchExtension,
	type ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "./fixture/create-pi-session.ts";

const entry = (name: string) =>
	fileURLToPath(new URL(`../extensions/${name}/index.ts`, import.meta.url));

const extensionPaths = [
	"ask",
	"mockup",
	"chrome-cdp",
	"firefox-bidi",
	"git-safe",
].map(entry);
const deferredNames = ["chrome_cdp", "firefox_bidi", "git_clone_safe"];

describe("custom tool discovery through Pi", () => {
	let session: PiSession | undefined;
	afterEach(async () => {
		await session?.dispose();
		session = undefined;
	});

	it("calls inactive browser tools from codemode without exposing interactive tools", async () => {
		const nested: { toolName: string; parentToolCallId: string | undefined }[] =
			[];
		const results: unknown[] = [];
		const observe = (pi: ExtensionAPI) => {
			pi.on("tool_call", (event) => {
				if (event.parentToolCallId)
					nested.push({
						toolName: event.toolName,
						parentToolCallId: event.parentToolCallId,
					});
			});
		};
		session = await createPiSession({
			extensionPaths,
			extensionFactories: [createCodemodeExtension(), observe],
			userSettings: { defaultTools: ["+codemode"] },
			env: { FIREFOX_BIN: "/nonexistent/firefox-bin" },
			responses: [
				fauxAssistantMessage(
					fauxToolCall(
						"codemode",
						{
							code: `
const found = await searchTools("git_clone_safe", { limit: 1 });
const closed = await tools.chrome_cdp({ operations: [{ kind: "close" }] });
let firefoxError;
try { await tools.firefox_bidi({ frames: [{ method: "session.status" }] }); }
catch (error) { firefoxError = error.message; }
return { found: found.map(tool => tool.name), closed, firefoxError,
  ask: "ask_user_question" in tools, mockup: "mockup" in tools };
`,
						},
						{ id: "script-1" },
					),
					{ stopReason: "toolUse" },
				),
				fauxAssistantMessage("done"),
			],
		});
		for (const name of deferredNames) {
			expect(session.session.getActiveToolNames()).not.toContain(name);
			expect(session.session.getCallableToolNames()).toContain(name);
		}
		for (const name of ["ask_user_question", "mockup"])
			expect(session.session.getCallableToolNames()).not.toContain(name);
		const unsubscribe = session.session.subscribe((event) => {
			if (event.type === "tool_execution_end" && event.toolName === "codemode")
				results.push(event.result);
		});
		await session.session.prompt(
			"Use scripts without manually activating browser tools.",
		);
		unsubscribe();
		expect(session.errors).toEqual([]);
		expect(nested).toEqual([
			{ toolName: "chrome_cdp", parentToolCallId: "script-1" },
			{ toolName: "firefox_bidi", parentToolCallId: "script-1" },
		]);
		const output = JSON.stringify(results);
		expect(output).toContain("Script completed");
		expect(output).toContain("git_clone_safe");
		expect(output).toContain("firefox not found");
		expect(output).toContain('\\"ask\\":false');
		expect(output).toContain('\\"mockup\\":false');
		for (const name of deferredNames)
			expect(session.session.getActiveToolNames()).not.toContain(name);
	});

	it("rejects named programmatic calls to model-only tools as error outcomes", async () => {
		const outcomes: { name: string; isError: boolean; result: unknown }[] = [];
		const probe = (pi: ExtensionAPI) => {
			pi.registerTool({
				name: "nested_probe",
				label: "Nested probe",
				description: "Probe model-only tool reachability.",
				parameters: Type.Object({}),
				async execute(_id, _args, _signal, _update, ctx) {
					for (const name of ["ask_user_question", "mockup"]) {
						const outcome = await ctx.executeTool(name, {});
						outcomes.push({
							name,
							isError: outcome.isError,
							result: outcome.result,
						});
					}
					return {
						content: [{ type: "text", text: "done" }],
						details: undefined,
					};
				},
			});
		};
		session = await createPiSession({
			extensionPaths,
			extensionFactories: [probe],
			responses: [
				fauxAssistantMessage(fauxToolCall("nested_probe", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage("done"),
			],
		});
		await session.session.prompt("Probe nested tool visibility.");
		expect(outcomes).toHaveLength(2);
		for (const outcome of outcomes) {
			expect(outcome.isError).toBe(true);
			expect(outcome.result).toMatchObject({
				content: [{ type: "text", text: `Tool ${outcome.name} not found` }],
			});
		}
	});

	it("warns once per session when deferred tools have no discovery route", async () => {
		session = await createPiSession({
			extensionPaths,
			responses: [fauxAssistantMessage("done"), fauxAssistantMessage("done")],
		});
		const notify = vi.spyOn(
			session.session.extensionRunner.createCommandContext().ui,
			"notify",
		);
		await session.session.prompt("Inspect available capabilities.");
		await session.session.prompt("Inspect them again.");
		const warnings = notify.mock.calls.filter(
			([message, level]) =>
				level === "warning" && /deferred|tool search/.test(String(message)),
		);
		expect(warnings).toHaveLength(3);
		for (const name of deferredNames)
			expect(warnings.some(([message]) => String(message).includes(name))).toBe(
				true,
			);
	});

	it.each(["codemode", "tool_search"])(
		"does not warn with active %s",
		async (name) => {
			session = await createPiSession({
				extensionPaths,
				extensionFactories: [
					createCodemodeExtension(),
					createToolSearchExtension(),
				],
				userSettings: { defaultTools: [`+${name}`] },
				responses: [fauxAssistantMessage("done")],
			});
			const notify = vi.spyOn(
				session.session.extensionRunner.createCommandContext().ui,
				"notify",
			);
			await session.session.prompt("Inspect available capabilities.");
			expect(
				notify.mock.calls.filter(
					([message, level]) =>
						level === "warning" && /deferred|tool search/.test(String(message)),
				),
			).toEqual([]);
		},
	);

	it.each(deferredNames)(
		"tool_search discovers and declares %s",
		async (name) => {
			session = await createPiSession({
				extensionPaths,
				extensionFactories: [createToolSearchExtension()],
				userSettings: { defaultTools: ["+tool_search"] },
				responses: [
					fauxAssistantMessage(
						fauxToolCall("tool_search", { query: name, limit: 1 }),
						{ stopReason: "toolUse" },
					),
					fauxAssistantMessage("done"),
				],
			});
			expect(session.session.getActiveToolNames()).not.toContain(name);
			await session.session.prompt(`Discover ${name}.`);
			expect(session.errors).toEqual([]);
			expect(session.session.getActiveToolNames()).toContain(name);
		},
	);
});