pi-ext

import { afterEach, describe, expect, it, vi } from "vitest";
import { parallelSearch } from "../parallel";

afterEach(() => vi.restoreAllMocks());

function server(toolResult: unknown, sse = false) {
	return vi
		.spyOn(globalThis, "fetch")
		.mockImplementation(async (_url, init) => {
			const request = JSON.parse(String(init?.body));
			if (request.method === "notifications/initialized")
				return new Response(null, { status: 202 });
			const result =
				request.method === "initialize"
					? { protocolVersion: "2024-11-05" }
					: toolResult;
			const message = { jsonrpc: "2.0", id: request.id, result };
			const body = sse
				? `event: message\r\ndata: {"jsonrpc":"2.0","method":"notifications/progress"}\r\n\r\nevent: message\r\ndata: ${JSON.stringify(message)}\r\n\r\n`
				: JSON.stringify(message);
			return new Response(body, {
				headers: { "mcp-session-id": "test-session" },
			});
		});
}

const entry = {
	title: "Docs",
	url: "https://example.com/docs",
	excerpts: ["a".repeat(30_000), "second excerpt"],
	publish_date: "2026-09-01",
};

describe("Parallel anonymous search", () => {
	it.each([false, true])(
		"handles MCP JSON/SSE=%s, session headers, and full excerpts",
		async (sse) => {
			const network = server(
				{
					content: [
						{
							type: "text",
							text: JSON.stringify({
								results: [
									entry,
									{ ...entry, url: "https://example.com/other" },
								],
							}),
						},
					],
				},
				sse,
			);
			const result = await parallelSearch("docs", undefined, 1);
			expect(result).toEqual([
				{
					title: "Docs",
					url: entry.url,
					snippet: entry.excerpts.join("\n\n"),
					date: entry.publish_date,
				},
			]);
			expect(network).toHaveBeenCalledTimes(3);
			for (const [url, init] of network.mock.calls) {
				expect(url).toBe("https://search.parallel.ai/mcp");
				expect(init?.redirect).toBe("error");
				expect(init?.headers).not.toHaveProperty("Authorization");
			}
			const request = network.mock.calls[2][1];
			expect(new Headers(request?.headers).get("mcp-session-id")).toBe(
				"test-session",
			);
			expect(JSON.parse(String(request?.body)).params).toEqual({
				name: "web_search",
				arguments: { objective: "docs", search_queries: ["docs"] },
			});
		},
	);

	it("accepts a genuine empty result set", async () => {
		server({ content: [{ type: "text", text: '{"results":[]}' }] });
		expect(await parallelSearch("no matches")).toEqual([]);
	});

	it.each([
		[
			{ isError: true, content: [{ type: "text", text: "rate limited" }] },
			"rate limited",
		],
		[{ content: [] }, "no content"],
		[{ content: [{ type: "text", text: "not json" }] }, "JSON"],
		[{ content: [{ type: "text", text: "{}" }] }, "invalid results"],
		[
			{
				content: [
					{
						type: "text",
						text: '{"results":[{"url":"https://example.com","excerpts":[3]}]}',
					},
				],
			},
			"malformed result",
		],
	])("rejects tool errors and malformed output", async (result, error) => {
		server(result);
		await expect(parallelSearch("docs")).rejects.toThrow(error as string);
	});

	it.each([401, 429, 503])("rejects HTTP %s", async (status) => {
		vi.spyOn(globalThis, "fetch").mockResolvedValue(
			new Response("error", { status }),
		);
		await expect(parallelSearch("docs")).rejects.toThrow(`HTTP ${status}`);
	});

	it("rejects JSON-RPC errors", async () => {
		vi.spyOn(globalThis, "fetch").mockResolvedValue(
			new Response('{"id":1,"error":{"message":"unavailable"}}'),
		);
		await expect(parallelSearch("docs")).rejects.toThrow(
			"Parallel MCP: unavailable",
		);
	});

	it("rejects mismatched response IDs", async () => {
		vi.spyOn(globalThis, "fetch").mockResolvedValue(
			new Response('{"id":99,"result":{}}'),
		);
		await expect(parallelSearch("docs")).rejects.toThrow(
			"no matching MCP response",
		);
	});

	it("cancels before I/O", async () => {
		const network = vi.spyOn(globalThis, "fetch");
		await expect(
			parallelSearch("docs", AbortSignal.abort(new Error("cancelled"))),
		).rejects.toThrow("cancelled");
		expect(network).not.toHaveBeenCalled();
	});

	it("enforces a deadline even with a caller signal", async () => {
		vi.spyOn(globalThis, "fetch").mockImplementation(
			(_url, init) =>
				new Promise((_resolve, reject) => {
					init?.signal?.addEventListener(
						"abort",
						() => reject(init.signal?.reason),
						{ once: true },
					);
				}),
		);
		await expect(
			parallelSearch("docs", new AbortController().signal, 5, 10),
		).rejects.toThrow(/timeout/i);
	});
});