pi-ext

import {
	access,
	mkdir,
	mkdtemp,
	rename,
	rm,
	symlink,
	writeFile,
} from "node:fs/promises";
import { createServer, type Server } from "node:net";
import { tmpdir } from "node:os";
import path from "node:path";
import { fauxToolCall } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";

const filesystemControl = vi.hoisted(() => ({
	afterRealpath: undefined as
		| ((requested: string, resolved: string) => Promise<void>)
		| undefined,
	beforeRead: undefined as ((openedPath: string) => Promise<void>) | undefined,
	readRequests: [] as Array<{ path: string; bytes: number }>,
}));
const protectionControl = vi.hoisted(() => ({ parts: new Set<string>() }));

vi.mock("node:fs/promises", async (importOriginal) => {
	const original = await importOriginal<typeof import("node:fs/promises")>();
	return {
		...original,
		realpath: async (value: string) => {
			const resolved = await original.realpath(value);
			await filesystemControl.afterRealpath?.(value, resolved);
			return resolved;
		},
		open: async (filePath: string, flags: number) => {
			const handle = await original.open(filePath, flags);
			return {
				fd: handle.fd,
				stat: handle.stat.bind(handle),
				close: handle.close.bind(handle),
				read: async (
					buffer: Buffer,
					offset: number,
					length: number,
					position: number | null,
				) => {
					const beforeRead = filesystemControl.beforeRead;
					filesystemControl.beforeRead = undefined;
					await beforeRead?.(filePath);
					filesystemControl.readRequests.push({
						path: filePath,
						bytes: length,
					});
					return handle.read(buffer, offset, length, position);
				},
			};
		},
	};
});

vi.mock("../git.js", async (importOriginal) => {
	const original = await importOriginal<typeof import("../git.js")>();
	return {
		...original,
		isProtectedPath: (value: string, nativePath = false) =>
			original.isProtectedPath(value, nativePath) ||
			value.split(/[\\/]+/u).some((part) => protectionControl.parts.has(part)),
	};
});

import { executeFilesystemTool, resolveAuthorizedPath } from "../filesystem.js";

const MAX_FILE_BYTES = 2 * 1024 * 1024;
const roots: string[] = [];

async function fixture(): Promise<{ base: string; repo: string }> {
	const base = await mkdtemp(path.join(tmpdir(), "strata-filesystem-"));
	roots.push(base);
	const repo = path.join(base, "repo");
	await mkdir(repo);
	return { base, repo };
}

function execute(
	cwd: string,
	name: "read" | "grep" | "ls",
	input: Record<string, unknown>,
	signal = new AbortController().signal,
) {
	return executeFilesystemTool(
		cwd,
		fauxToolCall(name, input),
		signal,
	) as Promise<Record<string, unknown>>;
}

afterEach(async () => {
	filesystemControl.afterRealpath = undefined;
	filesystemControl.beforeRead = undefined;
	filesystemControl.readRequests.length = 0;
	protectionControl.parts.clear();
	await Promise.all(
		roots.splice(0).map((root) => rm(root, { recursive: true, force: true })),
	);
});

function swapAfterPreOpenResolution(
	filePath: string,
	swap: () => Promise<void>,
): void {
	let resolutions = 0;
	filesystemControl.afterRealpath = async (requested) => {
		if (requested !== filePath || ++resolutions !== 2) return;
		filesystemControl.afterRealpath = undefined;
		await swap();
	};
}

describe("authorized Ask filesystem tools", () => {
	it("reads ordinary untracked text with bounded line pagination", async () => {
		const { repo } = await fixture();
		await writeFile(path.join(repo, "scratch.txt"), "one\ntwo\nthree\nfour\n");

		const first = await execute(repo, "read", {
			path: "scratch.txt",
			offset: 2,
			limit: 2,
		});
		expect(first).toMatchObject({
			evidence: "CURRENT FILESYSTEM",
			offset: 2,
			returnedLines: 2,
			content: "two\nthree",
			truncated: true,
			nextOffset: 4,
		});
		const second = await execute(repo, "read", {
			path: "scratch.txt",
			offset: 4,
			limit: 2,
		});
		expect(second.content).toBe("four\n");
		await expect(
			execute(repo, "read", { path: "scratch.txt", offset: 99 }),
		).rejects.toThrow(/beyond end of file/);

		await writeFile(path.join(repo, "long-line.txt"), "x".repeat(30 * 1024));
		const bounded = await execute(repo, "read", { path: "long-line.txt" });
		expect(bounded).toMatchObject({
			returnedLines: 1,
			truncated: true,
			lineTruncated: true,
		});
		expect(Buffer.byteLength(bounded.content as string)).toBeLessThanOrEqual(
			24 * 1024,
		);
	});

	it("accepts every nextOffset emitted at the maximum text-file boundary", async () => {
		const { repo } = await fixture();
		await writeFile(
			path.join(repo, "line-boundary.txt"),
			"\n".repeat(MAX_FILE_BYTES),
		);

		const penultimate = await execute(repo, "read", {
			path: "line-boundary.txt",
			offset: MAX_FILE_BYTES,
			limit: 1,
		});
		expect(penultimate).toMatchObject({
			totalLines: MAX_FILE_BYTES + 1,
			nextOffset: MAX_FILE_BYTES + 1,
			truncated: true,
		});
		const final = await execute(repo, "read", {
			path: "line-boundary.txt",
			offset: penultimate.nextOffset,
			limit: 1,
		});
		expect(final).toMatchObject({
			offset: MAX_FILE_BYTES + 1,
			returnedLines: 1,
			truncated: false,
		});
	});

	it("searches an explicit ignored node_modules directory and bounds matches", async () => {
		const { repo } = await fixture();
		const dependency = path.join(repo, "node_modules", "example-package");
		await mkdir(dependency, { recursive: true });
		await writeFile(path.join(repo, ".gitignore"), "node_modules/\n");
		await writeFile(
			path.join(dependency, "README.md"),
			"Package marker\npackage MARKER\nmarker again\n",
		);

		const result = await execute(repo, "grep", {
			path: "node_modules/example-package",
			pattern: "marker",
			ignoreCase: true,
			context: 1,
			limit: 1,
		});
		expect(result).toMatchObject({
			evidence: "CURRENT FILESYSTEM",
			returnedMatches: 1,
			totalMatches: 3,
			truncated: true,
		});
		expect(JSON.stringify(result)).toContain("README.md");
	});

	it("marks clipped match context as aggregate truncation", async () => {
		const { repo } = await fixture();
		await writeFile(
			path.join(repo, "context.txt"),
			`${"b".repeat(700)}\nneedle\n${"a".repeat(700)}\n`,
		);

		const result = await execute(repo, "grep", {
			path: "context.txt",
			pattern: "needle",
			context: 1,
		});
		expect(result).toMatchObject({
			returnedMatches: 1,
			linesTruncated: true,
			truncated: true,
		});
		const [match] = result.matches as Array<{
			before: string[];
			after: string[];
		}>;
		expect(match.before[0]).toHaveLength(501);
		expect(match.after[0]).toHaveLength(501);
		expect(match.before[0]).toMatch(/…$/u);
		expect(match.after[0]).toMatch(/…$/u);
	});

	it("reads absolute documentation outside the repository", async () => {
		const { base, repo } = await fixture();
		const docs = path.join(base, "external-docs");
		await mkdir(docs);
		const reference = path.join(docs, "reference.md");
		await writeFile(reference, "# External reference\nfilesystem evidence\n");

		const result = await execute(repo, "read", { path: reference });
		expect(result.content).toContain("filesystem evidence");
		expect(result.canonicalPath).toBe(
			await resolveAuthorizedPath(repo, reference).then(
				(value) => value.canonical,
			),
		);
	});

	it("follows safe package symlinks and terminates directory cycles", async () => {
		const { repo } = await fixture();
		const target = path.join(repo, "packages", "linked-package");
		await mkdir(target, { recursive: true });
		await writeFile(path.join(target, "README.md"), "safe symlink marker\n");
		const link = path.join(repo, "node_modules", "linked-package");
		await mkdir(path.dirname(link), { recursive: true });
		await symlink(
			target,
			link,
			process.platform === "win32" ? "junction" : "dir",
		);
		await symlink(
			target,
			path.join(target, "loop"),
			process.platform === "win32" ? "junction" : "dir",
		);

		const listed = await execute(repo, "ls", {
			path: "node_modules/linked-package",
		});
		expect(listed.entries).toEqual(
			expect.arrayContaining(["loop/", "README.md"]),
		);
		const searched = await execute(repo, "grep", {
			path: "node_modules/linked-package",
			pattern: "symlink marker",
		});
		expect(searched).toMatchObject({
			totalMatches: 1,
			directoriesVisited: 1,
		});
	});

	it.each(["read", "grep"] as const)(
		"rejects a final-file swap after resolution through %s before reading bytes",
		async (tool) => {
			const { repo } = await fixture();
			const visible = path.join(repo, "visible.txt");
			const blocked = path.join(repo, "blocked-target.txt");
			await writeFile(visible, "safe content\n");
			await writeFile(blocked, "blocked marker\n");
			protectionControl.parts.add(path.basename(blocked));
			swapAfterPreOpenResolution(visible, async () => {
				await rm(visible);
				await symlink(
					blocked,
					visible,
					process.platform === "win32" ? "file" : undefined,
				);
			});

			await expect(
				execute(
					repo,
					tool,
					tool === "read"
						? { path: "visible.txt" }
						: { path: ".", pattern: "blocked marker" },
				),
			).rejects.toThrow(/protected|authorization|safely opened/u);
			expect(filesystemControl.readRequests).toEqual([]);
		},
	);

	it.each(["read", "grep"] as const)(
		"rejects an ancestor-directory swap after resolution through %s before reading bytes",
		async (tool) => {
			const { repo } = await fixture();
			const visibleDirectory = path.join(repo, "visible-directory");
			const archivedDirectory = path.join(repo, "archived-directory");
			const blockedDirectory = path.join(repo, "blocked-directory");
			await mkdir(visibleDirectory);
			await mkdir(blockedDirectory);
			await writeFile(
				path.join(visibleDirectory, "document.txt"),
				"safe content\n",
			);
			await writeFile(
				path.join(blockedDirectory, "document.txt"),
				"blocked marker\n",
			);
			protectionControl.parts.add(path.basename(blockedDirectory));
			const visible = path.join(visibleDirectory, "document.txt");
			swapAfterPreOpenResolution(visible, async () => {
				await rename(visibleDirectory, archivedDirectory);
				await symlink(
					blockedDirectory,
					visibleDirectory,
					process.platform === "win32" ? "junction" : "dir",
				);
			});

			await expect(
				execute(
					repo,
					tool,
					tool === "read"
						? { path: "visible-directory/document.txt" }
						: { path: ".", pattern: "blocked marker" },
				),
			).rejects.toThrow(/protected|authorization/u);
			expect(filesystemControl.readRequests).toEqual([]);
		},
	);

	it("rejects protected lexical ancestors and canonical symlink targets before reads", async () => {
		const { repo } = await fixture();
		let resolved = false;
		await expect(
			resolveAuthorizedPath(
				repo,
				"private/.env.production/../public.txt",
				async () => {
					resolved = true;
					return path.join(repo, "public.txt");
				},
			),
		).rejects.toThrow(/protected path/);
		expect(resolved).toBe(false);

		await expect(
			resolveAuthorizedPath(repo, "ordinary-link", async () => {
				return path.join(repo, ".git", "objects", "pretend");
			}),
		).rejects.toThrow(/protected path/);
	});

	it("rejects URL-like paths, invalid arguments, protected traversal, and missing paths", async () => {
		const { repo } = await fixture();
		for (const unsafe of [
			"https://example.invalid/file",
			"file:///ordinary/file",
			"safe/../.git/config",
		]) {
			await expect(execute(repo, "read", { path: unsafe })).rejects.toThrow(
				/invalid filesystem path|protected path/,
			);
		}
		await expect(
			execute(repo, "read", { path: "missing.txt" }),
		).rejects.toThrow(/does not exist or cannot be resolved/);
		await expect(
			execute(repo, "grep", { path: ".", pattern: "x", context: 99 }),
		).rejects.toThrow();
	});

	it("bounds bytes read when a regular file grows after opened-handle checks", async () => {
		const { repo } = await fixture();
		const growing = path.join(repo, "growing.txt");
		await writeFile(growing, "small\n");
		filesystemControl.beforeRead = async () => {
			await writeFile(growing, Buffer.alloc(MAX_FILE_BYTES + 10, 0x61));
		};

		await expect(
			execute(repo, "read", { path: "growing.txt" }),
		).rejects.toThrow(/byte limit/u);
		expect(
			filesystemControl.readRequests.reduce(
				(total, request) => total + request.bytes,
				0,
			),
		).toBe(MAX_FILE_BYTES + 1);
	});

	it("stops grep at the actual total byte budget", async () => {
		const { repo } = await fixture();
		await Promise.all(
			Array.from({ length: 5 }, (_, index) =>
				writeFile(
					path.join(repo, `budget-${index}.txt`),
					Buffer.alloc(MAX_FILE_BYTES, 0x61),
				),
			),
		);

		const result = await execute(repo, "grep", {
			path: ".",
			pattern: "absent",
		});
		expect(result).toMatchObject({
			bytesVisited: 4 * MAX_FILE_BYTES,
			filesVisited: 4,
			truncated: true,
		});
		expect(result.limitReasons).toContain("total byte limit");
	});

	it("rejects directories, non-regular sockets, binary files, and oversized files", async () => {
		const { repo } = await fixture();
		await writeFile(path.join(repo, "binary.dat"), Buffer.from([0, 1, 2, 3]));
		await writeFile(
			path.join(repo, "oversized.txt"),
			Buffer.alloc(2 * 1024 * 1024 + 1, 0x61),
		);
		await expect(execute(repo, "read", { path: "." })).rejects.toThrow(
			/not a regular file/,
		);
		await expect(execute(repo, "read", { path: "binary.dat" })).rejects.toThrow(
			/binary file/,
		);
		await expect(
			execute(repo, "read", { path: "oversized.txt" }),
		).rejects.toThrow(/byte limit/);

		const binarySearch = await execute(repo, "grep", {
			path: "binary.dat",
			pattern: "absent",
		});
		expect(binarySearch).toMatchObject({
			bytesVisited: 4,
			skippedBinary: 1,
		});
		const oversizedSearch = await execute(repo, "grep", {
			path: "oversized.txt",
			pattern: "absent",
		});
		expect(oversizedSearch).toMatchObject({
			bytesVisited: MAX_FILE_BYTES + 1,
			truncated: true,
		});
		expect(oversizedSearch.limitReasons).toContain("per-file byte limit");

		if (process.platform !== "win32") {
			const socketPath = path.join(repo, "ordinary.sock");
			let server: Server | undefined;
			try {
				server = createServer();
				await new Promise<void>((resolve, reject) => {
					server?.once("error", reject);
					server?.listen(socketPath, resolve);
				});
				await expect(
					execute(repo, "read", { path: socketPath }),
				).rejects.toThrow(/not a regular file/);
			} finally {
				await new Promise<void>((resolve) => server?.close(() => resolve()));
			}
		}
	});

	it("returns bounded ls names with directory indicators", async () => {
		const { repo } = await fixture();
		await mkdir(path.join(repo, "directory"));
		await Promise.all(
			Array.from({ length: 12 }, (_, index) =>
				writeFile(path.join(repo, `file-${index}.txt`), "ordinary\n"),
			),
		);
		const result = await execute(repo, "ls", { limit: 3 });
		expect(result).toMatchObject({ returned: 3, truncated: true });
		expect(result.entries).toHaveLength(3);
		const complete = await execute(repo, "ls", { limit: 20 });
		expect(complete.entries).toContain("directory/");
	});

	it("reports recursive traversal limits instead of silently losing output", async () => {
		const { repo } = await fixture();
		let directory = repo;
		for (let depth = 0; depth < 27; depth++) {
			directory = path.join(directory, `level-${depth}`);
			await mkdir(directory);
		}
		await writeFile(path.join(directory, "deep.txt"), "deep marker\n");
		const result = await execute(repo, "grep", {
			path: ".",
			pattern: "deep marker",
		});
		expect(result).toMatchObject({ totalMatches: 0, truncated: true });
		expect(result.limitReasons).toContain("depth limit");
	});

	it("honors cancellation before filesystem access", async () => {
		const { repo } = await fixture();
		await writeFile(path.join(repo, "ordinary.txt"), "content\n");
		const controller = new AbortController();
		controller.abort(new Error("filesystem cancelled"));
		await expect(
			execute(repo, "read", { path: "ordinary.txt" }, controller.signal),
		).rejects.toThrow("filesystem cancelled");
		await expect(
			access(path.join(repo, "ordinary.txt")),
		).resolves.toBeUndefined();
	});

	it("treats backslashes according to the native platform", async () => {
		const { repo } = await fixture();
		if (process.platform === "win32") {
			await expect(
				resolveAuthorizedPath(repo, ".git\\config", async (value) => value),
			).rejects.toThrow(/protected path/);
		} else {
			const filename = "ordinary\\name.txt";
			await writeFile(path.join(repo, filename), "backslash filename\n");
			const result = await execute(repo, "read", { path: filename });
			expect(result.content).toContain("backslash filename");
		}
	});
});