pi-ext

import {
	chmodSync,
	copyFileSync,
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { delimiter, join } from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import uvExtension, { __test } from "../index";

let t: TestSession | undefined;
let fakeBin: string | undefined;

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

describe("uv extension", () => {
	it("loads in real Pi extension runtime", async () => {
		t = await createTestSession({ extensionFactories: [uvExtension] });

		expect(t.session.extensionRunner.extensions).toHaveLength(1);
		const bash =
			t.session.extensionRunner.extensions[0].tools.get("bash").definition;
		expect(bash.description).toContain("routes PATH-based Python through uv");
		expect(bash.promptSnippet).toContain("Execute bash commands");
		expect(bash.promptGuidelines).toEqual([
			expect.stringContaining("Use bash to inspect PI_*"),
			expect.stringContaining("Use bash with uv instead of pip"),
		]);
	});

	it("records availability outcome in its debug sandbox when enabled", async () => {
		fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-empty-bin-"));
		vi.stubEnv("PATH", fakeBin);
		const extensionPath = join(import.meta.dirname, "../index.ts");
		t = await createTestSession({
			env: { PI_UV_DEBUG: "1" },
			extensions: [extensionPath],
		});
		const extension = t.session.extensionRunner.extensions[0];
		const context = t.session.extensionRunner.createCommandContext();
		for (const handler of extension.handlers.get("session_start"))
			await handler({ type: "session_start" }, context);
		const directory = join(t.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,
		);
		t = await createTestSession({
			env: { PATH: fakeBin },
			extensionFactories: [uvExtension],
		});

		expect(t.events.uiCallsFor("notify")).toEqual([]);
	});

	it("warns through the extension when PATH has no uv executable", async () => {
		fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-bin-"));
		t = await createTestSession({
			env: { PATH: fakeBin },
			extensionFactories: [uvExtension],
		});

		await vi.waitFor(() =>
			expect(t?.events.uiCallsFor("notify").at(-1)?.args).toEqual([
				"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();
	});

	it("prepends shell shims to Unix PATH without mutating input", () => {
		const ctx = {
			command: "python --version",
			cwd: "/repo",
			env: { PATH: "/usr/bin", TOKEN: "preserved" },
		};

		const rewritten = __test.prependInterceptedCommands(ctx, "linux");
		expect(rewritten).not.toBe(ctx);
		expect(rewritten.env.PATH).toBe(
			`${join(import.meta.dirname, "../intercepted-commands")}${delimiter}/usr/bin`,
		);
		expect(rewritten.env.TOKEN).toBe("preserved");
		expect(ctx.env.PATH).toBe("/usr/bin");
	});

	it("does not expose extensionless shell shims to Windows", () => {
		const ctx = {
			command: "python --version",
			cwd: "C:/repo",
			env: { Path: "C:/Windows" },
		};

		expect(__test.prependInterceptedCommands(ctx, "win32")).toBe(ctx);
	});
});