my-pi/pi-web-search/test/exa-client.test.ts

import assert from "node:assert/strict";
import test from "node:test";

import { ExaError, exaFetchPages, exaSearch, parseRpcPayload, resetExaSession } from "../clients/exa.ts";

function sse(payload: Record<string, unknown>): string {
	return `event: message\ndata: ${JSON.stringify(payload)}\n\n`;
}

interface RecordedCall {
	method?: string;
	params?: Record<string, unknown>;
	headers: Record<string, string>;
}

/**
 * A fetch stub that speaks the Exa endpoint protocol: initialize answers with
 * the scripted session id and the request's own id, notifications get a bare
 * 202, and tools/call delegates to the supplied handler (default 404).
 */
function exaStub(options: {
	sessionIds?: string[];
	onToolCall?: (args: Record<string, unknown>, id: number, session: string | undefined) => Response;
} = {}) {
	const calls: RecordedCall[] = [];
	let sessionIndex = 0;
	const fetchImpl = async (_url: string | URL, init?: RequestInit): Promise<Response> => {
		const headers = (init?.headers ?? {}) as Record<string, string>;
		const body = JSON.parse(String(init?.body)) as { id?: number; method?: string; params?: Record<string, unknown> };
		calls.push({ method: body.method, params: body.params, headers });
		if (body.method === "initialize") {
			const sid = options.sessionIds?.[Math.min(sessionIndex++, (options.sessionIds?.length ?? 1) - 1)] ?? `sid-${sessionIndex}`;
			return new Response(
				sse({ jsonrpc: "2.0", id: body.id, result: { protocolVersion: "2025-03-26", capabilities: {}, serverInfo: { name: "exa" } } }),
				{ status: 200, headers: { "mcp-session-id": sid, "content-type": "text/event-stream" } },
			);
		}
		if (body.method === "notifications/initialized") return new Response("", { status: 202 });
		if (body.method === "tools/call") {
			return options.onToolCall?.(body.params ?? {}, body.id ?? 0, headers["mcp-session-id"]) ?? new Response("Not Found", { status: 404 });
		}
		return new Response("unexpected", { status: 500 });
	};
	return { fetchImpl, calls };
}

function okToolResult(id: number, text: string): Response {
	return new Response(
		sse({ jsonrpc: "2.0", id, result: { content: [{ type: "text", text }] } }),
		{ status: 200, headers: { "content-type": "text/event-stream" } },
	);
}

test("parseRpcPayload reads JSON bodies and SSE data lines, matching the request id", () => {
	const direct = JSON.stringify({ jsonrpc: "2.0", id: 7, result: { ok: true } });
	assert.deepEqual(parseRpcPayload(direct, 7).result, { ok: true });

	const body =
		`event: message\ndata: {"jsonrpc":"2.0","id":1,"result":{"a":1}}\n\n` +
		`event: message\ndata: {"jsonrpc":"2.0","id":2,"result":{"b":2}}\n\n`;
	assert.deepEqual(parseRpcPayload(body, 2).result, { b: 2 });

	assert.throws(() => parseRpcPayload("not json at all", 1), ExaError);
	assert.throws(() => parseRpcPayload(sse({ jsonrpc: "2.0", id: 99, result: {} }), 1), ExaError);
});

test("exaSearch initializes a session, notifies, then posts tools/call with the session id", async () => {
	resetExaSession();
	const { fetchImpl, calls } = exaStub({
		sessionIds: ["sid-A"],
		onToolCall: (_args, id) => okToolResult(id, "search results"),
	});

	const text = await exaSearch({ query: "docs", numResults: 3, includeDomains: ["python.org"] }, undefined, fetchImpl);
	assert.equal(text, "search results");
	assert.deepEqual(calls.map((c) => c.method), ["initialize", "notifications/initialized", "tools/call"]);

	const call = calls[2]!;
	assert.equal(call.headers["mcp-session-id"], "sid-A");
	assert.equal((call.params as { name: string }).name, "web_search_advanced_exa");
	const args = (call.params as { arguments: Record<string, unknown> }).arguments;
	assert.equal(args.numResults, 3);
	assert.equal(args.type, "auto");
	assert.deepEqual(args.includeDomains, ["python.org"]);
	assert.ok((args.textMaxCharacters as number) > 0);
});

test("exaSearch reuses the cached session across calls", async () => {
	resetExaSession();
	const { fetchImpl, calls } = exaStub({
		sessionIds: ["sid-1"],
		onToolCall: (args, id) => okToolResult(id, `result for ${(args as { name: string }).name}`),
	});

	await exaSearch({ query: "a" }, undefined, fetchImpl);
	const text = await exaFetchPages(["https://example.com"], 500, undefined, fetchImpl);
	assert.match(text, /web_fetch_exa/);
	assert.equal(calls.length, 4);
	assert.equal((calls[3]!.params as { name: string }).name, "web_fetch_exa");
	assert.equal(calls[3]!.headers["mcp-session-id"], "sid-1");
	assert.equal((calls[3]!.params as { arguments: { maxCharacters: number } }).arguments.maxCharacters, 500);
});

test("a 404 from tools/call drops the session, re-initializes, and retries once", async () => {
	resetExaSession();
	let toolCalls = 0;
	const { fetchImpl, calls } = exaStub({
		sessionIds: ["stale-sid", "fresh-sid"],
		onToolCall: (_args, id) => {
			toolCalls += 1;
			return toolCalls === 1 ? new Response("Not Found", { status: 404 }) : okToolResult(id, "after re-init");
		},
	});

	const text = await exaSearch({ query: "retry me" }, undefined, fetchImpl);
	assert.equal(text, "after re-init");
	assert.deepEqual(calls.map((c) => c.method), [
		"initialize",
		"notifications/initialized",
		"tools/call",
		"initialize",
		"notifications/initialized",
		"tools/call",
	]);
	assert.equal(calls[5]!.headers["mcp-session-id"], "fresh-sid");
});

test("isError tool results and HTTP failures raise ExaError", async () => {
	resetExaSession();
	const isError = exaStub({
		onToolCall: (_args, id) =>
			new Response(sse({ jsonrpc: "2.0", id, result: { isError: true, content: [{ type: "text", text: "bad args" }] } }), { status: 200 }),
	});
	await assert.rejects(exaFetchPages(["https://example.com"], undefined, undefined, isError.fetchImpl), /bad args/);

	resetExaSession();
	const httpFail = exaStub({ onToolCall: () => new Response("boom", { status: 503 }) });
	// Session init succeeds; tools/call hits a non-404 HTTP error.
	await assert.rejects(exaSearch({ query: "x" }, undefined, httpFail.fetchImpl), /HTTP 503/);
});