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

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

import {
	availableProviders,
	buildRouterPrompt,
	parseProviderChoice,
	routeSearch,
	ROUTE_CACHE_LIMIT,
	type RouterContext,
	type SearchProvider,
} from "../router.ts";

const ALL: SearchProvider[] = ["tavily", "exa", "keenable"];
const ALL_KEYS = () => "test-key";
const NO_KEYS = () => undefined;

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

test("availability: tavily needs a key, exa and keenable are keyless", () => {
	assert.deepEqual(availableProviders(NO_KEYS), ["exa", "keenable"]);
	assert.deepEqual(availableProviders(ALL_KEYS), ["tavily", "exa", "keenable"]);
});

test("router prompt lists only candidate rules plus the query and caller hints", () => {
	const prompt = buildRouterPrompt(
		{ query: "python 3.14 release notes", site: "python.org", published_after: "2025-10-01" },
		ALL,
	);
	assert.match(prompt, /- tavily: /);
	assert.match(prompt, /- exa: /);
	assert.match(prompt, /- keenable: /);
	assert.match(prompt, /Query: python 3\.14 release notes/);
	assert.match(prompt, /site filter.*python\.org|site: python\.org/);
	assert.match(prompt, /Provider:$/);

	const subset = buildRouterPrompt({ query: "hi" }, ["exa", "keenable"]);
	assert.doesNotMatch(subset, /- tavily/);
	assert.match(subset, /exactly one word: exa, keenable/);
});

test("parseProviderChoice takes the first available candidate token", () => {
	assert.equal(parseProviderChoice("tavily", ALL), "tavily");
	assert.equal(parseProviderChoice("I would pick Exa here.", ALL), "exa");
	assert.equal(parseProviderChoice("keenable!", ALL), "keenable");
	assert.equal(parseProviderChoice("google maybe", ALL), undefined);
	// Model picked an unavailable provider but mentioned a valid one later.
	assert.equal(parseProviderChoice("tavily or keenable", ["exa", "keenable"]), "keenable");
	assert.equal(parseProviderChoice("tavily", ["exa", "keenable"]), undefined);
});

test("routeSearch honors the model pick and caches identical requests", async () => {
	const calls = { count: 0 };
	const ctx = ctxAnswering("exa", calls);
	const cache = new Map<string, SearchProvider>();
	const request = { query: "official MCP spec" };

	const first = await routeSearch(request, ctx, cache, undefined, ALL_KEYS);
	assert.deepEqual(first, { provider: "exa", routedBy: "model" });

	const second = await routeSearch(request, ctx, cache, undefined, ALL_KEYS);
	assert.deepEqual(second, { provider: "exa", routedBy: "cache" });
	assert.equal(calls.count, 1);
});

test("routeSearch falls back to the first available provider on unroutable or failed model calls", async () => {
	const cache = new Map<string, SearchProvider>();
	const unparseable = await routeSearch({ query: "x" }, ctxAnswering("no idea"), cache, undefined, ALL_KEYS);
	assert.deepEqual(unparseable, { provider: "tavily", routedBy: "fallback" });

	const throwing: RouterContext = {
		model: {},
		modelRegistry: { complete: async () => { throw new Error("boom"); } },
	};
	const failed = await routeSearch({ query: "x" }, throwing, cache, undefined, ALL_KEYS);
	assert.deepEqual(failed, { provider: "tavily", routedBy: "fallback" });
});

test("routeSearch without a tavily key excludes it from candidates", async () => {
	const cache = new Map<string, SearchProvider>();
	const decision = await routeSearch({ query: "x" }, ctxAnswering("tavily"), cache, undefined, NO_KEYS);
	// "tavily" is not a candidate → unroutable → first available fallback.
	assert.deepEqual(decision, { provider: "exa", routedBy: "fallback" });
});

test("routeSearch works without a model", async () => {
	const cache = new Map<string, SearchProvider>();
	const decision = await routeSearch({ query: "x" }, {}, cache, undefined, ALL_KEYS);
	assert.deepEqual(decision, { provider: "tavily", routedBy: "fallback" });
});

test("route cache evicts the oldest entry at the limit", async () => {
	const cache = new Map<string, SearchProvider>();
	for (let i = 0; i < ROUTE_CACHE_LIMIT; i++) cache.set(`dummy-${i}`, "exa");
	const decision = await routeSearch({ query: "fresh query" }, ctxAnswering("keenable"), cache, undefined, ALL_KEYS);
	assert.equal(decision.provider, "keenable");
	assert.equal(cache.size, ROUTE_CACHE_LIMIT);
	assert.equal(cache.has("dummy-0"), false);
	assert.equal(cache.has("dummy-1"), true);
});