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(); }); });