pi-ext

/** Single-call MCP sessions shared by search providers, not a general MCP client. */
export interface McpToolResult {
	isError?: boolean;
	content?: Array<{ type?: string; text?: string }>;
}

export class McpError extends Error {
	readonly status: number;
	readonly code?: number | string;
	readonly response?: Response;

	constructor(
		message: string,
		status: number,
		code?: number | string,
		response?: Response,
	) {
		super(message);
		this.name = "McpError";
		this.status = status;
		this.code = code;
		this.response = response;
	}
}

interface RpcMessage {
	id?: number;
	result?: unknown;
	error?: { code?: number | string; message?: string };
}

function parseMessage(text: string, id: number, status: number): RpcMessage {
	let messages: RpcMessage[];
	try {
		messages = text.trimStart().startsWith("{")
			? [JSON.parse(text)]
			: text.split(/\r?\n\r?\n/).flatMap((event) => {
					const data = event
						.split(/\r?\n/)
						.filter((line) => line.startsWith("data:"))
						.map((line) => line.slice(5).trimStart())
						.join("\n");
					return data && data !== "[DONE]" ? [JSON.parse(data)] : [];
				});
	} catch {
		throw new McpError(`Malformed MCP response: ${text.slice(0, 500)}`, status);
	}
	const message = messages.find((entry) => entry?.id === id);
	if (!message) throw new McpError("no matching MCP response", status);
	if (message.error)
		throw new McpError(
			message.error.message ?? "request failed",
			status,
			message.error.code,
		);
	return message;
}

export async function callMcpTool(options: {
	url: string;
	protocolVersion: string;
	headers?: Record<string, string>;
	signal: AbortSignal;
	name: string;
	arguments: Record<string, unknown>;
}): Promise<McpToolResult> {
	let sessionId: string | undefined;
	let protocolVersion = options.protocolVersion;
	async function post(
		method: string,
		params: object,
		id?: number,
	): Promise<unknown> {
		options.signal.throwIfAborted();
		const response = await fetch(options.url, {
			method: "POST",
			redirect: "error",
			headers: {
				...options.headers,
				"Content-Type": "application/json",
				Accept: "application/json, text/event-stream",
				"MCP-Protocol-Version": protocolVersion,
				...(sessionId ? { "mcp-session-id": sessionId } : {}),
			},
			body: JSON.stringify({
				jsonrpc: "2.0",
				...(id === undefined ? {} : { id }),
				method,
				params,
			}),
			signal: options.signal,
		});
		if (!response.ok)
			throw new McpError(
				`HTTP ${response.status}: ${response.statusText}`,
				response.status,
				undefined,
				response,
			);
		sessionId = response.headers.get("mcp-session-id") ?? sessionId;
		const text = await response.text();
		if (id === undefined) return undefined;
		const result = parseMessage(text, id, response.status).result;
		if (method === "initialize") {
			const initialized = result as { protocolVersion?: unknown } | undefined;
			if (
				typeof initialized?.protocolVersion !== "string" ||
				!initialized.protocolVersion
			) {
				throw new McpError("MCP initialization failed", response.status);
			}
			protocolVersion = initialized.protocolVersion;
		}
		return result;
	}

	await post(
		"initialize",
		{
			protocolVersion,
			capabilities: {},
			clientInfo: { name: "pi-web-extension", version: "0.1.0" },
		},
		1,
	);
	await post("notifications/initialized", {});
	const result = await post(
		"tools/call",
		{ name: options.name, arguments: options.arguments },
		2,
	);
	if (!result || typeof result !== "object" || Array.isArray(result)) {
		throw new McpError("MCP tool returned an invalid result", 200);
	}
	return result as McpToolResult;
}