pi-ext

import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { fauxAssistantMessage } from "@earendil-works/pi-ai";
import { createMcpExtension } from "@earendil-works/pi-coding-agent";
import { createInMemoryTransportPair } from "@earendil-works/pi-mcp/testing";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

const extensionPath = fileURLToPath(new URL("../index.ts", import.meta.url));
type Call = { name: string; arguments: Record<string, unknown> };

function createNativeMcp(calls: Call[]) {
	return createMcpExtension({
		loadConfig: () => ({ servers: [], errors: [] }),
		createTransport: () => {
			const pair = createInMemoryTransportPair();
			pair.server.onMessage(async (message) => {
				if (!("id" in message) || !("method" in message)) return;
				let result: unknown;
				if (message.method === "initialize") {
					result = {
						protocolVersion: "2025-11-25",
						capabilities: { tools: {} },
						serverInfo: { name: "IntelliJ", version: "test" },
					};
				} else if (message.method === "tools/list") {
					result = {
						tools: [
							{
								name: "echo",
								description: "Echo project path.",
								inputSchema: {
									type: "object",
									properties: {
										value: { type: "string" },
										projectPath: { type: "string" },
									},
									additionalProperties: false,
								},
							},
						],
					};
				} else if (message.method === "tools/call") {
					const params = message.params as {
						name: string;
						arguments?: Record<string, unknown>;
					};
					calls.push({ name: params.name, arguments: params.arguments ?? {} });
					result = { content: [{ type: "text", text: "ok" }] };
				} else result = {};
				await pair.server.send({ jsonrpc: "2.0", id: message.id, result });
			});
			void pair.server.start();
			return pair.client;
		},
	});
}

function createProject(): string {
	const root = mkdtempSync(join(tmpdir(), "pi-ext-intellij-"));
	mkdirSync(join(root, ".idea"));
	return root;
}

async function activateDeferredTool(session: PiSession): Promise<void> {
	expect(session.session.getActiveToolNames()).not.toContain(
		"mcp__intellij__echo",
	);
	await vi.waitFor(() => {
		const active = session.session.getActiveToolNames();
		session.session.setActiveToolsByName([...active, "mcp__intellij__echo"]);
		expect(session.session.getActiveToolNames()).toContain(
			"mcp__intellij__echo",
		);
	});
}

async function call(session: PiSession, value: string): Promise<void> {
	session.setResponses([
		fauxAssistantMessage(
			[
				{
					type: "toolCall",
					id: `call-${value}`,
					name: "mcp__intellij__echo",
					arguments: { value, projectPath: "wrong-project" },
				},
			],
			{ stopReason: "toolUse" },
		),
		fauxAssistantMessage("done"),
	]);
	await session.session.prompt(value);
}

describe("intellij native MCP activation", () => {
	const sessions: PiSession[] = [];
	const roots: string[] = [];
	afterEach(async () => {
		// Fixture environments are nested process overrides: restore in reverse order.
		for (const session of sessions.splice(0).reverse()) await session.dispose();
		for (const root of roots.splice(0))
			rmSync(root, { recursive: true, force: true });
	});

	it("path-loads, detects parent .idea, and routes calls through native MCP", async () => {
		const root = createProject();
		roots.push(root);
		const cwd = join(root, "src", "nested");
		mkdirSync(cwd, { recursive: true });
		const calls: Call[] = [];
		const session = await createPiSession({
			cwd,
			userSettings: { intellij: { url: "http://127.0.0.1:63342/mcp" } },
			extensionPaths: [extensionPath],
			extensionFactories: [createNativeMcp(calls)],
		});
		sessions.push(session);
		await activateDeferredTool(session);
		await call(session, "one");
		expect(calls).toEqual([
			{ name: "echo", arguments: { value: "one", projectPath: root } },
		]);
		expect(session.extensions[0]?.commands.has("intellij")).toBe(false);
		expect(session.errors).toEqual([]);
		await session.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		await session.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});
		await activateDeferredTool(session);
		await call(session, "restarted");
		expect(calls.at(-1)).toEqual({
			name: "echo",
			arguments: { value: "restarted", projectPath: root },
		});
		expect(session.errors).toEqual([]);
	});

	it("keeps concurrent sessions isolated and closing one leaves the other callable", async () => {
		const calls: Call[] = [];
		const firstRoot = createProject();
		const secondRoot = createProject();
		roots.push(firstRoot, secondRoot);
		const options = {
			userSettings: { intellij: { url: "http://127.0.0.1:63342/mcp" } },
			extensionPaths: [extensionPath],
			extensionFactories: [createNativeMcp(calls)],
		};
		const first = await createPiSession({ cwd: firstRoot, ...options });
		sessions.push(first);
		const second = await createPiSession({ cwd: secondRoot, ...options });
		sessions.push(second);
		await Promise.all([
			activateDeferredTool(first),
			activateDeferredTool(second),
		]);
		await Promise.all([call(first, "first"), call(second, "second")]);
		await second.dispose();
		sessions.pop();
		await call(first, "again");
		expect(calls).toEqual(
			expect.arrayContaining([
				{ name: "echo", arguments: { value: "first", projectPath: firstRoot } },
				{
					name: "echo",
					arguments: { value: "second", projectPath: secondRoot },
				},
				{ name: "echo", arguments: { value: "again", projectPath: firstRoot } },
			]),
		);
		expect(calls).toHaveLength(3);
		expect(first.errors).toEqual([]);
	});
});