my-pi/pi-web-search/clients/exa.ts

/**
 * Exa client: direct JSON-RPC over the hosted endpoint https://mcp.exa.ai/mcp.
 * This is a plain HTTPS POST client (no MCP adapter, no SDK): initialize once,
 * cache the `mcp-session-id` response header, then POST `tools/call` per call.
 * Keyless works; the `exa` api_key credential from auth.json is sent as the
 * `x-api-key` header when present.
 */

import { getSearchApiKey } from "../auth.ts";

export class ExaError extends Error {}

const EXA_MCP_URL = "https://mcp.exa.ai/mcp?tools=web_search_exa,web_search_advanced_exa,web_fetch_exa";
const PROTOCOL_VERSION = "2025-03-26";
const SESSION_TIMEOUT_MS = 15_000;
const CALL_TIMEOUT_MS = 60_000;

/** Cap page text per result so advanced search responses stay bounded. */
export const EXA_SEARCH_TEXT_MAX_CHARACTERS = 2000;
export const EXA_FETCH_MAX_CHARACTERS = 3000;

export type FetchLike = (input: string | URL, init?: RequestInit) => Promise<Response>;

export interface ExaSearchParams {
	query: string;
	numResults?: number;
	includeDomains?: string[];
	startPublishedDate?: string;
	endPublishedDate?: string;
	textMaxCharacters?: number;
}

let cachedSessionId: string | undefined;
let pendingSession: Promise<string> | undefined;
let nextRequestId = 1;

/** Clear the cached MCP session (used on expiry and by tests). */
export function resetExaSession(): void {
	cachedSessionId = undefined;
}

function baseHeaders(): Record<string, string> {
	const headers: Record<string, string> = {
		"content-type": "application/json",
		accept: "application/json, text/event-stream",
	};
	const key = getSearchApiKey("exa");
	if (key) headers["x-api-key"] = key;
	return headers;
}

function withTimeout(signal: AbortSignal | undefined, timeoutMs: number): AbortSignal {
	const timeout = AbortSignal.timeout(timeoutMs);
	return signal ? AbortSignal.any([signal, timeout]) : timeout;
}

/**
 * Extract the JSON-RPC message for `id` from a raw JSON body or an SSE stream
 * (one or more `data: {...}` lines).
 */
export function parseRpcPayload(bodyText: string, id: number): Record<string, unknown> {
	const candidates: string[] = [];
	const trimmed = bodyText.trim();
	if (trimmed.startsWith("{")) candidates.push(trimmed);
	for (const line of trimmed.split(/\r?\n/u)) {
		if (line.startsWith("data:")) candidates.push(line.slice(5).trim());
	}
	for (const candidate of candidates) {
		try {
			const parsed = JSON.parse(candidate) as Record<string, unknown>;
			if (parsed && typeof parsed === "object" && !Array.isArray(parsed) && parsed.id === id) {
				return parsed;
			}
		} catch {
			// Not JSON — try the next candidate line.
		}
	}
	throw new ExaError("Exa endpoint returned an unreadable response.");
}

async function post(
	body: Record<string, unknown>,
	sessionId: string | undefined,
	timeoutMs: number,
	signal: AbortSignal | undefined,
	fetchImpl: FetchLike,
): Promise<{ status: number; sessionId?: string; text: string }> {
	const headers = baseHeaders();
	if (sessionId) headers["mcp-session-id"] = sessionId;
	let response: Response;
	try {
		response = await fetchImpl(EXA_MCP_URL, {
			method: "POST",
			headers,
			body: JSON.stringify(body),
			signal: withTimeout(signal, timeoutMs),
		});
	} catch (error) {
		throw new ExaError(`Could not reach the Exa endpoint: ${(error as Error).message}`);
	}
	return {
		status: response.status,
		sessionId: response.headers.get("mcp-session-id") ?? undefined,
		text: await response.text(),
	};
}

async function initialize(signal: AbortSignal | undefined, fetchImpl: FetchLike): Promise<string> {
	const id = nextRequestId++;
	const res = await post(
		{
			jsonrpc: "2.0",
			id,
			method: "initialize",
			params: {
				protocolVersion: PROTOCOL_VERSION,
				capabilities: {},
				clientInfo: { name: "my-pi-web-search", version: "0.1.0" },
			},
		},
		undefined,
		SESSION_TIMEOUT_MS,
		signal,
		fetchImpl,
	);
	if (res.status !== 200) throw new ExaError(`Exa initialize failed (HTTP ${res.status}).`);
	const sessionId = res.sessionId;
	if (!sessionId) throw new ExaError("Exa endpoint did not return an MCP session id.");
	const payload = parseRpcPayload(res.text, id);
	if (payload.error) throw new ExaError(`Exa initialize error: ${(payload.error as { message?: string }).message ?? "unknown"}`);

	const notice = await post(
		{ jsonrpc: "2.0", method: "notifications/initialized" },
		sessionId,
		SESSION_TIMEOUT_MS,
		signal,
		fetchImpl,
	);
	if (notice.status >= 400) {
		throw new ExaError(`Exa initialized notification failed (HTTP ${notice.status}).`);
	}
	return sessionId;
}

function sessionId(signal: AbortSignal | undefined, fetchImpl: FetchLike): Promise<string> {
	if (cachedSessionId) return Promise.resolve(cachedSessionId);
	pendingSession ??= initialize(signal, fetchImpl)
		.then((id) => {
			cachedSessionId = id;
			return id;
		})
		.finally(() => {
			pendingSession = undefined;
		});
	return pendingSession;
}

function contentText(result: Record<string, unknown>): string {
	const content = Array.isArray(result.content) ? (result.content as Record<string, unknown>[]) : [];
	return content
		.filter((part) => part?.type === "text" && typeof part.text === "string")
		.map((part) => part.text as string)
		.join("\n")
		.trim();
}

async function callTool(
	toolName: string,
	args: Record<string, unknown>,
	signal: AbortSignal | undefined,
	fetchImpl: FetchLike,
	allowRetry: boolean,
): Promise<Record<string, unknown>> {
	const sid = await sessionId(signal, fetchImpl);
	const id = nextRequestId++;
	const res = await post(
		{ jsonrpc: "2.0", id, method: "tools/call", params: { name: toolName, arguments: args } },
		sid,
		CALL_TIMEOUT_MS,
		signal,
		fetchImpl,
	);
	if (res.status === 404 && allowRetry) {
		// Session expired server-side: drop it, re-initialize, retry once.
		cachedSessionId = undefined;
		return callTool(toolName, args, signal, fetchImpl, false);
	}
	if (res.status !== 200) {
		throw new ExaError(`Exa ${toolName} failed (HTTP ${res.status}).`);
	}
	const payload = parseRpcPayload(res.text, id);
	if (payload.error) {
		const message = (payload.error as { message?: string }).message ?? "unknown";
		if (allowRetry && /session/iu.test(message)) {
			cachedSessionId = undefined;
			return callTool(toolName, args, signal, fetchImpl, false);
		}
		throw new ExaError(`Exa ${toolName} error: ${message}`);
	}
	const result = (payload.result ?? {}) as Record<string, unknown>;
	if (result.isError) {
		throw new ExaError(`Exa ${toolName} failed: ${contentText(result) || "unknown error"}`);
	}
	return result;
}

/** tools/call web_search_advanced_exa — precise search with bounded text per result. */
export async function exaSearch(
	params: ExaSearchParams,
	signal?: AbortSignal,
	fetchImpl: FetchLike = fetch,
): Promise<string> {
	const result = await callTool(
		"web_search_advanced_exa",
		{
			query: params.query,
			type: "auto",
			numResults: params.numResults ?? 8,
			textMaxCharacters: params.textMaxCharacters ?? EXA_SEARCH_TEXT_MAX_CHARACTERS,
			...(params.includeDomains?.length ? { includeDomains: params.includeDomains } : {}),
			...(params.startPublishedDate ? { startPublishedDate: params.startPublishedDate } : {}),
			...(params.endPublishedDate ? { endPublishedDate: params.endPublishedDate } : {}),
		},
		signal,
		fetchImpl,
		true,
	);
	return contentText(result) || "Exa returned no content.";
}

/** tools/call web_fetch_exa — full page content as clean markdown for a batch of URLs. */
export async function exaFetchPages(
	urls: string[],
	maxCharacters?: number,
	signal?: AbortSignal,
	fetchImpl: FetchLike = fetch,
): Promise<string> {
	const result = await callTool(
		"web_fetch_exa",
		{ urls, maxCharacters: maxCharacters ?? EXA_FETCH_MAX_CHARACTERS },
		signal,
		fetchImpl,
		true,
	);
	return contentText(result) || "Exa returned no content.";
}