path: root/packages/mcp/src/extension.test.ts

import type { Extension, HostAPI, Logger, ToolContract } from "@dispatch/kernel";
import type { ToolAssembly, toolsFilter } from "@dispatch/session-orchestrator";
import { describe, expect, it } from "vitest";
import { filterMcpTools, makeMcpExtension } from "./extension.js";
import { encode, FrameDecoder } from "./framing.js";
import type { SpawnedProcess, SpawnProcess } from "./transport.js";
import type { McpToolInfo } from "./types.js";

// ---------------------------------------------------------------------------
// Pure filterMcpTools
// ---------------------------------------------------------------------------

const stubTool = (name: string): ToolContract => ({
	name,
	description: "",
	parameters: { type: "object" },
	execute: async () => ({ content: "" }),
});

describe("filterMcpTools (pure)", () => {
	it("keeps non-MCP tools and connected-server tools, removes disconnected-server tools", () => {
		const toolToServer = new Map<string, string>([
			["a__x", "a"],
			["b__y", "b"],
		]);
		const connected = new Set<string>(["a"]);

		const result = filterMcpTools(
			{
				tools: [stubTool("a__x"), stubTool("b__y"), stubTool("other")],
				cwd: "/p",
				conversationId: "c",
			},
			toolToServer,
			connected,
		);

		expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]);
		expect(result.cwd).toBe("/p");
		expect(result.conversationId).toBe("c");
	});

	it("removes all MCP tools when no server is connected", () => {
		const result = filterMcpTools(
			{ tools: [stubTool("a__x")], conversationId: "c" },
			new Map<string, string>([["a__x", "a"]]),
			new Set<string>(),
		);
		expect(result.tools).toHaveLength(0);
		expect(result.conversationId).toBe("c");
		expect(result.cwd).toBeUndefined();
		expect(result.computerId).toBeUndefined();
	});

	it("preserves computerId when set (mirrors cwd/conversationId preservation)", () => {
		const toolToServer = new Map<string, string>([["a__x", "a"]]);
		const connected = new Set<string>(["a"]);

		const result = filterMcpTools(
			{
				tools: [stubTool("a__x"), stubTool("other")],
				cwd: "/p",
				computerId: "ssh-host",
				conversationId: "c",
			},
			toolToServer,
			connected,
		);

		expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]);
		expect(result.computerId).toBe("ssh-host");
		expect(result.cwd).toBe("/p");
		expect(result.conversationId).toBe("c");
	});
});

// ---------------------------------------------------------------------------
// In-memory MCP server + fake spawn (integration)
// ---------------------------------------------------------------------------

interface FakeServer {
	tools: McpToolInfo[];
	failInitialize: boolean;
	emitListChanged: () => void;
}

/**
 * Build a fake spawn backed by a shared `FakeServer`. Incoming framed
 * JSON-RPC on stdin is answered with framed responses on stdout, exercising
 * transport → framing → rpc → client → manager end to end.
 */
function makeFakeSpawn(server: FakeServer): SpawnProcess {
	const decoder = new FrameDecoder();
	let dataListeners: Array<(data: Uint8Array) => void> = [];

	const emit = (frame: Uint8Array) => {
		for (const cb of dataListeners) cb(frame);
	};

	const spawn: SpawnProcess = (_command, _opts) => {
		// Each spawn is a fresh process; reset listeners so a reconnect (after
		// shutdown) doesn't feed closed rpc instances.
		dataListeners = [];
		const process: SpawnedProcess = {
			stdin: {
				write: (bytes: Uint8Array) => {
					for (const msg of decoder.decode(bytes)) {
						const parsed = JSON.parse(msg) as {
							id?: number;
							method?: string;
							params?: unknown;
						};
						const id = parsed.id ?? 0;
						const method = parsed.method;
						if (method === "initialize") {
							if (server.failInitialize) {
								emit(
									encode(
										JSON.stringify({
											jsonrpc: "2.0",
											id,
											error: { code: -32603, message: "initialize failed" },
										}),
									),
								);
							} else {
								emit(
									encode(
										JSON.stringify({
											jsonrpc: "2.0",
											id,
											result: {
												protocolVersion: "2025-11-25",
												capabilities: { tools: { listChanged: true } },
												serverInfo: { name: "fake", version: "0.0.0" },
											},
										}),
									),
								);
							}
						} else if (method === "tools/list") {
							emit(encode(JSON.stringify({ jsonrpc: "2.0", id, result: { tools: server.tools } })));
						} else if (method === "tools/call") {
							emit(
								encode(
									JSON.stringify({
										jsonrpc: "2.0",
										id,
										result: { content: [{ type: "text", text: "ok" }], isError: false },
									}),
								),
							);
						}
						// notifications (notifications/initialized): no response.
					}
				},
			},
			stdout: {
				on: (event: string, cb: (data: Uint8Array) => void) => {
					if (event === "data") dataListeners.push(cb);
				},
			},
			pid: 7000,
			kill: () => {},
		};
		return process;
	};

	server.emitListChanged = () => {
		emit(encode(JSON.stringify({ jsonrpc: "2.0", method: "notifications/tools/list_changed" })));
	};

	return spawn;
}

// ---------------------------------------------------------------------------
// Minimal fake HostAPI
// ---------------------------------------------------------------------------

function makeFakeHost(): {
	host: HostAPI;
	tools: Map<string, ToolContract>;
	getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null;
	getService: () => unknown;
} {
	const tools = new Map<string, ToolContract>();
	let filterFn: ((a: ToolAssembly) => Promise<ToolAssembly>) | null = null;
	let service: unknown = null;

	const noopSpan = {
		id: "s",
		log: {} as Logger,
		setAttributes: () => {},
		addLink: () => {},
		child: () => noopSpan,
		end: () => {},
	};
	const noopLogger: Logger = {
		info: () => {},
		warn: () => {},
		error: () => {},
		debug: () => {},
		child: () => noopLogger,
		span: () => noopSpan,
	};

	const host = {
		defineTool: (t: ToolContract) => {
			tools.set(t.name, t);
		},
		addFilter: (_hook: typeof toolsFilter, fn: (a: ToolAssembly) => Promise<ToolAssembly>) => {
			filterFn = fn;
			return () => {
				filterFn = null;
			};
		},
		provideService: (_handle: unknown, impl: unknown) => {
			service = impl;
		},
		getService: () => service,
		getTools: () => tools,
		logger: noopLogger,
	} as unknown as HostAPI;

	return { host, tools, getFilter: () => filterFn, getService: () => service };
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const dispatchConfig = (servers: Record<string, unknown>): string => JSON.stringify({ servers });

const tool = (name: string, description = name): McpToolInfo => ({
	name,
	description,
	inputSchema: { type: "object" },
});

const assembly = (tools: ToolContract[], cwd = "/proj"): ToolAssembly => ({
	tools,
	cwd,
	conversationId: "conv-1",
});

const flush = () => new Promise((r) => setTimeout(r, 0));

function makeServer(initialTools: McpToolInfo[]): FakeServer {
	return { tools: [...initialTools], failInitialize: false, emitListChanged: () => {} };
}

function makeExt(server: FakeServer, configJson: string): Extension {
	return makeMcpExtension({
		spawn: makeFakeSpawn(server),
		readFile: async (path) => (path.endsWith(".dispatch/mcp.json") ? configJson : null),
		getCwd: () => "/proj",
	});
}

// ---------------------------------------------------------------------------
// Lifecycle tests
// ---------------------------------------------------------------------------
describe("mcp extension lifecycle", () => {
	/** Get the registered filter, throwing if activation did not register one. */
	function requireFilter(getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null) {
		const filter = getFilter();
		if (!filter) throw new Error("toolsFilter was not registered");
		return filter;
	}

	/** Look up a registered tool, throwing if absent. */
	function requireTool(tools: Map<string, ToolContract>, name: string) {
		const t = tools.get(name);
		if (!t) throw new Error(`tool ${name} not registered`);
		return t;
	}

	it("registers tools on connect", async () => {
		const server = makeServer([tool("create_object", "Create an object")]);
		const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } }));
		const { host, tools, getFilter } = makeFakeHost();

		ext.activate(host);
		const filter = requireFilter(getFilter);

		// Running the filter triggers lazy connect + register.
		await filter(assembly([]));

		expect(tools.has("freecad__create_object")).toBe(true);
		const t = requireTool(tools, "freecad__create_object");
		expect(t.description).toBe("[freecad] Create an object");
		expect(t.concurrencySafe).toBe(false);
		ext.deactivate?.();
	});

	it("toolsFilter keeps connected-server tools and removes disconnected-server tools", async () => {
		const server = makeServer([tool("create_object")]);
		const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } }));
		const { host, tools, getFilter } = makeFakeHost();
		ext.activate(host);
		const filter = requireFilter(getFilter);

		// Connect + register the tool.
		await filter(assembly([]));
		const registered = requireTool(tools, "freecad__create_object");

		// Connected server → tool passes through the filter.
		const kept = await filter(assembly([registered]));
		expect(kept.tools.map((t) => t.name)).toContain("freecad__create_object");

		// Disconnect: deactivate shuts down the client (clearing it), then make
		// the server fail to reconnect. toolToServer still maps the tool, so the
		// filter drops it because the server is no longer connected.
		ext.deactivate?.();
		server.failInitialize = true;

		const removed = await filter(assembly([registered]));
		expect(removed.tools.map((t) => t.name)).not.toContain("freecad__create_object");
	});

	it("re-registers tools on list_changed", async () => {
		const server = makeServer([tool("first_tool")]);
		const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } }));
		const { host, tools, getFilter } = makeFakeHost();
		ext.activate(host);
		const filter = requireFilter(getFilter);

		await filter(assembly([]));
		expect(tools.has("freecad__first_tool")).toBe(true);

		// Server changes its tool set, then announces list_changed.
		server.tools = [tool("first_tool"), tool("second_tool", "The second")];
		server.emitListChanged();

		// Let the async onToolsChanged handler (re-list + re-register) flush.
		await flush();

		expect(tools.has("freecad__second_tool")).toBe(true);
		expect(requireTool(tools, "freecad__second_tool").description).toBe("[freecad] The second");
		ext.deactivate?.();
	});

	it("deactivate shuts down all clients", async () => {
		const server = makeServer([tool("create_object")]);
		const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } }));
		const { host, getFilter, getService } = makeFakeHost();
		ext.activate(host);
		const filter = requireFilter(getFilter);

		await filter(assembly([]));

		const service = getService() as {
			status: (cwd: string) => Promise<readonly { state: string }[]>;
		};
		const before = await service.status("/proj");
		expect(before[0].state).toBe("connected");

		ext.deactivate?.();

		const after = await service.status("/proj");
		expect(after[0].state).toBe("disconnected");
	});
});