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 importlib
import json
import os
import sys

# cloud 包优先、租户 flat copy 回落（与 playwright_mcp 同款理由：
# 两次 from-import 会让 mypy 报 no-redef，importlib 属性绑定
# 在两条路径下都是同一组名字；cloud 导入本身的 ImportError
# 同样回落，与原来 except 语义一致）。
try:
    _mcp_backend = importlib.import_module("cloud.websearch")
except ImportError:  # flat copy in a tenant codex-home
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    _mcp_backend = importlib.import_module("websearch")
WebSearchError = _mcp_backend.WebSearchError
fetch_best = _mcp_backend.fetch_best
search = _mcp_backend.search

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' (frozen default: keyless SearXNG "
            "first, then DuckDuckGo as the free fallback, falling "
            "back on failure; the paid backends Perplexity, Exa, "
            "Tavily and Brave are NOT in the default chain even "
            "when their keys are set, and a SearXNG failure never "
            "falls back to a paid backend, so default use never "
            "incurs paid cost) or a single provider name to pin it "
            "(pinning a paid backend still calls it: that is an "
            "explicit override, not a default). Set "
            "WEBSEARCH_ENABLE_PAID=1 to restore the old default "
            "chain (Perplexity, Exa, Tavily, Brave, SearXNG, "
            "DuckDuckGo); unfreezing is explicit and may incur "
            "paid cost. "
            "scope 'code' restricts results to programming sources "
            "only (official language docs, GitHub, package "
            "registries, Stack Overflow, CVE/NVD); anything else "
            "is filtered out, never returned. Default scope "
            "'general' is unfiltered. With scope 'code' only, "
            "allowed_hosts may name extra project domains for "
            "this call (explicit allowlist extension, subdomains "
            "included, dot-boundary matched); it is not an "
            "arbitrary-site search, and it is rejected with "
            "scope 'general'. This tool does not do "
            "shopping, social, places, travel, general news, "
            "weather, finance, or browser checkout."
        ),
        "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"},
                "scope": {"type": "string", "default": "general",
                          "enum": ["general", "code"],
                          "description": "general: unfiltered web "
                                         "search. code: programming "
                                         "sources only, see tool "
                                         "description."},
                "allowed_hosts": {"type": "array",
                                  "items": {"type": "string"},
                                  "description": "Only with scope "
                                                 "'code': extra "
                                                 "domains for this "
                                                 "call only, on top "
                                                 "of the builtin "
                                                 "programming list. "
                                                 "Omit for no "
                                                 "extension; empty, "
                                                 "wildcard, or "
                                                 "non-hostname "
                                                 "values are "
                                                 "rejected, never "
                                                 "silently ignored."},
            },
            "required": ["query"],
        },
    },
    {
        "name": "web_fetch",
        "description": (
            "Fetch one URL and return clean text plus its links, "
            "truncated head+tail when over char_limit. "
            "Extraction first: Tavily extract, then Exa contents "
            "(only backends with a configured key are tried), "
            "direct fetch last. Under the egress proxy, direct "
            "fetch of a domain outside the Tinyproxy whitelist "
            "fails 403 fail-closed; that error is reported, not "
            "bypassed. Only http(s) URLs; localhost / private IPs "
            "are blocked. This is not a sandbox boundary."
        ),
        "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_best
    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":
                # allowed_hosts is passed through only when the
                # caller supplied it, so older search_fn fakes
                # without that parameter keep working.
                search_kwargs = {}
                if "allowed_hosts" in args:
                    search_kwargs["allowed_hosts"] = \
                        args.get("allowed_hosts")
                result = search_fn(
                    args.get("query"),
                    limit=args.get("limit", 5),
                    provider=args.get("provider", "auto"),
                    scope=args.get("scope", "general"),
                    **search_kwargs)
            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()