my-pi/pi-web-search/test/dispatch.test.ts

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

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

import webSearchExtension, {
	createFetchRunners,
	createProviderRunners,
	dispatchFetch,
	dispatchSearch,
	normalizeUrls,
	toExaParams,
	toKeenableParams,
	toTavilyParams,
	type WebSearchParams,
} from "../index.ts";

const PARAMS: WebSearchParams = { query: "release notes" };

function ctxAnswering(answer: string, calls: { count: number } = { count: 0 }) {
	return {
		model: { provider: "test", id: "main" },
		modelRegistry: {
			complete: async () => {
				calls.count += 1;
				return { content: [{ type: "text", text: answer }] };
			},
		},
	};
}

function withEnv<T>(key: string, value: string | undefined, fn: () => Promise<T> | T): Promise<T> | T {
	const previous = process.env[key];
	if (value === undefined) delete process.env[key];
	else process.env[key] = value;
	const restore = () => {
		if (previous === undefined) delete process.env[key];
		else process.env[key] = previous;
	};
	try {
		const out = fn();
		if (out instanceof Promise) return out.finally(restore);
		restore();
		return out;
	} catch (error) {
		restore();
		throw error;
	}
}

test("extension registers only web_search and web_fetch, with no provider in the schema", () => {
	const tools = new Map<string, { parameters?: { properties?: Record<string, unknown> }; description?: string }>();
	const pi = {
		registerTool(definition: { name: string }) {
			tools.set(definition.name, definition as never);
		},
		on() {},
	} as unknown as ExtensionAPI;

	webSearchExtension(pi);
	assert.deepEqual([...tools.keys()].sort(), ["web_fetch", "web_search"]);
	const schemaText = JSON.stringify([...tools.values()].map((t) => t.parameters));
	assert.doesNotMatch(schemaText, /tavily|exa|keenable|"provider"/iu);
});

test("dispatchSearch runs the model-picked provider first and records details", async () => {
	await withEnv("TAVILY_API_KEY", "k", async () => {
		const ran: string[] = [];
		const runners = createProviderRunners({
			exa: async () => {
				ran.push("exa");
				return "exa results";
			},
		});
		const result = await dispatchSearch(PARAMS, ctxAnswering("exa"), new Map(), undefined, runners);
		assert.deepEqual(ran, ["exa"]);
		assert.equal(result.content[0]!.type === "text" && result.content[0]!.text, "exa results");
		assert.deepEqual(result.details, { provider: "exa", routedBy: "model" });
	});
});

test("dispatchSearch falls back across remaining providers when the pick fails", async () => {
	await withEnv("TAVILY_API_KEY", "k", async () => {
		const ran: string[] = [];
		const runners = createProviderRunners({
			tavily: async () => {
				ran.push("tavily");
				throw new Error("tavily down");
			},
			exa: async () => {
				ran.push("exa");
				return "exa results";
			},
		});
		const result = await dispatchSearch(PARAMS, ctxAnswering("tavily"), new Map(), undefined, runners);
		assert.deepEqual(ran, ["tavily", "exa"]);
		assert.equal((result.details as Record<string, unknown>).provider, "exa");
		assert.equal((result.details as Record<string, unknown>).routedBy, "provider-fallback");
	});
});

test("dispatchSearch reports an error result when every provider fails", async () => {
	await withEnv("TAVILY_API_KEY", "k", async () => {
		const runners = createProviderRunners({
			tavily: async () => { throw new Error("t"); },
			exa: async () => { throw new Error("e"); },
			keenable: async () => { throw new Error("k"); },
		});
		const result = await dispatchSearch(PARAMS, ctxAnswering("exa"), new Map(), undefined, runners);
		assert.equal((result.details as Record<string, unknown>).error, true);
		assert.match(result.content[0]!.type === "text" ? result.content[0]!.text : "", /failed on every available provider/);
	});
});

test("normalizeUrls accepts a string or array, trims, and deduplicates", () => {
	assert.deepEqual(normalizeUrls(" https://a.example "), ["https://a.example"]);
	assert.deepEqual(normalizeUrls(["https://a", " https://a ", "https://b"]), ["https://a", "https://b"]);
});

test("dispatchFetch prefers exa and degrades to keenable without a model call", async () => {
	await withEnv("TAVILY_API_KEY", undefined, async () => {
		const fetchers = createFetchRunners({
			exa: async (urls) => {
				throw new Error(`exa down for ${urls.length}`);
			},
			keenable: async (urls) => `keenable pages: ${urls.join(",")}`,
		});
		const result = await dispatchFetch({ urls: ["https://a.example", "https://b.example"] }, undefined, fetchers);
		assert.equal((result.details as Record<string, unknown>).provider, "keenable");
		assert.match(result.content[0]!.type === "text" ? result.content[0]!.text : "", /keenable pages: https:\/\/a\.example,https:\/\/b\.example/);
	});
});

test("dispatchFetch validates input and caps the URL count", async () => {
	const empty = await dispatchFetch({ urls: [] }, undefined);
	assert.equal((empty.details as Record<string, unknown>).error, true);

	const tooMany = await dispatchFetch({ urls: Array.from({ length: 11 }, (_, i) => `https://e${i}.example`) }, undefined);
	assert.equal((tooMany.details as Record<string, unknown>).error, true);
	assert.match(tooMany.content[0]!.type === "text" ? tooMany.content[0]!.text : "", /at most 10/);
});

test("unified params map onto each provider's native arguments", () => {
	const params: WebSearchParams = {
		query: "q",
		max_results: 4,
		site: "a.dev",
		published_after: "2025-01-01",
		published_before: "2025-06-01",
		depth: "advanced",
	};
	assert.deepEqual(toTavilyParams(params), {
		query: "q",
		max_results: 4,
		search_depth: "advanced",
		include_domains: ["a.dev"],
		start_date: "2025-01-01",
		end_date: "2025-06-01",
	});
	assert.deepEqual(toExaParams(params), {
		query: "q",
		numResults: 4,
		includeDomains: ["a.dev"],
		startPublishedDate: "2025-01-01",
		endPublishedDate: "2025-06-01",
	});
	assert.deepEqual(toKeenableParams(params), {
		query: "q",
		site: "a.dev",
		published_after: "2025-01-01",
		published_before: "2025-06-01",
	});

	const defaults = toTavilyParams({ query: "q" });
	assert.equal(defaults.max_results, 8);
	assert.equal(defaults.search_depth, "basic");
	assert.equal(defaults.include_domains, undefined);
});