polysearch/src/mcp.js

import { childLogger } from "./utils/logger.js";

const log = childLogger({ component: "mcp" });

const TOOLS = [
  {
    name: "web_search",
    description: "Search the web via DuckDuckGo. Returns organic results with title, URL, snippet, and domain.",
    inputSchema: {
      type: "object",
      properties: {
        query: { type: "string", description: "Search query" },
        limit: { type: "number", description: "Max results (1-50)", default: 10 }
      },
      required: ["query"]
    }
  },
  {
    name: "image_search",
    description: "Search for images via DuckDuckGo. Returns results with title, image_url, source_url, and dimensions.",
    inputSchema: {
      type: "object",
      properties: {
        query: { type: "string", description: "Image search query" },
        limit: { type: "number", description: "Max results (1-50)", default: 10 }
      },
      required: ["query"]
    }
  }
];

export function createMCPHandler(runner) {
  return async function handleMCP(body) {
    if (!body || body.jsonrpc !== "2.0") {
      return { response: mcpError(null, -32600, "Invalid Request: must be JSON-RPC 2.0"), isNotification: false };
    }

    const { id, method, params } = body;
    const isNotification = id === undefined || id === null;

    switch (method) {
      case "tools/list":
        return {
          response: { jsonrpc: "2.0", id, result: { tools: TOOLS } },
          isNotification: false
        };

      case "tools/call":
        return { response: await handleToolCall(id, params, runner), isNotification: false };

      case "initialize":
        return {
          response: {
            jsonrpc: "2.0",
            id,
            result: {
              protocolVersion: "2024-11-05",
              capabilities: { tools: {} },
              serverInfo: { name: "polysearch", version: "2.0.0" }
            }
          },
          isNotification: false
        };

      case "notifications/initialized":
      case "notifications/cancelled":
        log.debug({ method }, "notification received, no response");
        return { response: null, isNotification: true };

      default:
        if (isNotification) {
          log.debug({ method }, "unknown notification, silently ignored");
          return { response: null, isNotification: true };
        }
        return { response: mcpError(id, -32601, `Method not found: ${method}`), isNotification: false };
    }
  };
}

async function handleToolCall(id, params, runner) {
  if (!params || !params.name || !params.arguments) {
    return mcpError(id, -32602, "Invalid params: name and arguments required");
  }

  const { name, arguments: args } = params;

  switch (name) {
    case "web_search":
    case "image_search": {
      const query = args.query;
      if (!query || typeof query !== "string" || query.trim().length === 0) {
        return mcpError(id, -32602, "Missing required argument: query");
      }

      const type = name === "image_search" ? "image" : "web";
      const limit = Math.min(Math.max(parseInt(args.limit, 10) || 10, 1), 50);

      log.info({ query, type, limit }, "MCP tool call");

      try {
        const data = await runner.run({ query, type, limit });
        const content = formatMCPResult(data, type);
        return {
          jsonrpc: "2.0",
          id,
          result: { content, isError: false }
        };
      } catch (err) {
        log.error({ error: err.message, name }, "MCP tool failed");
        return {
          jsonrpc: "2.0",
          id,
          result: {
            content: [{ type: "text", text: JSON.stringify({ error: err.message }) }],
            isError: true
          }
        };
      }
    }

    default:
      return mcpError(id, -32601, `Unknown tool: ${name}`);
  }
}

function formatMCPResult(data, type) {
  const section = data[type];
  if (!section || !section.results || section.results.length === 0) {
    return [{ type: "text", text: JSON.stringify({ query: data.query, total: 0, results: [] }) }];
  }

  const results = section.results.map(r => {
    const base = { title: r.title, domain: r.domain || "unknown" };
    if (type === "image") {
      return { ...base, image_url: r.image_url, source_url: r.source_url, width: r.width, height: r.height };
    }
    return { ...base, url: r.url, snippet: r.snippet };
  });

  return [{
    type: "text",
    text: JSON.stringify({
      query: data.query,
      total: section.total,
      engine: section.engine,
      execution_time_ms: data.executionTimeMs,
      results
    }, null, 2)
  }];
}

function mcpError(id, code, message) {
  return {
    jsonrpc: "2.0",
    id,
    error: { code, message }
  };
}

export { mcpError };