pi-ext

import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import {
	type ExtensionAPI,
	initTheme,
	type ToolDefinition,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

type ToolRenderContext = Parameters<
	NonNullable<ToolDefinition["renderCall"]>
>[2];

import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { findNu } from "../runtime.ts";

const extensionPath = fileURLToPath(new URL("../index.ts", import.meta.url));

describe("nushell contract", () => {
	let session: PiSession | undefined;
	let root: string | undefined;
	afterEach(async () => {
		await session?.dispose();
		session = undefined;
		if (root) rmSync(root, { recursive: true, force: true });
		root = undefined;
		vi.restoreAllMocks();
	});

	it("registers inert tools and toggles without changing unrelated tools", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		const runner = session.session.extensionRunner;
		const command = runner.getCommand("nushell");
		const ctx = runner.createCommandContext();
		expect(session.extensions[0]?.tools.has("nu")).toBe(true);
		expect(session.extensions[0]?.tools.has("nu_session")).toBe(false);
		expect(
			session.extensions[0]?.tools.get("nu")?.definition.annotations
				?.readOnlyHint,
		).toBe(false);
		expect(session.session.getActiveToolNames()).not.toContain("nu");
		expect(command?.getArgumentCompletions?.("o")).toEqual([
			{ value: "on", label: "on" },
			{ value: "off", label: "off" },
		]);
		expect(command?.getArgumentCompletions?.("re")).toBeNull();
		const initial = session.session.getActiveToolNames();
		expect(
			session.session.getAllTools().find((tool) => tool.name === "nu")
				?.exposure,
		).toBe("direct");
		expect(
			session.session.getAllTools().map((tool) => tool.name),
		).not.toContain("nu_session");
		await command?.handler("on", ctx);
		const enabled = session.session.getActiveToolNames();
		if (!findNu()) {
			expect(enabled).toEqual(initial);
			return;
		}
		expect(enabled).toContain("nu");
		expect(enabled).not.toContain("nu_session");
		expect(enabled).not.toContain("bash");
		expect(enabled).not.toContain("powershell");
		for (const name of initial.filter(
			(name) => name !== "bash" && name !== "powershell",
		))
			expect(enabled).toContain(name);
		await command?.handler("off", ctx);
		expect([...session.session.getActiveToolNames()].sort()).toEqual(
			[...initial].sort(),
		);
	});

	it("keeps shell tools active and warns when PATH has no executable Nu", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		root = mkdtempSync(path.join(tmpdir(), "pi-ext-nushell-path-"));
		mkdirSync(path.join(root, process.platform === "win32" ? "nu.exe" : "nu"));
		session.env["PATH"] = root;
		process.env["PATH"] = root;
		const runner = session.session.extensionRunner;
		const notify = vi.spyOn(runner.createCommandContext().ui, "notify");
		const initial = session.session.getActiveToolNames();
		runner.setFlagValue("nushell", true);
		await runner.emit({ type: "session_start", reason: "startup" });
		expect(session.session.getActiveToolNames()).toEqual(initial);
		expect(notify.mock.calls.at(-1)).toEqual([
			"Nushell executable not found on PATH; tools unchanged",
			"warning",
		]);
		await runner
			.getCommand("nushell")
			?.handler("on", runner.createCommandContext());
		expect(session.session.getActiveToolNames()).toEqual(initial);
		expect(notify.mock.calls.at(-1)).toEqual([
			"Nushell executable not found on PATH; tools unchanged",
			"warning",
		]);
	});

	it("registers --nushell to enable fresh Nu only", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		const runner = session.session.extensionRunner;
		expect(runner.getFlags().has("nushell")).toBe(true);
		runner.setFlagValue("nushell", true);
		await runner.emit({ type: "session_start", reason: "startup" });
		const active = session.session.getActiveToolNames();
		if (!findNu()) {
			expect(active).not.toContain("nu");
			return;
		}
		expect(active).toContain("nu");
		expect(active).not.toContain("nu_session");
		expect(active).not.toContain("bash");
		expect(active).not.toContain("powershell");
	});

	it.runIf(Boolean(findNu()))(
		"does not claim activation when direct Nu is excluded",
		async () => {
			session = await createPiSession({ extensionPaths: [extensionPath] });
			const runner = session.session.extensionRunner;
			const notify = vi.spyOn(runner.createCommandContext().ui, "notify");
			const initial = session.session.getActiveToolNames();
			const tools = session.session.getAllTools();
			vi.spyOn(session.session, "getAllTools").mockReturnValue(
				tools.filter((tool) => tool.name !== "nu"),
			);
			runner.setFlagValue("nushell", true);
			await runner.emit({ type: "session_start", reason: "startup" });
			expect(session.session.getActiveToolNames()).toEqual(initial);
			expect(notify.mock.calls.at(-1)).toEqual([
				expect.stringContaining("nu"),
				"warning",
			]);
		},
	);

	it.runIf(Boolean(findNu()))(
		"renders fresh Nu with syntax highlighting and shell-style results",
		async () => {
			// The syntax highlighter spawns `nu` in Pi's own cwd (the repo root),
			// where a mise shim would refuse the untrusted mise.toml under the
			// sandbox HOME. Trusting the repo config restores the real behavior and
			// is inert when `nu` is a plain executable.
			session = await createPiSession({
				extensionPaths: [extensionPath],
				env: { MISE_TRUSTED_CONFIG_PATHS: process.cwd() },
			});
			initTheme("default", false);
			const tools = session.extensions[0]?.tools;
			const theme = {
				fg: (color: string, value: string) =>
					value ? `<${color}>${value}</${color}>` : "",
				bold: (value: string) => value,
			} as never;
			const tool = tools?.get("nu")?.definition;
			const icon = "";
			expect(tool?.label).toBe(`${icon} nu`);
			const context = {
				state: {},
				executionStarted: false,
				argsComplete: false,
				invalidate: vi.fn(),
			} as unknown as ToolRenderContext;

			const partial = tool?.renderCall?.(
				{ command: "let value" },
				theme,
				context,
			);
			expect(partial?.render(1000).join("\n")).not.toContain(
				"<syntaxVariable>",
			);
			context.argsComplete = true;
			context.executionStarted = true;
			const command = 'let value = "safe"';
			const call = tool?.renderCall?.({ command }, theme, context);
			const lines = call?.render(1000).join("\n") ?? "";
			expect(lines.replace(/<\/?[a-zA-Z]+>/g, "")).toContain(command);
			expect(lines).not.toContain("<syntaxVariable>");
			await vi.waitFor(() =>
				expect(context.invalidate).toHaveBeenCalledTimes(1),
			);
			const highlighted =
				tool
					?.renderCall?.({ command }, theme, context)
					.render(1000)
					.join("\n") ?? "";
			expect(highlighted).toContain("<syntaxVariable>value</syntaxVariable>");
			expect(highlighted).toContain('<syntaxString>"safe"</syntaxString>');
			expect(highlighted).not.toMatch(/\x1b\[/);
			const paletteCommand =
				'if true { let value = "safe"; $value | str uppercase; 2 > 1 }';
			tool?.renderCall?.({ command: paletteCommand }, theme, context);
			await vi.waitFor(() =>
				expect(context.invalidate).toHaveBeenCalledTimes(2),
			);
			const palette =
				tool
					?.renderCall?.({ command: paletteCommand }, theme, context)
					.render(2000)
					.join("\n") ?? "";
			for (const token of [
				"<syntaxFunction>if</syntaxFunction>",
				"<syntaxKeyword>true</syntaxKeyword>",
				"<syntaxPunctuation>{</syntaxPunctuation>",
				"<syntaxVariable>$value</syntaxVariable>",
				"<syntaxPunctuation>|</syntaxPunctuation>",
				"<syntaxOperator>></syntaxOperator>",
				'<syntaxString>"safe"</syntaxString>',
			])
				expect(palette).toContain(token);
			expect(lines).toContain(`${icon} nu>`);
			const result = tool?.renderResult?.(
				{
					content: [{ type: "text", text: "one\ntwo\nthree\nfour\nfive\nsix" }],
					details: undefined,
				},
				{ expanded: false, isPartial: false },
				theme,
				{ ...context, showImages: false, isError: false },
			);
			expect(result?.render(120).join("\n")).toContain("to expand");
		},
	);

	it.runIf(Boolean(findNu()))(
		"keeps fresh Nu callable through a real nested Pi tool call",
		async () => {
			const calls: Array<{
				toolName: string;
				parentToolCallId: string | undefined;
			}> = [];
			const nestedProbe = (pi: ExtensionAPI) => {
				pi.on("tool_call", (event) => {
					calls.push({
						toolName: event.toolName,
						parentToolCallId: event.parentToolCallId,
					});
				});
				pi.registerTool({
					name: "nu_nested_probe",
					label: "Nu nested probe",
					description: "Run a fresh Nushell command.",
					parameters: Type.Object({ command: Type.String() }),
					async execute(_id, { command }, signal, _onUpdate, ctx) {
						const nested = await ctx.executeTool(
							"nu",
							{ command },
							{
								...(signal ? { signal } : {}),
							},
						);
						return nested.result;
					},
				});
			};
			session = await createPiSession({
				extensionPaths: [extensionPath],
				extensionFactories: [nestedProbe],
				responses: [
					fauxAssistantMessage(
						fauxToolCall(
							"nu_nested_probe",
							{ command: "1 + 1" },
							{ id: "probe-1" },
						),
						{ stopReason: "toolUse" },
					),
					fauxAssistantMessage("done"),
				],
			});
			const runner = session.session.extensionRunner;
			await runner
				.getCommand("nushell")
				?.handler("on", runner.createCommandContext());
			await session.session.prompt("run fresh Nu");
			expect(session.errors).toEqual([]);
			expect(calls).toContainEqual(
				expect.objectContaining({
					toolName: "nu",
					parentToolCallId: "probe-1",
				}),
			);
		},
	);

	it.runIf(Boolean(findNu()))(
		"isolates variables, functions, environment, and cwd across fresh calls",
		async () => {
			root = mkdtempSync(path.join(tmpdir(), "pi-ext-nushell-isolation-"));
			const subdirectory = path.join(root, "subdirectory");
			mkdirSync(subdirectory);
			session = await createPiSession({
				cwd: root,
				extensionPaths: [extensionPath],
			});
			const runner = session.session.extensionRunner;
			const tool = session.extensions[0]?.tools.get("nu")?.definition;
			const run = (command: string) =>
				tool?.execute(
					"call",
					{ command },
					undefined,
					undefined,
					runner.createToolContext("call", undefined),
				);
			const first = await run(
				'let value = 41; def private_test [] { 42 }; $env.PI_NUSHELL_TEST = "set"; cd subdirectory; pwd',
			);
			expect(first?.content[0]).toMatchObject({
				type: "text",
				text: expect.stringContaining(subdirectory),
			});
			const second = await run(
				"[(scope variables | where name == '$value' | length), (scope commands | where name == 'private_test' | length), $env.PI_NUSHELL_TEST?, $env.PWD] | to json --raw",
			);
			const secondText = second?.content[0];
			const text = secondText?.type === "text" ? secondText.text : "";
			expect(JSON.parse(text)).toEqual([0, 0, null, root]);
		},
	);
});