pi-ext

import { chmod, mkdtemp, readFile, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import {
	afterAll,
	afterEach,
	beforeAll,
	describe,
	expect,
	it,
	vi,
} from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import { fetchUrlStream } from "../fetch";
import webExtension from "../index";
import { filterLineMatches } from "../line-match";
import {
	discoverSiteMap,
	parseLlmsLinks,
	parseRobotsSitemaps,
	parseSitemapXml,
} from "../map";
import {
	assertSafeHttpUrl,
	isPrivateOrReservedIp,
	resolveRedirectUrl,
} from "../safety";
import { searchWeb } from "../search";
import { spawnToText } from "../spawn";
import { platformStateDir, projectStoreDir, saveJsonResult } from "../store";

const originalFetch = globalThis.fetch;

// Restore `fetch` after EVERY test in this file. Several tests below swap
// `globalThis.fetch` for a stub, and they live in two sibling `describe`
// blocks — a describe-scoped hook only covers its own block, so the last stub
// set in the second block would otherwise never be restored. Under vitest's
// `isolate: false` that leaked stub bleeds into later test files (e.g.
// diff-review's SSR server tests), which then see the stub's `text/plain` body
// instead of real responses. A file-level hook guarantees the global is always
// handed back clean, regardless of which block (or future block) mutated it.
afterEach(() => {
	globalThis.fetch = originalFetch;
});

async function makeFakeCommand(
	dir: string,
	name: string,
	stdout: string,
): Promise<void> {
	if (process.platform === "win32") {
		await writeFile(
			path.join(dir, `${name}.cmd`),
			`@echo off\r\necho ${stdout}\r\n`,
		);
		return;
	}
	await writeFile(
		path.join(dir, name),
		`#!/bin/sh\nprintf '%s\\n' '${stdout}'\n`,
		{
			mode: 0o755,
		},
	);
}

async function makeEarlyExitCommand(dir: string, name: string): Promise<void> {
	if (process.platform === "win32") {
		await writeFile(
			path.join(dir, `${name}.cmd`),
			"@echo off\r\nexit /b 0\r\n",
		);
		return;
	}
	const file = path.join(dir, name);
	await writeFile(file, "#!/bin/sh\nexit 0\n");
	await chmod(file, 0o755);
}

function pathEnv(dir: string): Record<string, string> {
	return process.platform === "win32"
		? { PATH: dir, Path: dir, PATHEXT: ".COM;.EXE;.BAT;.CMD" }
		: { PATH: dir };
}

describe("web fetch/map context contract", () => {
	let defaultSession: TestSession;
	let t: TestSession | undefined;

	beforeAll(async () => {
		defaultSession = await createTestSession({
			extensionFactories: [webExtension],
		});
	});

	afterEach(() => {
		t?.dispose();
		t = undefined;
	});

	afterAll(() => defaultSession.dispose());

	it("stores full output but returns only capped visible text and metadata", async () => {
		globalThis.fetch = vi.fn(
			async () =>
				new Response("x".repeat(50_000), {
					status: 200,
					headers: { "content-type": "text/plain" },
				}),
		) as typeof fetch;

		const [extension] = defaultSession.session.extensionRunner.extensions;
		const tool = extension.tools.get("web").definition;
		const result = await tool.execute(
			"test",
			{ action: "fetch", url: "https://93.184.216.34/plain", maxChars: 1_000 },
			new AbortController().signal,
			undefined,
			{ cwd: defaultSession.cwd },
		);

		const text = result.content[0].text;
		expect(text.length).toBeLessThan(1_800);
		expect(text).toContain("responseId:");
		expect(result.details.responseId).toBeTypeOf("string");
		expect(result.details.fullOutputPath).toBeTypeOf("string");
		expect(text).toContain(`fullOutputPath: ${result.details.fullOutputPath}`);
		expect(JSON.stringify(result.details)).not.toContain("x".repeat(200));
		const saved = JSON.parse(
			await readFile(result.details.fullOutputPath, "utf8"),
		);
		expect(saved.output).toContain("x".repeat(10_000));
	});

	it("handles commands that exit before stdin is written", async () => {
		const binDir = await mkdtemp(path.join(tmpdir(), "pi-web-early-exit-"));
		await makeEarlyExitCommand(binDir, "early");
		const oldPath = process.env.PATH;
		process.env.PATH = binDir;
		try {
			await expect(
				spawnToText("early", [], "x".repeat(10_000_000)),
			).resolves.toBe("");
		} finally {
			process.env.PATH = oldPath;
		}
	});

	it("converts HTML with PATH-resolved pandoc and no Unix which", async () => {
		globalThis.fetch = vi.fn(
			async () =>
				new Response("<html><main><h1>Hello</h1></main></html>", {
					status: 200,
					headers: { "content-type": "text/html" },
				}),
		) as typeof fetch;

		const binDir = await mkdtemp(path.join(tmpdir(), "pi-web-bin-"));
		await makeFakeCommand(binDir, "pandoc", "converted markdown");
		t = await createTestSession({
			extensionFactories: [webExtension],
			env: pathEnv(binDir),
		});
		const [extension] = t.session.extensionRunner.extensions;
		const result = await extension.tools
			.get("web")
			.definition.execute(
				"test",
				{ action: "fetch", url: "https://93.184.216.34/html" },
				new AbortController().signal,
				undefined,
				{ cwd: t.cwd },
			);

		expect(result.isError).not.toBe(true);
		expect(result.details.kind).toBe("html");
		expect(result.content[0].text).toContain("converted markdown");
		expect(result.content[0].text).not.toContain("fullOutputPath:");
	});

	it("notifies missing required platform tools only on session start", async () => {
		const emptyDir = await mkdtemp(path.join(tmpdir(), "pi-web-empty-bin-"));
		t = await createTestSession({
			extensionFactories: [webExtension],
			env: pathEnv(emptyDir),
		});
		await vi.waitFor(() => {
			const calls = t?.events.uiCallsFor("notify").map((call) => call.args);
			expect(calls).toEqual(
				expect.arrayContaining([
					expect.arrayContaining([
						expect.stringContaining("pandoc"),
						"warning",
					]),
					expect.arrayContaining([
						expect.stringContaining("pdftotext"),
						"warning",
					]),
				]),
			);
			expect(calls).not.toEqual(
				expect.arrayContaining([
					expect.arrayContaining([expect.stringContaining("identify"), "info"]),
					expect.arrayContaining([expect.stringContaining("ffprobe"), "info"]),
				]),
			);
		});
	});

	it("throws search failures without duplicate notify", async () => {
		globalThis.fetch = vi.fn(
			async () => new Response("no", { status: 503, statusText: "offline" }),
		) as typeof fetch;
		const notify = vi.fn();

		await expect(
			searchWeb({
				query: "nope",
				count: 1,
				timeout: 50,
				notify,
			} as any),
		).rejects.toThrow("Search unavailable");
		expect(notify).not.toHaveBeenCalled();
	});

	it("line matching returns excerpts instead of full page", async () => {
		globalThis.fetch = vi.fn(
			async () =>
				new Response(
					["one", "before", "install with mise", "after", "tail"].join("\n"),
					{ status: 200, headers: { "content-type": "text/plain" } },
				),
		) as typeof fetch;

		const [extension] = defaultSession.session.extensionRunner.extensions;
		const result = await extension.tools.get("web").definition.execute(
			"test",
			{
				action: "fetch",
				url: "https://93.184.216.34/plain",
				linesMatching: ["install"],
				contextLines: 1,
			},
			new AbortController().signal,
			undefined,
			{ cwd: defaultSession.cwd },
		);

		const text = result.content[0].text;
		expect(text).toContain("2: before");
		expect(text).toContain("3: install with mise");
		expect(text).toContain("4: after");
		expect(text).not.toContain("1: one");
		expect(text).not.toContain("5: tail");
		expect(text).toContain(`fullOutputPath: ${result.details.fullOutputPath}`);
		expect(
			JSON.parse(await readFile(result.details.fullOutputPath, "utf8")).output,
		).toContain("tail");
	});

	it("provides saved-output recovery when no lines match", async () => {
		globalThis.fetch = vi.fn(
			async () =>
				new Response("complete page body", {
					status: 200,
					headers: { "content-type": "text/plain" },
				}),
		) as typeof fetch;

		const [extension] = defaultSession.session.extensionRunner.extensions;
		const result = await extension.tools.get("web").definition.execute(
			"test",
			{
				action: "fetch",
				url: "https://93.184.216.34/plain",
				linesMatching: ["missing"],
			},
			new AbortController().signal,
			undefined,
			{ cwd: defaultSession.cwd },
		);

		expect(result.content[0].text).toContain("No lines matched: missing");
		expect(result.content[0].text).toContain(
			`fullOutputPath: ${result.details.fullOutputPath}`,
		);
		expect(
			JSON.parse(await readFile(result.details.fullOutputPath, "utf8")).output,
		).toContain("complete page body");
	});

	it.each([
		{
			name: "entry count",
			urlCount: 60,
			maxChars: 8_000,
			notice: "at most 50 of 60",
		},
		{
			name: "character",
			urlCount: 10,
			maxChars: 100,
			notice: "capped at 100 characters",
		},
	])(
		"marks $name-limited web map listings as partial with saved recovery",
		async ({ urlCount, maxChars, notice }) => {
			const sitemap = Array.from(
				{ length: urlCount },
				(_, index) =>
					`<url><loc>https://example.com/documentation/${index.toString().padStart(3, "0")}</loc></url>`,
			).join("");
			globalThis.fetch = vi.fn(async (input) => {
				const url = String(input);
				return new Response(url.endsWith("/sitemap.xml") ? sitemap : "", {
					status: 200,
					headers: { "content-type": "text/plain" },
				});
			}) as typeof fetch;

			const [extension] = defaultSession.session.extensionRunner.extensions;
			const result = await extension.tools
				.get("web")
				.definition.execute(
					"test",
					{ action: "map", url: "https://93.184.216.34/docs", maxChars },
					new AbortController().signal,
					undefined,
					{ cwd: defaultSession.cwd },
				);

			expect(result.content[0].text).toContain("Partial URL listing:");
			expect(result.content[0].text).toContain(notice);
			expect(result.content[0].text).toContain(
				`fullOutputPath: ${result.details.fullOutputPath}`,
			);
			const saved = JSON.parse(
				await readFile(result.details.fullOutputPath, "utf8"),
			);
			expect(saved.urls).toHaveLength(urlCount);
		},
	);

	it("does not add a recovery notice to a complete web map listing", async () => {
		globalThis.fetch = vi.fn(async (input) => {
			const url = String(input);
			const sitemap =
				"<url><loc>https://example.com/a</loc></url><url><loc>https://example.com/b</loc></url>";
			return new Response(url.endsWith("/sitemap.xml") ? sitemap : "", {
				status: 200,
				headers: { "content-type": "text/plain" },
			});
		}) as typeof fetch;

		const [extension] = defaultSession.session.extensionRunner.extensions;
		const result = await extension.tools
			.get("web")
			.definition.execute(
				"test",
				{ action: "map", url: "https://93.184.216.34/docs" },
				new AbortController().signal,
				undefined,
				{ cwd: defaultSession.cwd },
			);

		expect(result.content[0].text).not.toContain("Partial URL listing:");
		expect(result.content[0].text).not.toContain("fullOutputPath:");
		expect(result.details.fullOutputPath).toBeTypeOf("string");
	});
});

describe("web safety/store/map primitives", () => {
	it("blocks unsafe URLs and DNS answers", async () => {
		await expect(assertSafeHttpUrl("file:///etc/passwd")).rejects.toThrow(
			/Only http: and https:/,
		);
		await expect(assertSafeHttpUrl("http://127.0.0.1")).rejects.toThrow(
			/private address/,
		);
		await expect(
			assertSafeHttpUrl("https://public.test", {
				resolver: async () => ["10.0.0.2"],
			}),
		).rejects.toThrow(/blocked private address/);
		await expect(
			resolveRedirectUrl("https://example.com/a", "http://127.0.0.1/b"),
		).rejects.toThrow(/private address/);
		expect(isPrivateOrReservedIp("8.8.8.8")).toBe(false);
	});

	it("fetch validates redirects and max bytes", async () => {
		globalThis.fetch = vi.fn(
			async () =>
				new Response(null, {
					status: 302,
					headers: { location: "http://127.0.0.1/private" },
				}),
		) as typeof fetch;
		await expect(
			fetchUrlStream("https://example.com", undefined, undefined, {
				resolver: async () => ["93.184.216.34"],
			}),
		).rejects.toThrow(/private address/);
		globalThis.fetch = vi.fn(
			async () =>
				new Response("x".repeat(20), {
					status: 200,
					headers: { "content-type": "text/plain" },
				}),
		) as typeof fetch;
		await expect(
			fetchUrlStream("https://example.com", undefined, undefined, {
				maxBytes: 10,
				resolver: async () => ["93.184.216.34"],
			}),
		).rejects.toThrow(/exceeds maxBytes/);
	});

	it("stores full JSON under XDG state dir", async () => {
		const root = await mkdtemp(path.join(tmpdir(), "pi-web-store-"));
		expect(
			platformStateDir("linux", { XDG_STATE_HOME: "/state" }, "/home/me"),
		).toBe("/state/pi");
		expect(
			platformStateDir(
				"win32",
				{
					XDG_STATE_HOME: "C:\\state",
					LOCALAPPDATA: "C:\\Users\\me\\AppData\\Local",
				},
				"C:\\Users\\me",
			),
		).toBe("C:\\state\\pi");
		expect(
			projectStoreDir("web", "/work/project", {
				platform: "linux",
				env: { XDG_STATE_HOME: "/state" },
				home: "/home/me",
			}),
		).toMatch(/^\/state\/pi\/web\/[a-f0-9]{64}$/);
		const saved = await saveJsonResult(
			"web",
			root,
			{ body: "x".repeat(50_000) },
			{ platform: "linux", env: { XDG_STATE_HOME: root }, home: root },
		);
		expect(
			JSON.parse(await readFile(saved.fullOutputPath, "utf8")).body,
		).toHaveLength(50_000);
		expect(JSON.stringify(saved)).not.toContain("xxxxx");
	});

	it("filters matching lines with context", () => {
		const result = filterLineMatches(
			["a", "before", "Install with mise", "after", "z"].join("\n"),
			{ needles: ["install"], contextLines: 1 },
		);
		expect(result.text).toContain("2: before");
		expect(result.text).toContain("3: Install with mise");
		expect(result.text).toContain("4: after");
		expect(filterLineMatches("huge\nbody", { needles: ["missing"] }).text).toBe(
			"",
		);
	});

	it("maps robots/sitemap/llms without page bodies", async () => {
		expect(
			parseRobotsSitemaps("Sitemap: https://x.test/sitemap.xml\n"),
		).toEqual(["https://x.test/sitemap.xml"]);
		expect(parseSitemapXml("<url><loc>https://x.test/a</loc></url>")).toEqual([
			"https://x.test/a",
		]);
		expect(
			parseLlmsLinks(
				"[Docs](https://x.test/docs) [Rel](/rel)",
				"https://x.test/base",
			),
		).toEqual(["https://x.test/docs", "https://x.test/rel"]);
		const fetched: string[] = [];
		const result = await discoverSiteMap("https://x.test/docs", {
			fetchText: async (url) => {
				fetched.push(url);
				if (url.endsWith("robots.txt"))
					return "Sitemap: https://x.test/sitemap.xml";
				if (url.endsWith("sitemap.xml"))
					return "<url><loc>https://x.test/a</loc></url>";
				if (url.endsWith("llms.txt")) return "[LLM](https://x.test/llm)";
				throw new Error(`unexpected ${url}`);
			},
		});
		expect(fetched).toEqual([
			"https://x.test/robots.txt",
			"https://x.test/sitemap.xml",
			"https://x.test/llms.txt",
		]);
		expect(result.urls.map((x) => x.url)).toEqual([
			"https://x.test/a",
			"https://x.test/llm",
		]);
	});
});