pi-ext

import {
	chmodSync,
	mkdirSync,
	mkdtempSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { afterEach, describe, expect, it } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

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

// Avoid Node startup latency inside the 100ms version-probe timeout.
// Rewrites answer after 300ms; the version probe optionally does too.
const FAKE_RTK = `#!/bin/sh
if [ "$1" = "--version" ]; then
	if [ -n "$FAKE_RTK_VERSION_DELAY" ]; then sleep "$FAKE_RTK_VERSION_DELAY"; fi
	printf '%s\\n' 'rtk 0.37.2'
else
	sleep 0.3
	printf '%s\\n' 'rtk git status'
fi
`;

type Case = {
	user?: unknown;
	project?: unknown;
	env?: string;
	untrusted?: boolean;
	versionDelay?: number;
};

// Recording UI context so the extension's notify calls are observable, capturing
// notify calls.
function recordingUI() {
	const calls: Array<{ args: [string, string | undefined] }> = [];
	const ui = {
		notify: (message: string, type?: "info" | "warning" | "error") => {
			calls.push({ args: [message, type] });
		},
	};
	return { ui, calls };
}

describe.skipIf(process.platform === "win32")(
	"rtk.rewriteTimeoutMs resolution",
	() => {
		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;
		});

		async function start(input: Case) {
			root = mkdtempSync(path.join(tmpdir(), "pi-ext-rtk-settings-"));
			const bin = path.join(root, "bin");
			mkdirSync(bin);
			writeFileSync(path.join(bin, "rtk"), FAKE_RTK);
			chmodSync(path.join(bin, "rtk"), 0o755);
			if ("project" in input) {
				const dir = path.join(root, ".pi");
				mkdirSync(dir, { recursive: true });
				writeFileSync(
					path.join(dir, "settings.json"),
					JSON.stringify({ rtk: { rewriteTimeoutMs: input.project } }),
				);
			}
			session = await createPiSession({
				cwd: root,
				extensionPaths: [extensionPath],
				...("user" in input
					? { userSettings: { rtk: { rewriteTimeoutMs: input.user } } }
					: {}),
				env: {
					PATH: `${bin}${path.delimiter}${process.env["PATH"] ?? ""}`,
					PI_RTK_REWRITE_TIMEOUT_MS: input.env,
					FAKE_RTK_VERSION_DELAY:
						input.versionDelay === undefined
							? undefined
							: String(input.versionDelay / 1000),
				},
			});
			expect(session.errors).toEqual([]);
			if (input.untrusted)
				session.session.settingsManager.setProjectTrusted(false);
			const runner = session.session.extensionRunner;
			const { ui, calls } = recordingUI();
			runner.setUIContext(ui as never, "tui");
			const before = calls.length;
			await runner.emit({ type: "session_start", reason: "startup" });
			const event = {
				type: "tool_call",
				toolName: "bash",
				toolCallId: "1",
				input: { command: "git status" },
			};
			await runner.emitToolCall(event as never);
			const warnings = calls
				.slice(before)
				.filter((call) => call.args[1] === "warning")
				.map((call) => call.args[0]);
			return { rewritten: event.input.command !== "git status", warnings };
		}

		it.each([
			["default waits for rewrite", {}, true],
			["user timeout applied", { user: 100 }, false],
			["env over user", { user: 100, env: "5000" }, true],
			["trusted project over user", { user: 5000, project: 100 }, false],
			[
				"untrusted project ignored",
				{ user: 5000, project: 100, untrusted: true },
				true,
			],
		] satisfies [string, Case, boolean][])(
			"%s",
			async (_name, input, expected) => {
				const { rewritten, warnings } = await start(input);
				expect(rewritten).toBe(expected);
				expect(warnings).toEqual([]);
			},
		);

		it("applies the timeout to the version probe", async () => {
			const { rewritten, warnings } = await start({
				user: 100,
				versionDelay: 300,
			});
			expect(rewritten).toBe(false);
			// A killed probe reports empty output, which rtk treats as an unknown version.
			expect(warnings).toEqual([
				"rtk disabled: too old (unknown); need >= 0.23.0",
			]);
		});

		it.each([
			["string user", { user: "100" }, "user settings"],
			["zero project", { user: 100, project: 0 }, "project settings"],
			["fractional user", { user: 100.5 }, "user settings"],
			[
				"non-numeric env",
				{ user: 100, env: "fast" },
				"environment variable PI_RTK_REWRITE_TIMEOUT_MS",
			],
		] satisfies [string, Case, string][])(
			"rejects %s without falling through",
			async (_name, input, source) => {
				const { rewritten, warnings } = await start(input);
				// The 2000ms default outlasts the 300ms rewrite; a 100ms fallthrough would not.
				expect(rewritten).toBe(true);
				expect(warnings).toEqual([
					`Invalid rtk.rewriteTimeoutMs from ${source}: expected an integer from 1 to 2147483647; using 2000ms`,
				]);
			},
		);
	},
);