pi-ext

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
import { afterEach, describe, expect, it, vi } from "vitest";
import type { KlausTool } from "../src/protocol";
import { TOOL_ID_FIELD, ToolBridge } from "../src/tool-bridge";

const echo: KlausTool = {
	name: "echo",
	description: "Echo one value exactly.",
	inputSchema: {
		type: "object",
		properties: { value: { type: "string" } },
		required: ["value"],
		additionalProperties: false,
	},
};

const cleanup: Array<() => Promise<void>> = [];

async function connectedBridge(
	tools: KlausTool[] = [echo],
	onActivity: () => void = () => undefined,
): Promise<{ bridge: ToolBridge; client: Client }> {
	const bridge = new ToolBridge(tools, onActivity);
	const client = new Client({ name: "klaus-test", version: "1" });
	const [clientTransport, serverTransport] =
		InMemoryTransport.createLinkedPair();
	await bridge.server.connect(serverTransport);
	await client.connect(clientTransport);
	cleanup.push(async () => {
		await client.close().catch(() => undefined);
		await bridge.close("Test cleanup.").catch(() => undefined);
	});
	return { bridge, client };
}

async function park(
	bridge: ToolBridge,
	client: Client,
	id: string,
	input: Record<string, unknown>,
	name = "echo",
) {
	bridge.allow(name, input, id);
	const result = client.callTool({ name, arguments: input });
	await bridge.waitForPending([id]);
	return result;
}

afterEach(async () => {
	await Promise.all(cleanup.splice(0).map((close) => close()));
});

describe("Klaus query-owned MCP tool bridge", () => {
	it("lists exact Pi names, descriptions, schemas, and order", async () => {
		const second: KlausTool = {
			name: "second tool",
			description: "Second description.",
			inputSchema: { type: "object", properties: {} },
		};
		const { bridge, client } = await connectedBridge([echo, second]);

		expect(await client.listTools()).toMatchObject({ tools: [echo, second] });
		expect(bridge.toolNames()).toEqual(["echo", "second tool"]);
	});

	it("keeps the first identity for a repeated tool use id", async () => {
		const { bridge, client } = await connectedBridge();
		/** The `PreToolUse` hook reports every call, then `canUseTool` reports the
		 * same call again whenever permissions prompt. */
		bridge.allow("echo", { value: "hook" }, "call-1");
		bridge.allow("echo", { value: "hook" }, "call-1");

		const result = client.callTool({
			name: "echo",
			arguments: { value: "hook", [TOOL_ID_FIELD]: "call-1" },
		});
		await bridge.waitForPending(["call-1"]);
		expect(bridge.pendingIds()).toEqual(["call-1"]);
		expect(
			bridge.settle({
				id: "call-1",
				content: [{ type: "text", text: "hook" }],
				isError: false,
			}),
		).toBe(true);

		expect(await result).toMatchObject({
			isError: false,
			content: [{ type: "text", text: "hook" }],
		});
	});

	it("truncates a tool result above Claude's result size ceiling", async () => {
		const { bridge, client } = await connectedBridge();
		const oversized = `${"x".repeat(600_000)}TAIL`;
		bridge.allow("echo", { value: "big" }, "call-big");
		const result = client.callTool({
			name: "echo",
			arguments: { value: "big", [TOOL_ID_FIELD]: "call-big" },
		});
		await bridge.waitForPending(["call-big"]);
		bridge.settle({
			id: "call-big",
			content: [{ type: "text", text: oversized }],
			isError: false,
		});

		const settled = (await result) as {
			content: Array<{ type: string; text: string }>;
		};
		const text = settled.content[0]?.text ?? "";
		expect(text.length).toBeLessThanOrEqual(500_000);
		expect(text).toContain("Klaus truncated this tool result");
		expect(text).not.toContain("TAIL");
	});

	it("raises the per-tool result size ceiling above Claude's default", async () => {
		const { client } = await connectedBridge();

		const listed = await client.listTools();

		/** Claude persists any MCP result above this limit to a file Klaus deletes,
		 * so every tool must declare the ceiling instead of the 100000 default. */
		for (const tool of listed.tools) {
			expect(tool._meta).toEqual({
				"anthropic/maxResultSizeChars": 500_000,
			});
		}
	});

	it("counts tool listing, permission, call, and settlement as activity", async () => {
		const activity = vi.fn();
		const { bridge, client } = await connectedBridge([echo], activity);
		await client.listTools();
		const result = park(bridge, client, "call-1", { value: "hello" });
		await bridge.waitForPending(["call-1"]);
		expect(
			bridge.settle({
				id: "call-1",
				content: [{ type: "text", text: "hello" }],
				isError: false,
			}),
		).toBe(true);
		await result;

		expect(activity).toHaveBeenCalledTimes(4);
	});

	it("parks an allowed call and forwards text, image, and error state", async () => {
		const { bridge, client } = await connectedBridge();
		const result = park(bridge, client, "call-1", { value: "hello" });
		await bridge.waitForPending(["call-1"]);

		expect(bridge.pendingIds()).toEqual(["call-1"]);
		expect(
			bridge.settle({
				id: "call-1",
				content: [
					{ type: "text", text: "failed" },
					{ type: "image", data: "base64", mimeType: "image/png" },
				],
				isError: true,
			}),
		).toBe(true);
		expect(await result).toMatchObject({
			content: [
				{ type: "text", text: "failed" },
				{ type: "image", data: "base64", mimeType: "image/png" },
			],
			isError: true,
		});
		expect(bridge.pendingIds()).toEqual([]);
	});

	it("matches fallback identity despite recursively reordered object keys", async () => {
		const { bridge, client } = await connectedBridge();
		const allowed = {
			outer: { beta: [1, { right: true, left: false }], alpha: null },
			value: "hello",
		};
		bridge.allow("echo", allowed, "call-1");
		const result = client.callTool({
			name: "echo",
			arguments: {
				value: "hello",
				outer: { alpha: null, beta: [1, { left: false, right: true }] },
			},
		});
		await bridge.waitForPending(["call-1"]);
		bridge.settle({ id: "call-1", content: [], isError: false });

		expect(await result).toMatchObject({ isError: false });
	});

	it("uses the embedded tool ID and removes it before parking", async () => {
		const { bridge, client } = await connectedBridge();
		bridge.allow("mcp__klaus__echo", { value: "original" }, "call-exact");
		const result = client.callTool({
			name: "echo",
			arguments: { value: "updated", [TOOL_ID_FIELD]: "call-exact" },
		});
		await bridge.waitForPending(["call-exact"]);

		expect(bridge.pending.get("call-exact")?.input).toEqual({
			value: "original",
		});
		bridge.settle({ id: "call-exact", content: [], isError: false });
		expect(await result).toMatchObject({ isError: false });
	});

	it("maps prefixed Claude names back to original Pi names", async () => {
		const { bridge, client } = await connectedBridge();
		bridge.allow("mcp__klaus__echo", { value: "hello" }, "call-1");
		const result = client.callTool({
			name: "mcp__klaus__echo",
			arguments: { value: "hello" },
		});
		await bridge.waitForPending(["call-1"]);
		bridge.settle({ id: "call-1", content: [], isError: false });

		expect(await result).toMatchObject({ isError: false });
	});

	it("rejects inactive, unpermitted, and mismatched tool identities", async () => {
		const { bridge, client } = await connectedBridge();
		bridge.allow("echo", { value: "expected" }, "call-1");

		const inactive = await client.callTool({ name: "missing", arguments: {} });
		const unpermitted = await client.callTool({
			name: "echo",
			arguments: { value: "other" },
		});
		const wrongEmbeddedId = await client.callTool({
			name: "echo",
			arguments: { value: "expected", [TOOL_ID_FIELD]: "wrong" },
		});

		expect(inactive).toMatchObject({
			isError: true,
			content: [{ type: "text", text: "Tool missing is no longer active." }],
		});
		expect(unpermitted).toMatchObject({
			isError: true,
			content: [{ type: "text", text: "No Klaus tool identity for echo." }],
		});
		expect(wrongEmbeddedId).toMatchObject({
			isError: true,
			content: [{ type: "text", text: "No Klaus tool identity for echo." }],
		});
		expect(bridge.pendingIds()).toEqual([]);
	});

	it("preserves parallel calls and accepts results in any order", async () => {
		const { bridge, client } = await connectedBridge();
		const first = park(bridge, client, "first", { value: "one" });
		const second = park(bridge, client, "second", { value: "two" });
		await bridge.waitForPending(["first", "second"]);

		expect(new Set(bridge.pendingIds())).toEqual(new Set(["first", "second"]));
		expect(
			bridge.settle({
				id: "second",
				content: [{ type: "text", text: "two" }],
				isError: false,
			}),
		).toBe(true);
		expect(
			bridge.settle({
				id: "first",
				content: [{ type: "text", text: "one" }],
				isError: false,
			}),
		).toBe(true);
		expect(await first).toMatchObject({
			content: [{ type: "text", text: "one" }],
		});
		expect(await second).toMatchObject({
			content: [{ type: "text", text: "two" }],
		});
	});

	it("rejects a duplicate in-flight call ID without disturbing the original", async () => {
		const { bridge, client } = await connectedBridge();
		bridge.allow("echo", { value: "one" }, "duplicate");
		const first = client.callTool({
			name: "echo",
			arguments: { value: "one", [TOOL_ID_FIELD]: "duplicate" },
		});
		await bridge.waitForPending(["duplicate"]);
		/** Claude re-dispatches an identity after a transport gap. */
		bridge.allow("echo", { value: "two" }, "duplicate");
		const duplicate = await client.callTool({
			name: "echo",
			arguments: { value: "two", [TOOL_ID_FIELD]: "duplicate" },
		});

		expect(duplicate).toMatchObject({
			isError: true,
			content: [{ type: "text", text: "Duplicate Klaus tool call duplicate." }],
		});
		bridge.settle({ id: "duplicate", content: [], isError: false });
		expect(await first).toMatchObject({ isError: false });
	});

	it("settles each call ID exactly once", async () => {
		const { bridge, client } = await connectedBridge();
		const result = park(bridge, client, "call-1", { value: "hello" });
		await bridge.waitForPending(["call-1"]);
		const settlement = {
			id: "call-1",
			content: [{ type: "text" as const, text: "hello" }],
			isError: false,
		};

		expect(bridge.settle({ ...settlement, id: "unknown" })).toBe(false);
		expect(bridge.settle(settlement)).toBe(true);
		expect(bridge.settle(settlement)).toBe(false);
		await result;
	});

	it("waits until every requested call is parked", async () => {
		const { bridge, client } = await connectedBridge();
		let resolved = false;
		const waiting = bridge.waitForPending(["first", "second"]).then(() => {
			resolved = true;
		});
		const first = park(bridge, client, "first", { value: "one" });
		await bridge.waitForPending(["first"]);
		await Promise.resolve();
		expect(resolved).toBe(false);
		const second = park(bridge, client, "second", { value: "two" });
		await waiting;
		expect(resolved).toBe(true);
		bridge.settle({ id: "first", content: [], isError: false });
		bridge.settle({ id: "second", content: [], isError: false });
		await Promise.all([first, second]);
	});

	it("aborts a pending wait without requiring MCP activity", async () => {
		const { bridge } = await connectedBridge();
		const controller = new AbortController();
		const waiting = bridge.waitForPending(["never"], controller.signal);
		controller.abort();

		await expect(waiting).rejects.toThrow("Klaus tool wait aborted");
	});

	it("close is idempotent, rejects parked calls, and wakes pending waits", async () => {
		const { bridge, client } = await connectedBridge();
		const call = park(bridge, client, "call-1", { value: "hello" });
		await bridge.waitForPending(["call-1"]);
		const waiting = bridge.waitForPending(["never"]);

		await Promise.all([
			bridge.close("Deliberate close."),
			bridge.close("Ignored second reason."),
		]);
		await expect(call).rejects.toThrow();
		await expect(waiting).rejects.toThrow(
			"Klaus query closed before tools parked",
		);
		expect(bridge.pendingIds()).toEqual([]);
		expect(bridge.settle({ id: "call-1", content: [], isError: false })).toBe(
			false,
		);
	});

	it("detects any tool list, order, description, or schema change", async () => {
		const { bridge } = await connectedBridge();
		expect(bridge.hasSchemaChange([structuredClone(echo)])).toBe(false);
		expect(bridge.hasSchemaChange([])).toBe(true);
		expect(bridge.hasSchemaChange([{ ...echo, description: "changed" }])).toBe(
			true,
		);
		expect(
			bridge.hasSchemaChange([
				{
					...echo,
					inputSchema: { ...echo.inputSchema, additionalProperties: true },
				},
			]),
		).toBe(true);
		expect(bridge.hasSchemaChange([echo, { ...echo, name: "second" }])).toBe(
			true,
		);
	});
});