pi-ext

import {
	chmodSync,
	copyFileSync,
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { initTheme } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { __test } from "../index.ts";

function extensionPath(): string {
	return fileURLToPath(new URL("../index.ts", import.meta.url));
}

let session: PiSession | undefined;
let fakeBin: string | undefined;

afterEach(async () => {
	await session?.dispose();
	if (fakeBin) rmSync(fakeBin, { recursive: true, force: true });
	vi.unstubAllEnvs();
	session = undefined;
	fakeBin = undefined;
});

describe("uv extension", () => {
	it("loads in real Pi extension runtime without owning a shell tool", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });

		expect(session.extensions).toHaveLength(1);
		expect(session.extensions[0]!.tools.size).toBe(0);
		expect(session.extensions[0]!.handlers.get("tool_call")).toHaveLength(1);
	});

	it("blocks a shell tool_call running pip and leaves other commands alone", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const extension = session.extensions[0]!;
		const [handler] = extension.handlers.get("tool_call") ?? [];

		const blocked = (await handler!({
			type: "tool_call",
			toolCallId: "1",
			toolName: "bash",
			input: { command: "pip install rich" },
		})) as { block?: boolean; reason?: string } | undefined;
		expect(blocked?.block).toBe(true);
		expect(blocked?.reason).toContain("uv add PACKAGE");

		const allowed = await handler!({
			type: "tool_call",
			toolCallId: "2",
			toolName: "bash",
			input: { command: "uv run python script.py" },
		});
		expect(allowed).toBeUndefined();
	});

	it("records availability outcome in its debug sandbox when enabled", async () => {
		fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-empty-bin-"));
		vi.stubEnv("PATH", fakeBin);
		session = await createPiSession({
			env: { PI_UV_DEBUG: "1" },
			extensionPaths: [extensionPath()],
		});
		const extension = session.extensions[0]!;
		const context = session.session.extensionRunner.createCommandContext();
		for (const handler of extension.handlers.get("session_start")!)
			await handler({ type: "session_start" }, context);
		const directory = join(session.cwd, ".test-state", "pi-ext", "debug", "uv");
		await vi.waitFor(() => {
			const lines = readdirSync(directory).flatMap((file) =>
				readFileSync(join(directory, file), "utf8").trim().split("\n"),
			);
			expect(
				lines.some((line) => line.includes("availability.check.finish")),
			).toBe(true);
		});
		for (const handler of extension.handlers.get("session_shutdown")!)
			await handler({ type: "session_shutdown" }, context);
		const records = readFileSync(
			join(directory, readdirSync(directory)[0]!),
			"utf8",
		)
			.trim()
			.split("\n")
			.map((line) => JSON.parse(line));
		expect(records.map((record) => record.event)).toEqual(
			expect.arrayContaining([
				"session.start",
				"session.shutdown",
				"availability.check.finish",
			]),
		);
		expect(records).toContainEqual(
			expect.objectContaining({
				event: "availability.check.finish",
				installed: false,
			}),
		);
	});

	it("detects a controlled executable on PATH through the extension", async () => {
		fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-bin-"));
		if (process.platform === "win32") {
			copyFileSync(process.execPath, join(fakeBin, "uv.exe"));
		} else {
			const uv = join(fakeBin, "uv");
			writeFileSync(uv, "#!/bin/sh\nexit 0\n", "utf8");
			chmodSync(uv, 0o755);
		}

		expect(await __test.uvInstalled({ PATH: fakeBin }, process.platform)).toBe(
			true,
		);
		session = await createPiSession({
			env: { PATH: fakeBin },
			extensionPaths: [extensionPath()],
		});
		const extension = session.extensions[0]!;
		const context = session.session.extensionRunner.createContext();
		Object.defineProperty(context, "hasUI", { value: true });
		initTheme(undefined, false);
		const notify = vi.spyOn(context.ui, "notify");
		for (const handler of extension.handlers.get("session_start")!)
			await handler({ type: "session_start" }, context);
		await new Promise((resolve) => setImmediate(resolve));

		expect(notify).not.toHaveBeenCalled();
	});

	it("warns through the extension when PATH has no uv executable", async () => {
		fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-bin-"));
		session = await createPiSession({
			env: { PATH: fakeBin },
			extensionPaths: [extensionPath()],
		});
		const extension = session.extensions[0]!;
		const context = session.session.extensionRunner.createContext();
		Object.defineProperty(context, "hasUI", { value: true });
		initTheme(undefined, false);
		const notify = vi.spyOn(context.ui, "notify");
		for (const handler of extension.handlers.get("session_start")!)
			await handler({ type: "session_start" }, context);

		await vi.waitFor(() =>
			expect(notify).toHaveBeenCalledWith(
				"uv extension: uv not found on PATH",
				"warning",
			),
		);
	});

	it("blocks package-manager commands with uv guidance", () => {
		expect(__test.getBlockedCommandMessage("pip install rich")).toContain(
			"uv add PACKAGE",
		);
		expect(__test.getBlockedCommandMessage("pip3 install rich")).toContain(
			"pip3 is disabled",
		);
		expect(__test.getBlockedCommandMessage("poetry add rich")).toContain(
			"uv sync",
		);
		expect(
			__test.getBlockedCommandMessage("python -m pip install rich"),
		).toContain("python -m pip");
	});

	it("blocks venv and py_compile bytecode writes", () => {
		expect(__test.getBlockedCommandMessage("python3 -m venv .venv")).toContain(
			"uv venv",
		);
		expect(
			__test.getBlockedCommandMessage(".venv/bin/python -m py_compile x.py"),
		).toContain("__pycache__");
	});

	it("leaves normal Python execution alone", () => {
		expect(__test.getBlockedCommandMessage("python script.py")).toBeNull();
		expect(
			__test.getBlockedCommandMessage("uv run python script.py"),
		).toBeNull();
	});
});