pi-ext

import {
	chmod,
	mkdir,
	mkdtemp,
	readFile,
	rm,
	writeFile,
} from "node:fs/promises";
import { createServer, type Server as HttpServer } from "node:http";
import { tmpdir } from "node:os";
import { delimiter, join } from "node:path";
import { type ExtensionAPI, initTheme } from "@earendil-works/pi-coding-agent";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import {
	CallToolRequestSchema,
	ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { Type } from "typebox";
import { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import intellijExtension from "../index";

function occupiedToolExtension(pi: ExtensionAPI): void {
	pi.registerTool({
		name: "intellij_echo",
		label: "Occupied",
		description: "Occupy the default bridge name.",
		parameters: Type.Object({}),
		async execute() {
			return { content: [{ type: "text", text: "occupied" }] };
		},
	});
}

async function fakeJetBrainsServer(
	calls: Array<Record<string, unknown>>,
	catalog: { tools: string[]; noProjectPathTools?: string[] },
	projectOpen: () => boolean | Promise<boolean> = () => true,
): Promise<{
	server: HttpServer;
	url: string;
}> {
	const server = createServer(async (request, response) => {
		if (request.method !== "POST") {
			response.writeHead(405).end();
			return;
		}
		const chunks: Buffer[] = [];
		for await (const chunk of request) chunks.push(Buffer.from(chunk));
		const body = JSON.parse(Buffer.concat(chunks).toString("utf8"));
		const mcp = new Server(
			{ name: "JetBrains IDE", version: "2025.2" },
			{ capabilities: { tools: {} } },
		);
		mcp.setRequestHandler(ListToolsRequestSchema, async () => ({
			tools: [
				{
					name: "get_project_modules",
					description: "Resolve the open project",
					inputSchema: {
						type: "object",
						properties: { projectPath: { type: "string" } },
						required: ["projectPath"],
					},
				},
				...catalog.tools.map((name) => {
					const acceptsProjectPath =
						!catalog.noProjectPathTools?.includes(name);
					return {
						name,
						description: "Echo project routing and an image",
						inputSchema: {
							type: "object" as const,
							properties: {
								...(acceptsProjectPath
									? { projectPath: { type: "string" } }
									: {}),
								value: { type: "string" },
							},
							required: ["value"],
						},
					};
				}),
			],
		}));
		mcp.setRequestHandler(CallToolRequestSchema, async ({ params }, extra) => {
			calls.push({ name: params.name, arguments: params.arguments });
			if (params.name === "get_project_modules")
				return (await projectOpen())
					? { content: [{ type: "text", text: "project open" }] }
					: {
							content: [{ type: "text", text: "project missing" }],
							isError: true,
						};
			if (params.name === "get_file_problems")
				return {
					content: [
						{
							type: "text",
							text: JSON.stringify({
								filePath: params.arguments?.filePath,
								errors:
									params.arguments?.value === "problems"
										? [
												{
													severity: "WARNING",
													description: "Nullable result is unchecked",
													line: 7,
													column: 2,
												},
											]
										: [],
							}),
						},
					],
				};
			if (params.arguments?.value === "semantic-fail")
				return {
					content: [
						{
							type: "text",
							text: JSON.stringify({
								isSuccess: false,
								problems: [
									{
										kind: "ERROR",
										message: "Cannot resolve symbol",
										file: "src/Main.java",
										line: 7,
									},
								],
							}),
						},
					],
				};
			if (params.arguments?.value === "structured")
				return {
					content: [{ type: "text", text: "summary" }],
					structuredContent: { answer: 42 },
				};
			if (params.arguments?.value === "fail")
				return {
					content: [{ type: "text", text: "remote failure" }],
					isError: true,
				};
			if (params.arguments?.value === "progress")
				await extra.sendNotification({
					method: "notifications/progress",
					params: {
						progressToken: extra._meta?.progressToken ?? "progress",
						progress: 1,
						total: 1,
						message: "indexed",
					},
				});
			if (params.arguments?.value === "wait")
				await new Promise<void>((resolve, reject) => {
					const timer = setTimeout(resolve, 5_000);
					extra.signal.addEventListener(
						"abort",
						() => {
							clearTimeout(timer);
							reject(new Error("remote cancelled"));
						},
						{ once: true },
					);
				});
			return {
				content: [
					{ type: "text", text: `echo ${params.arguments?.value}` },
					{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" },
					{
						type: "resource_link",
						uri: "file:///project/result.txt",
						name: "result",
					},
				],
			};
		});
		const transport = new StreamableHTTPServerTransport({
			sessionIdGenerator: undefined,
		});
		await mcp.connect(transport);
		await transport.handleRequest(request, response, body);
		response.once("close", () => {
			void transport.close();
			void mcp.close();
		});
	});
	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
	const address = server.address();
	if (!address || typeof address === "string")
		throw new Error("missing address");
	return { server, url: `http://127.0.0.1:${address.port}/mcp` };
}

async function writeFakeLauncher(bin: string, marker: string): Promise<void> {
	await mkdir(bin, { recursive: true });
	if (process.platform === "win32") {
		await writeFile(
			join(bin, "idea.cmd"),
			`@echo off\r\n>>"${marker}" echo(%~1\r\n`,
		);
		return;
	}
	const launcher = join(bin, "idea");
	await writeFile(launcher, `#!/bin/sh\nprintf '%s\\n' "$1" >> '${marker}'\n`);
	await chmod(launcher, 0o700);
}

async function readEventually(file: string): Promise<string> {
	for (let attempt = 0; attempt < 50; attempt++) {
		try {
			return await readFile(file, "utf8");
		} catch {
			await new Promise((resolve) => setTimeout(resolve, 20));
		}
	}
	throw new Error(`Timed out waiting for ${file}`);
}

describe("intellij pi-test-harness", () => {
	let t: TestSession | undefined;
	let server: HttpServer | undefined;
	const roots: string[] = [];

	afterEach(async () => {
		t?.dispose();
		t = undefined;
		if (server)
			await new Promise<void>((resolve, reject) =>
				server?.close((error) => (error ? reject(error) : resolve())),
			);
		server = undefined;
		await Promise.all(
			roots.splice(0).map((root) => rm(root, { recursive: true, force: true })),
		);
	});

	it("loads disconnected with no startup IntelliJ tools", async () => {
		t = await createTestSession({ extensionFactories: [intellijExtension] });
		const [extension] = t.session.extensionRunner.extensions;
		expect(extension.commands.has("intellij")).toBe(true);
		expect([...extension.tools.keys()]).toEqual([]);
		expect(
			t.session
				.getActiveToolNames()
				.some((name) => name.startsWith("intellij_")),
		).toBe(false);
	});

	it("cancels an in-flight connection without launching IntelliJ", async () => {
		const root = await mkdtemp(join(tmpdir(), "intellij-cancel-"));
		roots.push(root);
		const bin = join(root, "bin");
		const marker = join(root, "launcher-args.txt");
		await writeFakeLauncher(bin, marker);
		let requestArrived!: () => void;
		const received = new Promise<void>((resolve) => {
			requestArrived = resolve;
		});
		server = createServer(() => requestArrived());
		await new Promise<void>((resolve) =>
			server?.listen(0, "127.0.0.1", resolve),
		);
		const address = server.address();
		if (!address || typeof address === "string")
			throw new Error("missing address");
		t = await createTestSession({
			extensionFactories: [intellijExtension],
			env: { PATH: `${bin}${delimiter}${process.env.PATH ?? ""}` },
		});
		const [extension] = t.session.extensionRunner.extensions;
		const command = extension.commands.get("intellij");
		const ctx = t.session.extensionRunner.createCommandContext();

		const connecting = command.handler(
			`connect http://127.0.0.1:${address.port}/mcp`,
			ctx,
		);
		await received;
		await command.handler("disconnect", ctx);
		await connecting;

		await expect(readFile(marker, "utf8")).rejects.toMatchObject({
			code: "ENOENT",
		});
		expect(
			t.session
				.getActiveToolNames()
				.some((name) => name.startsWith("intellij_")),
		).toBe(false);
	});

	it("launches a missing project once, retains the catalog, and reports a stale saved URL", async () => {
		const root = await mkdtemp(join(tmpdir(), "intellij-project-sync-"));
		roots.push(root);
		const project = join(root, "project & exact");
		const bin = join(root, "bin");
		const marker = join(root, "launcher-args.txt");
		await mkdir(project);
		await writeFakeLauncher(bin, marker);
		const calls: Array<Record<string, unknown>> = [];
		const fake = await fakeJetBrainsServer(
			calls,
			{ tools: ["echo"] },
			async () => {
				try {
					return (await readFile(marker, "utf8")).trim() === project;
				} catch {
					return false;
				}
			},
		);
		server = fake.server;
		t = await createTestSession({
			cwd: project,
			extensionFactories: [intellijExtension],
			env: { PATH: `${bin}${delimiter}${process.env.PATH ?? ""}` },
		});
		const [extension] = t.session.extensionRunner.extensions;
		const command = extension.commands.get("intellij");
		const ctx = t.session.extensionRunner.createCommandContext();

		await command.handler(`connect ${fake.url}`, ctx);
		expect((await readEventually(marker)).trim().split(/\r?\n/)).toEqual([
			project,
		]);
		expect(t.session.getActiveToolNames()).toContain("intellij_echo");

		await command.handler("disconnect", ctx);
		await command.handler("tools", ctx);
		expect(t.events.uiCallsFor("notify").at(-1)?.args[0]).toContain(
			"stale intellij_echo",
		);
		await new Promise<void>((resolve, reject) =>
			server?.close((error) => (error ? reject(error) : resolve())),
		);
		server = undefined;

		await command.handler("reconnect", ctx);
		const failure = String(t.events.uiCallsFor("notify").at(-1)?.args[0]);
		expect(failure).toContain("The saved URL is stale.");
		expect(failure).toContain("Copy HTTP Stream Config");
		expect(t.session.getActiveToolNames()).not.toContain("intellij_echo");
		expect((await readEventually(marker)).trim().split(/\r?\n/)).toHaveLength(
			1,
		);
	});

	it("launches IntelliJ and retries the saved endpoint when the IDE is closed", async () => {
		const root = await mkdtemp(join(tmpdir(), "intellij-closed-"));
		roots.push(root);
		const project = join(root, "project");
		const bin = join(root, "bin");
		const marker = join(root, "launcher-args.txt");
		await mkdir(project);
		await writeFakeLauncher(bin, marker);
		const fake = await fakeJetBrainsServer([], { tools: ["echo"] });
		const port = Number(new URL(fake.url).port);
		await new Promise<void>((resolve, reject) =>
			fake.server.close((error) => (error ? reject(error) : resolve())),
		);
		server = fake.server;
		let reopening = false;
		const watcher = setInterval(async () => {
			if (reopening) return;
			try {
				await readFile(marker);
				reopening = true;
				server?.listen(port, "127.0.0.1");
			} catch {
				// Launcher has not run yet.
			}
		}, 20);
		try {
			t = await createTestSession({
				cwd: project,
				extensionFactories: [intellijExtension],
				env: { PATH: `${bin}${delimiter}${process.env.PATH ?? ""}` },
			});
			const [extension] = t.session.extensionRunner.extensions;
			const command = extension.commands.get("intellij");
			await command.handler(
				`connect ${fake.url}`,
				t.session.extensionRunner.createCommandContext(),
			);
			expect(t.session.getActiveToolNames()).toContain("intellij_echo");
			expect((await readEventually(marker)).trim()).toBe(project);
		} finally {
			clearInterval(watcher);
		}
	});

	it("resolves tool-name collisions deterministically", async () => {
		const fake = await fakeJetBrainsServer([], { tools: ["echo"] });
		server = fake.server;
		t = await createTestSession({
			extensionFactories: [occupiedToolExtension, intellijExtension],
		});
		const extension = t.session.extensionRunner.extensions[1];
		await extension.commands
			.get("intellij")
			.handler(
				`connect ${fake.url}`,
				t.session.extensionRunner.createCommandContext(),
			);
		expect(extension.tools.has("intellij_echo_2")).toBe(true);
		expect(t.session.getActiveToolNames()).toContain("intellij_echo_2");
		await extension.commands
			.get("intellij")
			.handler("disconnect", t.session.extensionRunner.createCommandContext());
		expect(t.session.getActiveToolNames()).toContain("intellij_echo");
		expect(t.session.getActiveToolNames()).not.toContain("intellij_echo_2");
	});

	it("routes concurrent session clients without cross-project leakage", async () => {
		const calls: Array<Record<string, unknown>> = [];
		const fake = await fakeJetBrainsServer(calls, { tools: ["echo"] });
		server = fake.server;
		const first = await createTestSession({
			extensionFactories: [intellijExtension],
		});
		const second = await createTestSession({
			extensionFactories: [intellijExtension],
		});
		try {
			await Promise.all(
				[first, second].map(async (session) => {
					const [extension] = session.session.extensionRunner.extensions;
					const ctx = session.session.extensionRunner.createCommandContext();
					await extension.commands
						.get("intellij")
						.handler(`connect ${fake.url}`, ctx);
					await extension.tools
						.get("intellij_echo")
						.definition.execute(
							"call",
							{ value: session.cwd },
							undefined,
							undefined,
							ctx,
						);
				}),
			);
			const echoCalls = calls
				.filter((call) => call.name === "echo")
				.sort((left, right) =>
					String(
						(left.arguments as Record<string, unknown>).value,
					).localeCompare(
						String((right.arguments as Record<string, unknown>).value),
					),
				);
			const expected = [first, second]
				.map((session) => ({
					name: "echo",
					arguments: { value: session.cwd, projectPath: session.cwd },
				}))
				.sort((left, right) =>
					left.arguments.value.localeCompare(right.arguments.value),
				);
			expect(echoCalls).toEqual(expected);
		} finally {
			second.dispose();
			first.dispose();
		}
	});

	it("presents routed tools and inspection results with useful compact detail", async () => {
		const fake = await fakeJetBrainsServer([], {
			tools: ["get_file_problems", "global_status"],
			noProjectPathTools: ["global_status"],
		});
		server = fake.server;
		t = await createTestSession({ extensionFactories: [intellijExtension] });
		const [extension] = t.session.extensionRunner.extensions;
		const command = extension.commands.get("intellij");
		const ctx = t.session.extensionRunner.createCommandContext();

		await command.handler(`connect ${fake.url}`, ctx);
		const status = String(t.events.uiCallsFor("setStatus").at(-1)?.args[1]);
		expect(status).toContain(" ✓");
		expect(status).not.toContain("tools");

		const probe = extension.tools.get(
			"intellij_get_project_modules",
		).definition;
		expect(probe.promptGuidelines).toEqual([
			expect.stringContaining("intellij_*"),
			expect.stringContaining("available intellij_* inspection tools"),
		]);
		const definition = extension.tools.get(
			"intellij_get_file_problems",
		).definition;
		expect(definition.description).toContain(
			"injects projectPath automatically",
		);
		expect(definition.parameters.properties).not.toHaveProperty("projectPath");
		expect(definition.parameters.required).toEqual(["value"]);
		const globalStatus = extension.tools.get(
			"intellij_global_status",
		).definition;
		expect(globalStatus.description).toContain("connected IntelliJ session");
		expect(globalStatus.description).not.toContain("projectPath");
		expect(globalStatus.parameters.properties).not.toHaveProperty(
			"projectPath",
		);

		const args = {
			value: "problems",
			filePath: "src/Main.java",
		};
		const result = await definition.execute(
			"problems",
			args,
			undefined,
			undefined,
			ctx,
		);
		initTheme("default", false);
		const theme = {
			fg: (_style: string, text: string) => text,
			bold: (text: string) => text,
		};
		const call = definition
			.renderCall(args, theme as never, { expanded: false } as never)
			.render(120)
			.join("\n");
		const collapsed = definition
			.renderResult(
				result,
				{ expanded: false } as never,
				theme as never,
				{ isError: false } as never,
			)
			.render(120)
			.join("\n");
		const expanded = definition
			.renderResult(
				result,
				{ expanded: true } as never,
				theme as never,
				{ isError: false } as never,
			)
			.render(120)
			.join("\n");
		expect(call).toContain("get file problems");
		expect(call).toContain("src/Main.java");
		expect(collapsed).toContain("1 warning");
		expect(collapsed).not.toContain("Nullable result");
		expect(expanded).toContain("WARNING src/Main.java:7:2");
		expect(expanded).toContain("Nullable result is unchecked");

		const progress = {
			content: [
				{
					type: "text",
					text: "IntelliJ progress: indexing END-PROGRESS-DIAGNOSTIC",
				},
			],
			details: { source: "intellij", tool: "get_file_problems" },
		};
		const collapsedProgress = definition
			.renderResult(
				progress,
				{ expanded: false, isPartial: true } as never,
				theme as never,
				{} as never,
			)
			.render(120)
			.join("\n");
		const expandedProgress = definition
			.renderResult(
				progress,
				{ expanded: true, isPartial: true } as never,
				theme as never,
				{} as never,
			)
			.render(120)
			.join("\n");
		expect(collapsedProgress).toContain("to expand");
		expect(collapsedProgress).not.toContain("END-PROGRESS-DIAGNOSTIC");
		expect(expandedProgress).toContain("END-PROGRESS-DIAGNOSTIC");
	});

	it("deactivates the live catalog before rejecting a replacement URL", async () => {
		const fake = await fakeJetBrainsServer([], { tools: ["echo"] });
		server = fake.server;
		t = await createTestSession({ extensionFactories: [intellijExtension] });
		const [extension] = t.session.extensionRunner.extensions;
		const command = extension.commands.get("intellij");
		const ctx = t.session.extensionRunner.createCommandContext();

		await command.handler(`connect ${fake.url}`, ctx);
		expect(t.session.getActiveToolNames()).toContain("intellij_echo");
		await command.handler("connect https://example.com/mcp", ctx);
		expect(t.session.getActiveToolNames()).not.toContain("intellij_echo");
		expect(String(t.events.uiCallsFor("notify").at(-1)?.args[0])).toContain(
			"must use HTTP",
		);
	});

	it("connects to a trusted project URL without persisting it globally", async () => {
		const fake = await fakeJetBrainsServer([], { tools: ["echo"] });
		server = fake.server;
		t = await createTestSession({ extensionFactories: [intellijExtension] });
		const userFile = join(
			t.env.HOME as string,
			".pi",
			"agent",
			"settings.json",
		);
		const user = { intellij: { url: "http://127.0.0.1:1/user" } };
		await mkdir(join(t.env.HOME as string, ".pi", "agent"), {
			recursive: true,
		});
		await writeFile(userFile, JSON.stringify(user));
		await mkdir(join(t.cwd, ".pi"), { recursive: true });
		await writeFile(
			join(t.cwd, ".pi", "settings.json"),
			JSON.stringify({ intellij: { url: fake.url } }),
		);
		const [extension] = t.session.extensionRunner.extensions;

		await extension.commands
			.get("intellij")
			.handler("connect", t.session.extensionRunner.createCommandContext());

		expect(t.session.getActiveToolNames()).toContain("intellij_echo");
		expect(JSON.parse(await readFile(userFile, "utf8"))).toEqual(user);
	});

	it("connects through Streamable HTTP, pins projectPath, maps content, and disconnects", async () => {
		const calls: Array<Record<string, unknown>> = [];
		const catalog = { tools: ["echo.image"] };
		const fake = await fakeJetBrainsServer(calls, catalog);
		server = fake.server;
		t = await createTestSession({ extensionFactories: [intellijExtension] });
		const settingsDir = join(t.env.HOME as string, ".pi", "agent");
		const settingsFile = join(settingsDir, "settings.json");
		await mkdir(settingsDir, { recursive: true });
		await writeFile(
			settingsFile,
			JSON.stringify({ theme: "nugu", intellij: { keep: true } }),
		);
		const [extension] = t.session.extensionRunner.extensions;
		const command = extension.commands.get("intellij");
		const ctx = t.session.extensionRunner.createCommandContext();

		await command.handler(`connect ${fake.url}`, ctx);

		expect(JSON.parse(await readFile(settingsFile, "utf8"))).toEqual({
			theme: "nugu",
			intellij: { keep: true, url: fake.url },
		});
		expect(t.session.getActiveToolNames()).toContain("intellij_echo_image");
		const definition = extension.tools.get("intellij_echo_image").definition;
		const result = await definition.execute(
			"call-1",
			{ value: "hello", projectPath: "C:/wrong-project" },
			undefined,
			undefined,
			ctx,
		);
		expect(calls.at(-1)).toEqual({
			name: "echo.image",
			arguments: { value: "hello", projectPath: t.cwd },
		});
		expect(result.content).toEqual([
			{ type: "text", text: "echo hello" },
			{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" },
			{
				type: "text",
				text: '{"name":"result","uri":"file:///project/result.txt","type":"resource_link"}',
			},
		]);
		initTheme("default", false);
		const theme = {
			fg: (_style: string, text: string) => text,
			bold: (text: string) => text,
		};
		const collapsed = definition
			.renderResult(
				result,
				{ expanded: false } as never,
				theme as never,
				{} as never,
			)
			.render(120)
			.join("\n");
		const expanded = definition
			.renderResult(
				result,
				{ expanded: true } as never,
				theme as never,
				{} as never,
			)
			.render(120)
			.join("\n");
		expect(collapsed).not.toContain("echo hello");
		expect(expanded).toContain("echo hello");
		const onUpdate = vi.fn();
		await definition.execute(
			"progress",
			{ value: "progress" },
			undefined,
			onUpdate,
			ctx,
		);
		expect(onUpdate).toHaveBeenCalledWith(
			expect.objectContaining({
				content: [
					expect.objectContaining({ text: expect.stringContaining("indexed") }),
				],
			}),
		);
		const abort = new AbortController();
		const waiting = definition.execute(
			"cancel",
			{ value: "wait" },
			abort.signal,
			undefined,
			ctx,
		);
		abort.abort();
		await expect(waiting).rejects.toThrow();
		await expect(
			definition.execute(
				"failure",
				{ value: "fail" },
				undefined,
				undefined,
				ctx,
			),
		).rejects.toThrow("remote failure");
		await expect(
			definition.execute(
				"semantic-failure",
				{ value: "semantic-fail" },
				undefined,
				undefined,
				ctx,
			),
		).rejects.toThrow(
			/IntelliJ echo\.image failed[\s\S]*Cannot resolve symbol[\s\S]*Next:/,
		);
		const structured = await definition.execute(
			"structured",
			{ value: "structured" },
			undefined,
			undefined,
			ctx,
		);
		expect(structured.content).toEqual([
			{ type: "text", text: "summary" },
			{ type: "text", text: '{"answer":42}' },
		]);
		await expect(
			definition.execute(
				"outside",
				{ value: "blocked" },
				undefined,
				undefined,
				{ ...ctx, cwd: join(t.cwd, "other") },
			),
		).rejects.toThrow("outside this session cwd");

		catalog.tools = ["new_tool"];
		await command.handler("reconnect", ctx);
		expect(t.session.getActiveToolNames()).not.toContain("intellij_echo_image");
		expect(t.session.getActiveToolNames()).toContain("intellij_new_tool");
		await expect(
			definition.execute(
				"call-2",
				{ value: "later" },
				undefined,
				undefined,
				ctx,
			),
		).rejects.toThrow("disconnected or stale");

		await command.handler("disconnect", ctx);
		expect(t.session.getActiveToolNames()).not.toContain("intellij_new_tool");
		expect(t.events.uiCallsFor("setStatus").at(-1)?.args).toEqual([
			"intellij",
			undefined,
		]);
	});
});