codex-app-server/cloud/websearch_mcp.py

"""Stdio MCP server exposing web_search / web_fetch to codex.

Protocol: JSON-RPC 2.0 over stdio, one JSON object per line.
Logs go to stderr only -- anything on stdout would corrupt the
JSON-RPC stream.

This file is deliberately thin: all provider logic lives in
``websearch.py``. It can run two ways:
- inside the repo: ``python3 -m cloud.websearch_mcp`` (imports
  ``cloud.websearch``);
- copied flat into a tenant codex-home next to ``websearch.py``
  (imports plain ``websearch`` from its own directory).
The ``ensure_codex_home_websearch`` helper in websearch.py does the
copying; it never writes API keys into config.toml -- keys are read
from the MCP server process environment by websearch.py.
"""

import json
import os
import sys

try:
    from cloud.websearch import WebSearchError, fetch_url, search
except ImportError:  # flat copy in a tenant codex-home
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from websearch import WebSearchError, fetch_url, search  # noqa: E402

SERVER_NAME = "websearch"
SERVER_VERSION = "0.1.0"

_TOOLS = [
    {
        "name": "web_search",
        "description": (
            "Search the web and return ranked results with title, URL "
            "and snippet, plus a synthesised answer when the provider "
            "supplies one. Use web_fetch to read specific pages. "
            "provider may be 'auto' (walk Perplexity, Exa, Tavily, "
            "Brave, then keyless SearXNG / DuckDuckGo, falling back "
            "on failure) or a single provider name to pin it."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string",
                          "description": "Search query"},
                "limit": {"type": "integer", "default": 5,
                          "description": "Max results, 1-10"},
                "provider": {"type": "string", "default": "auto",
                             "description": "auto or a provider name"},
            },
            "required": ["query"],
        },
    },
    {
        "name": "web_fetch",
        "description": (
            "Fetch one URL and return clean text plus its links, "
            "truncated head+tail when over char_limit. "
            "Only http(s) URLs; localhost / private IPs are blocked."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "url": {"type": "string",
                        "description": "URL to fetch"},
                "char_limit": {"type": "integer", "default": 15000,
                               "description": "Max characters, "
                                              "clamped 2000-500000"},
            },
            "required": ["url"],
        },
    },
]


def _err(code, message):
    return {"code": code, "message": message}


def handle_request(request, search_fn=None, fetch_fn=None):
    """Handle one JSON-RPC request dict. Returns a response dict or
    None for notifications (no id). Pure function for testability:
    the stdio loop passes the real search/fetch, tests pass fakes."""
    search_fn = search_fn or search
    fetch_fn = fetch_fn or fetch_url
    if not isinstance(request, dict):
        return {"jsonrpc": "2.0", "id": None,
                "error": _err(-32700, "request is not an object")}
    method = request.get("method")
    req_id = request.get("id")
    is_notification = "id" not in request

    if method == "initialize":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {
                    "protocolVersion": request.get("params", {}).get(
                        "protocolVersion", "2024-11-05"),
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": SERVER_NAME,
                                   "version": SERVER_VERSION}}}
    if method in ("notifications/initialized", "notifications/cancelled"):
        return None
    if method == "ping":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}
    if method == "tools/list":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"tools": _TOOLS}}
    if method == "tools/call":
        if is_notification:
            return None
        params = request.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        try:
            if name == "web_search":
                result = search_fn(
                    args.get("query"),
                    limit=args.get("limit", 5),
                    provider=args.get("provider", "auto"))
            elif name == "web_fetch":
                result = fetch_fn(
                    args.get("url"),
                    char_limit=args.get("char_limit", 15000))
            else:
                return {"jsonrpc": "2.0", "id": req_id,
                        "error": _err(-32602, "unknown tool: %r" % name)}
        except Exception as exc:  # noqa: BLE001 - tool errors are data
            # WebSearchError messages are already key-free by
            # construction in websearch.py; other exceptions are
            # rendered as type + message, never with env or headers.
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text",
                                            "text": "error: %s" % exc}],
                               "isError": True}}
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"content": [{"type": "text",
                                        "text": json.dumps(
                                            result, ensure_ascii=False)}],
                           "isError": False}}
    if is_notification:
        return None
    return {"jsonrpc": "2.0", "id": req_id,
            "error": _err(-32601, "method not found: %r" % method)}


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
        except json.JSONDecodeError as exc:
            sys.stdout.write(json.dumps(
                {"jsonrpc": "2.0", "id": None,
                 "error": _err(-32700, "parse error: %s" % exc)}) + "\n")
            sys.stdout.flush()
            continue
        try:
            response = handle_request(request)
        except Exception as exc:  # noqa: BLE001 - server must not die
            print("websearch_mcp: handler error: %s" % exc, file=sys.stderr)
            if "id" in request:
                response = {"jsonrpc": "2.0", "id": request["id"],
                            "error": _err(-32603, "internal error")}
            else:
                response = None
        if response is not None:
            sys.stdout.write(json.dumps(response) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()