pi-ext

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import {
	CallToolRequestSchema,
	type CallToolResult,
	ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { dbg } from "./debug.js";
import type { KlausTool } from "./protocol.js";

interface PermissionCall {
	id: string;
	name: string;
	input: Record<string, unknown>;
}

interface PendingCall extends PermissionCall {
	resolve(result: CallToolResult): void;
	reject(error: Error): void;
}

export interface PiToolResult {
	id: string;
	content: Array<
		| { type: "text"; text: string }
		| { type: "image"; data: string; mimeType: string }
	>;
	isError: boolean;
}

export const TOOL_ID_FIELD = "__klausToolUseID";

/** Claude caps MCP tool results at 100000 characters by default and replaces
 * anything larger with a `<persisted-output>` pointer into its own temp
 * directory, which Klaus deletes when the query ends. Declaring the per-tool
 * limit raises that threshold to Claude's ceiling so Pi results reach the model
 * intact; Pi already truncates its own tool output. */
const MAX_RESULT_SIZE_CHARS = 500_000;

const TRUNCATION_NOTE =
	"\n[Klaus truncated this tool result at Claude's 500000-character result limit.]";

/** Beyond the declared ceiling Claude persists the result to a file inside the
 * per-query config directory and shows the model a pointer, and Klaus deletes
 * that directory when the query ends. Truncating here keeps the loss visible to
 * the model instead of handing it a dead path. */
function boundedContent(
	content: PiToolResult["content"],
): PiToolResult["content"] {
	let remaining = MAX_RESULT_SIZE_CHARS;
	let truncated = false;
	const bounded: PiToolResult["content"] = [];
	for (const item of content) {
		if (item.type !== "text") {
			bounded.push({ ...item });
			continue;
		}
		if (truncated) continue;
		if (item.text.length <= remaining) {
			remaining -= item.text.length;
			bounded.push({ ...item });
			continue;
		}
		const kept = item.text.slice(
			0,
			Math.max(0, remaining - TRUNCATION_NOTE.length),
		);
		dbg?.("bridge.settle.truncated", { originalLength: item.text.length });
		bounded.push({ type: "text", text: `${kept}${TRUNCATION_NOTE}` });
		truncated = true;
		remaining = 0;
	}
	return bounded;
}

const originalName = (name: string): string =>
	name.startsWith("mcp__klaus__") ? name.slice("mcp__klaus__".length) : name;

function canonical(value: unknown): string {
	if (Array.isArray(value)) return `[${value.map(canonical).join(",")}]`;
	if (typeof value === "object" && value !== null) {
		return `{${Object.entries(value)
			.sort(([left], [right]) => left.localeCompare(right))
			.map(([key, item]) => `${JSON.stringify(key)}:${canonical(item)}`)
			.join(",")}}`;
	}
	return JSON.stringify(value) ?? "undefined";
}

const sameInput = (
	left: Record<string, unknown>,
	right: Record<string, unknown>,
): boolean => canonical(left) === canonical(right);

export class ToolBridge {
	readonly server: McpServer;
	readonly pending = new Map<string, PendingCall>();
	private readonly settled = new Set<string>();
	private permissions: PermissionCall[] = [];
	private tools = new Map<string, KlausTool>();
	private closed = false;
	private waiters: Array<() => void> = [];

	constructor(
		tools: KlausTool[],
		private readonly onActivity: () => void = () => undefined,
	) {
		dbg?.("bridge.create");
		this.tools = new Map(tools.map((tool) => [tool.name, tool]));
		this.server = new McpServer(
			{ name: "klaus", version: "1" },
			{ capabilities: { tools: {} } },
		);
		this.server.server.setRequestHandler(ListToolsRequestSchema, async () => {
			dbg?.("bridge.listTools");
			this.onActivity();
			return {
				tools: [...this.tools.values()].map((tool) => ({
					name: tool.name,
					description: tool.description,
					inputSchema: tool.inputSchema,
					_meta: {
						"anthropic/maxResultSizeChars": MAX_RESULT_SIZE_CHARS,
					},
				})),
			};
		});
		this.server.server.setRequestHandler(
			CallToolRequestSchema,
			async (request) => {
				dbg?.("bridge.callTool.start", { closed: this.closed });
				this.onActivity();
				if (this.closed) return this.errorResult("Klaus query is closed.");
				const name = originalName(request.params.name);
				if (!this.tools.has(name))
					return this.errorResult(`Tool ${name} is no longer active.`);
				const input = { ...(request.params.arguments ?? {}) };
				const embeddedId = input[TOOL_ID_FIELD];
				delete input[TOOL_ID_FIELD];
				const permissionIndex = this.permissions.findIndex((call) =>
					typeof embeddedId === "string"
						? call.id === embeddedId && call.name === name
						: call.name === name && sameInput(call.input, input),
				);
				const permission =
					permissionIndex < 0
						? undefined
						: this.permissions.splice(permissionIndex, 1)[0];
				if (!permission) {
					return this.errorResult(`No Klaus tool identity for ${name}.`);
				}
				if (this.pending.has(permission.id)) {
					return this.errorResult(
						`Duplicate Klaus tool call ${permission.id}.`,
					);
				}
				dbg?.("bridge.callTool.park");
				return await new Promise<CallToolResult>((resolve, reject) => {
					const pending: PendingCall = { ...permission, resolve, reject };
					this.pending.set(permission.id, pending);
					for (const notify of this.waiters.splice(0)) notify();
				});
			},
		);
	}

	/** Claude reaches this through both the `PreToolUse` hook and `canUseTool`,
	 * and only the hook is guaranteed, so the first identity for a tool use id
	 * wins and later duplicates are ignored. */
	allow(name: string, input: Record<string, unknown>, id: string): void {
		const duplicate = this.permissions.some((call) => call.id === id);
		dbg?.("bridge.allow");
		this.onActivity();
		if (duplicate) return;
		this.permissions.push({ id, name: originalName(name), input });
	}

	/** The tool list in the exact order Claude receives it, which is the order the
	 * API reports a rejected tool schema by position. */
	activeTools(): KlausTool[] {
		return [...this.tools.values()];
	}

	toolNames(): string[] {
		const names = [...this.tools.keys()];
		dbg?.("bridge.toolNames");
		return names;
	}

	/** Pi owns the active tool list, so a live child adopts the new list instead
	 * of forcing a canonical replay. */
	replaceTools(tools: KlausTool[]): void {
		dbg?.("bridge.replaceTools");
		this.tools = new Map(tools.map((tool) => [tool.name, tool]));
	}

	hasSchemaChange(tools: KlausTool[]): boolean {
		const changed =
			JSON.stringify([...this.tools.values()]) !== JSON.stringify(tools);
		dbg?.("bridge.hasSchemaChange", { changed });
		return changed;
	}

	pendingIds(): string[] {
		const ids = [...this.pending.keys()];
		dbg?.("bridge.pendingIds");
		return ids;
	}

	async waitForPending(ids: string[], signal?: AbortSignal): Promise<void> {
		dbg?.("bridge.waitForPending.start");
		while (ids.some((id) => !this.pending.has(id))) {
			if (this.closed)
				throw new Error("Klaus query closed before tools parked.");
			if (signal?.aborted) throw new Error("Klaus tool wait aborted.");
			await new Promise<void>((resolve, reject) => {
				const abort = () => reject(new Error("Klaus tool wait aborted."));
				signal?.addEventListener("abort", abort, { once: true });
				this.waiters.push(() => {
					signal?.removeEventListener("abort", abort);
					resolve();
				});
			});
		}
		dbg?.("bridge.waitForPending.end");
	}

	settle(result: PiToolResult): boolean {
		dbg?.("bridge.settle.start");
		this.onActivity();
		if (this.settled.has(result.id)) return false;
		const pending = this.pending.get(result.id);
		if (!pending) return false;
		this.pending.delete(result.id);
		this.settled.add(result.id);
		pending.resolve({
			content: boundedContent(result.content),
			isError: result.isError,
		});
		dbg?.("bridge.settle.end");
		return true;
	}

	async close(reason: string): Promise<void> {
		dbg?.("bridge.close.start");
		if (this.closed) return;
		this.closed = true;
		const error = new Error(reason);
		for (const pending of this.pending.values()) pending.reject(error);
		this.pending.clear();
		for (const notify of this.waiters.splice(0)) notify();
		await this.server.close();
		dbg?.("bridge.close.end");
	}

	private errorResult(message: string): CallToolResult {
		dbg?.("bridge.errorResult");
		return { content: [{ type: "text", text: message }], isError: true };
	}
}