feat(mcp): register omniroute_web_search tool in MCP server (#951)

The tool was fully defined in schemas/tools.ts and backed by the working /v1/search endpoint, but server.registerTool() was never called, leaving it absent from tools/list. Changes: - server.ts: add webSearchInput import, handleWebSearch handler, and server.registerTool("omniroute_web_search") after sync_pricing block - schemas/tools.ts: align webSearchInput with /v1/search contract -- query max reduced 1000->500, provider narrowed to explicit enum - essentialTools.test.ts: rewrite web_search stubs to dispatch through a real McpServer+InMemoryTransport+Client, providing actual handler coverage (POST method, body fields, error paths, tools/list check) - vitest.mcp.config.ts: dedicated vitest config for MCP server tests; update test:vitest script to use it Note: omniRouteFetch hardcodes AbortSignal.timeout(10000) unconditionally (server.ts line 126), so caller signals are silently discarded -- the effective search timeout is 10s. Follow-up PR can add timeoutMs param. cacheStatsTool and cacheFlushTool are also unregistered; out of scope. 🤖 Generated with Claude Sonnet 4.6 via Claude Code (https://claude.com/claude-code) + Compound Engineering v2.58.1 Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-03 14:17:27 +01:00
parent db0af86085
commit a0cfae214d
5 changed files with 221 additions and 71 deletions
@@ -1,11 +1,16 @@
 /**
 * Unit tests for MCP Essential Tools (Phase 1)
 *
- * Tests all 8 essential tool handlers via the tool handler functions.
+ * Tests all 9 essential tool handlers via the tool handler functions.
+ * The omniroute_web_search tests use InMemoryTransport + Client to exercise
+ * the actual registered handler (not mockFetch directly).
 */

-import { describe, it, expect, vi, beforeEach } from "vitest";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
 import { MCP_ESSENTIAL_TOOLS } from "../schemas/tools";
+import { createMcpServer } from "../server";

 // Mock fetch globally
 const mockFetch = vi.fn();
@@ -17,7 +22,7 @@ describe("MCP Essential Tools", () => {
  });

  describe("Tool schema validation", () => {
-    it("should have exactly 9 essential tools", () => {
+    it("should have exactly 9 essential tools (includes web_search)", () => {
      const schemas = MCP_ESSENTIAL_TOOLS;
      expect(schemas).toHaveLength(9);
    });
@@ -136,77 +141,157 @@ describe("MCP Essential Tools", () => {
      expect(data).toHaveProperty("requestCount");
    });
  });
+});

-  describe("web_search handler", () => {
-    it("should return search results when API is available", async () => {
-      mockFetch.mockResolvedValueOnce({
-        ok: true,
-        json: async () => ({
-          id: "search-123",
-          provider: "serper",
-          query: "typescript best practices",
-          results: [
-            {
-              title: "TypeScript Best Practices 2024",
-              url: "https://example.com/ts-best",
-              display_url: "https://example.com/ts-best",
-              snippet: "Best practices for TypeScript development...",
-              position: 1,
-            },
-            {
-              title: "Advanced TypeScript Patterns",
-              url: "https://example.com/ts-advanced",
-              snippet: "Advanced patterns and techniques...",
-              position: 2,
-            },
-          ],
-          cached: false,
-          usage: { queries_used: 1, search_cost_usd: 0.002 },
-        }),
-      });
+// ── omniroute_web_search: handler dispatch tests ──────────────────────────────
+// These tests use InMemoryTransport + Client to exercise the actual registered
+// handler (not mockFetch directly), ensuring real handler coverage.

-      const response = await mockFetch(
-        "http://localhost:20128/v1/search?query=typescript%20best%20practices&max_results=5"
-      );
-      const data = await response.json();
-      expect(data.results).toHaveLength(2);
-      expect(data.results[0].title).toBe("TypeScript Best Practices 2024");
-      expect(data.provider).toBe("serper");
+vi.mock("../audit.ts", () => ({
+  logToolCall: vi.fn().mockResolvedValue(undefined),
+}));
+
+describe("omniroute_web_search handler (via MCP dispatch)", () => {
+  let client: Client;
+
+  beforeEach(async () => {
+    mockFetch.mockReset();
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+    const server = createMcpServer();
+    await server.connect(serverTransport);
+    client = new Client({ name: "test-client", version: "1.0.0" });
+    await client.connect(clientTransport);
+  });
+
+  afterEach(async () => {
+    await client.close();
+  });
+
+  it("should appear in tools/list after registration", async () => {
+    const { tools } = await client.listTools();
+    const webSearch = tools.find((t) => t.name === "omniroute_web_search");
+    expect(webSearch).toBeDefined();
+    expect(webSearch?.description).toContain("web search");
+  });
+
+  it("should return search results on success", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({
+        id: "search-123",
+        provider: "serper-search",
+        query: "typescript best practices",
+        results: [
+          {
+            title: "TypeScript Best Practices 2024",
+            url: "https://example.com/ts-best",
+            display_url: "https://example.com/ts-best",
+            snippet: "Best practices for TypeScript development...",
+            position: 1,
+          },
+        ],
+        cached: false,
+        usage: { queries_used: 1, search_cost_usd: 0.002 },
+      }),
    });

-    it("should handle API failure gracefully", async () => {
-      mockFetch.mockRejectedValueOnce(new Error("Search service unavailable"));
-
-      await expect(mockFetch("http://localhost:20128/v1/search?query=test")).rejects.toThrow(
-        "Search service unavailable"
-      );
+    const result = await client.callTool({
+      name: "omniroute_web_search",
+      arguments: { query: "typescript best practices" },
    });

-    it("should pass correct parameters to /v1/search", async () => {
-      mockFetch.mockResolvedValueOnce({
-        ok: true,
-        json: async () => ({
-          id: "search-456",
-          provider: "brave",
-          query: "react hooks tutorial",
-          results: [],
-          cached: false,
-          usage: { queries_used: 1, search_cost_usd: 0.003 },
-        }),
-      });
+    expect(result.isError).toBeFalsy();
+    const content = result.content[0] as { type: string; text: string };
+    const data = JSON.parse(content.text);
+    expect(data.results).toHaveLength(1);
+    expect(data.results[0].title).toBe("TypeScript Best Practices 2024");
+    expect(data.provider).toBe("serper-search");
+  });

-      const query = "react hooks tutorial";
-      const response = await mockFetch(
-        `http://localhost:20128/v1/search?query=${encodeURIComponent(query)}&max_results=10&search_type=news&provider=brave`
-      );
-      const data = await response.json();
-
-      expect(mockFetch).toHaveBeenCalledWith(expect.stringContaining("/v1/search"));
-      expect(mockFetch).toHaveBeenCalledWith(
-        expect.stringContaining("query=react%20hooks%20tutorial")
-      );
-      expect(mockFetch).toHaveBeenCalledWith(expect.stringContaining("max_results=10"));
-      expect(data.provider).toBe("brave");
+  it("should POST to /v1/search with correct body fields", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({
+        id: "s1",
+        provider: "brave-search",
+        query: "react hooks tutorial",
+        results: [],
+        cached: false,
+        usage: { queries_used: 1, search_cost_usd: 0.003 },
+      }),
    });
+
+    await client.callTool({
+      name: "omniroute_web_search",
+      arguments: {
+        query: "react hooks tutorial",
+        max_results: 10,
+        search_type: "news",
+        provider: "brave-search",
+      },
+    });
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      expect.stringContaining("/v1/search"),
+      expect.objectContaining({ method: "POST" })
+    );
+    const [, options] = mockFetch.mock.calls[0];
+    const body = JSON.parse(options.body as string);
+    expect(body.query).toBe("react hooks tutorial");
+    expect(body.max_results).toBe(10);
+    expect(body.search_type).toBe("news");
+    expect(body.provider).toBe("brave-search");
+  });
+
+  it("should omit provider from POST body when not specified", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({
+        id: "s2",
+        provider: "serper-search",
+        query: "test query",
+        results: [],
+        cached: false,
+        usage: { queries_used: 1, search_cost_usd: 0 },
+      }),
+    });
+
+    await client.callTool({
+      name: "omniroute_web_search",
+      arguments: { query: "test query" },
+    });
+
+    const [, options] = mockFetch.mock.calls[0];
+    const body = JSON.parse(options.body as string);
+    expect(body).not.toHaveProperty("provider");
+  });
+
+  it("should return isError on backend non-OK response", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: false,
+      status: 500,
+      text: async () => "Internal server error",
+    });
+
+    const result = await client.callTool({
+      name: "omniroute_web_search",
+      arguments: { query: "test" },
+    });
+
+    expect(result.isError).toBe(true);
+    const content = result.content[0] as { type: string; text: string };
+    expect(content.text).toContain("Error");
+  });
+
+  it("should return isError on fetch abort/timeout", async () => {
+    mockFetch.mockRejectedValueOnce(new DOMException("signal timed out", "TimeoutError"));
+
+    const result = await client.callTool({
+      name: "omniroute_web_search",
+      arguments: { query: "test" },
+    });
+
+    expect(result.isError).toBe(true);
  });
 });
@@ -405,7 +405,7 @@ export const webSearchInput = z.object({
  query: z
    .string()
    .min(1, "Query is required")
-    .max(1000, "Query must be 1000 characters or fewer")
+    .max(500, "Query must be 500 characters or fewer")
    .describe("The search query string"),
  max_results: z
    .number()
@@ -416,9 +416,9 @@ export const webSearchInput = z.object({
    .describe("Maximum number of search results to return"),
  search_type: z.enum(["web", "news"]).default("web").describe("Type of search to perform"),
  provider: z
-    .string()
+    .enum(["serper-search", "brave-search", "perplexity-search", "exa-search", "tavily-search"])
    .optional()
-    .describe("Specific search provider to use (serper, brave, perplexity, exa, tavily)"),
+    .describe("Specific search provider to use"),
 });

 export const webSearchOutput = z.object({
@@ -23,6 +23,7 @@ import {
  routeRequestInput,
  costReportInput,
  listModelsCatalogInput,
+  webSearchInput,
  simulateRouteInput,
  setBudgetGuardInput,
  setRoutingStrategyInput,
@@ -500,6 +501,39 @@ async function handleListModelsCatalog(args: { provider?: string; capability?: s
  }
 }

+async function handleWebSearch(args: {
+  query: string;
+  max_results?: number;
+  search_type?: "web" | "news";
+  provider?:
+    | "serper-search"
+    | "brave-search"
+    | "perplexity-search"
+    | "exa-search"
+    | "tavily-search";
+}) {
+  const start = Date.now();
+  try {
+    const body: Record<string, unknown> = {
+      query: args.query,
+      max_results: args.max_results ?? 5,
+      search_type: args.search_type ?? "web",
+    };
+    if (args.provider) body.provider = args.provider;
+
+    const result = await omniRouteFetch("/v1/search", {
+      method: "POST",
+      body: JSON.stringify(body),
+    });
+    await logToolCall("omniroute_web_search", args, result, Date.now() - start, true);
+    return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }] };
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    await logToolCall("omniroute_web_search", args, null, Date.now() - start, false, msg);
+    return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
+  }
+}
+
 // ============ MCP Server Setup ============

 /**
@@ -725,6 +759,18 @@ export function createMcpServer(): McpServer {
    )
  );

+  server.registerTool(
+    "omniroute_web_search",
+    {
+      description:
+        "Performs a web search using OmniRoute's search gateway. Supports multiple providers (Serper, Brave, Perplexity, Exa, Tavily) with automatic failover. Returns search results with titles, URLs, snippets, and position data.",
+      inputSchema: webSearchInput,
+    },
+    withScopeEnforcement("omniroute_web_search", (args) =>
+      handleWebSearch(webSearchInput.parse(args))
+    )
+  );
+
  // ── Memory Tools ──────────────────────────────
  Object.values(memoryTools).forEach((toolDef) => {
    server.registerTool(
@@ -70,7 +70,7 @@
    "test:integration": "node --import tsx/esm --test tests/integration/*.test.mjs",
    "test:e2e": "node scripts/run-playwright-tests.mjs test tests/e2e/*.spec.ts",
    "test:protocols:e2e": "node scripts/run-protocol-clients-tests.mjs",
-    "test:vitest": "vitest run open-sse/mcp-server/__tests__/*.test.ts open-sse/services/autoCombo/__tests__/*.test.ts",
+    "test:vitest": "vitest run --config vitest.mcp.config.ts",
    "test:ecosystem": "node scripts/run-ecosystem-tests.mjs",
    "test:coverage": "c8 --exclude=tests/** --exclude=**/*.test.* --reporter=text-summary --reporter=html --reporter=json-summary --reporter=lcov --check-coverage --statements 55 --lines 55 --functions 55 --branches 60 node --import tsx/esm --test tests/unit/*.test.mjs",
    "test:coverage:legacy": "c8 --exclude=open-sse --check-coverage --lines 50 --functions 50 --branches 50 node --import tsx/esm --test tests/unit/*.test.mjs",
@@ -0,0 +1,19 @@
+import { defineConfig } from "vitest/config";
+import path from "path";
+
+export default defineConfig({
+  test: {
+    environment: "node",
+    globals: true,
+    include: [
+      "open-sse/mcp-server/__tests__/**/*.test.ts",
+      "open-sse/services/autoCombo/__tests__/**/*.test.ts",
+    ],
+    exclude: ["**/node_modules/**", "**/.git/**"],
+  },
+  resolve: {
+    alias: {
+      "@": path.resolve(__dirname, "./src"),
+    },
+  },
+});