codex-py/codex/mcp_server.py

"""MCP server — read-only KB tools exposed via FastMCP (F-14).

This module is a *thin facade* over existing codex domain modules.
No new domain logic lives here — each tool delegates to the relevant
module and reshapes the result into a JSON-serialisable dict/list.

Graceful degradation contract
------------------------------
Every tool that depends on an optional feature (wiki, leads/synthesis,
provenance-verify) catches all exceptions and returns
``{"error": "feature not available"}`` instead of crashing.

Transport
---------
* **stdio** (default):
  ``python -m codex.mcp_server``  or  ``codex-mcp``

* **HTTP** (config-gated):
  Set ``MCP_TRANSPORT=http`` *and* ``MCP_AUTH_TOKEN=<token>`` in your
  environment (or ``.env``).  The server refuses to start without a
  token — no unauthenticated HTTP endpoints are ever exposed.
"""

from __future__ import annotations

import logging

from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings

logger = logging.getLogger(__name__)

# DNS rebinding protection disabled: the server is guarded by Bearer token auth
# (mcp_server.py _TokenAuth middleware) and UFW firewall (trusted LAN IPs only).
mcp = FastMCP(
    "codex",
    transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)


# ---------------------------------------------------------------------------
# Tool: search
# ---------------------------------------------------------------------------


@mcp.tool()
def search(query: str, limit: int = 10) -> list[dict[str, object]]:
    """Chunk-level hybrid search (dense BGE-M3 + FTS) over the KB.

    Returns one dict per hit:
    ``{bibkey, paper_id, locator, score, snippet}``.
    """
    try:
        from codex.db import get_conn
        from codex.embed import get_embedder

        embedder = get_embedder()
        dense_vec = embedder.encode_dense([query])[0].tolist()

        sql = """
            SELECT
                c.id,
                c.paper_id,
                c.ord,
                c.content,
                p.bibkey,
                c.embedding <-> %(emb)s::vector AS dist
            FROM chunks c
            JOIN papers p ON p.id = c.paper_id
            WHERE c.embedding IS NOT NULL
              AND p.bibkey IS NOT NULL
            ORDER BY c.embedding <-> %(emb)s::vector
            LIMIT %(limit)s
        """
        with get_conn() as conn:
            rows = conn.execute(sql, {"emb": dense_vec, "limit": limit}).fetchall()

        results: list[dict[str, object]] = []
        for row in rows:
            results.append(
                {
                    "bibkey": row["bibkey"],
                    "paper_id": row["paper_id"],
                    "locator": f"chunk {row['ord']}",
                    "score": float(1.0 - row["dist"]),  # convert distance → similarity
                    "snippet": row["content"][:300],
                }
            )
        return results
    except Exception as exc:  # noqa: BLE001
        logger.warning("search tool error: %s", exc)
        return [{"error": str(exc)}]


# ---------------------------------------------------------------------------
# Tool: ask
# ---------------------------------------------------------------------------


@mcp.tool()
def ask(question: str) -> dict[str, object]:
    """RAG answer over the KB.

    Returns ``{"answer": str}`` when implemented,
    or ``{"error": "not available"}`` until a RAG layer is added.
    """
    # RAG not yet implemented (CLI parity with codex ask)
    return {"error": "not available"}


# ---------------------------------------------------------------------------
# Tool: wiki_read
# ---------------------------------------------------------------------------


@mcp.tool()
def wiki_read(concept_slug: str) -> dict[str, object]:
    """Compiled concept page (F-12) as ``{markdown, sources}``.

    Gracefully returns ``{"error": "feature not available"}`` when the wiki
    directory or the requested page does not exist.
    """
    try:
        import json
        from pathlib import Path

        from codex.config import get_settings

        settings = get_settings()
        wiki_dir = Path(settings.wiki_dir)
        page_path = wiki_dir / f"{concept_slug}.md"

        if not page_path.exists():
            return {"error": "feature not available"}

        markdown = page_path.read_text(encoding="utf-8")

        # Collect cited bibkeys from the compile state
        state_path = wiki_dir / ".compile-state.json"
        sources: list[str] = []
        if state_path.exists():
            try:
                state: dict[str, str] = json.loads(state_path.read_text(encoding="utf-8"))
                if concept_slug in state:
                    sources = [concept_slug]
            except (json.JSONDecodeError, OSError):
                pass

        return {"markdown": markdown, "sources": sources}
    except Exception as exc:  # noqa: BLE001
        logger.warning("wiki_read tool error: %s", exc)
        return {"error": "feature not available"}


# ---------------------------------------------------------------------------
# Tool: wiki_list
# ---------------------------------------------------------------------------


@mcp.tool()
def wiki_list() -> list[dict[str, object]]:
    """All compiled concept pages with freshness metadata.

    Returns a list of ``{slug, mtime, hash_prefix}`` dicts.
    Gracefully returns ``[]`` when the wiki directory does not exist.
    """
    try:
        import json
        from datetime import UTC, datetime
        from pathlib import Path

        from codex.config import get_settings

        settings = get_settings()
        wiki_dir = Path(settings.wiki_dir)

        if not wiki_dir.exists():
            return []

        state_path = wiki_dir / ".compile-state.json"
        state: dict[str, str] = {}
        if state_path.exists():
            try:
                state = json.loads(state_path.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                state = {}

        pages = sorted(wiki_dir.glob("*.md"))
        result: list[dict[str, object]] = []
        for page in pages:
            if page.name in ("index.md", "log.md"):
                continue
            slug = page.stem
            hash_prefix = state.get(slug, "")[:8]
            mtime = datetime.fromtimestamp(page.stat().st_mtime, tz=UTC).isoformat()
            result.append({"slug": slug, "mtime": mtime, "hash_prefix": hash_prefix})

        return result
    except Exception as exc:  # noqa: BLE001
        logger.warning("wiki_list tool error: %s", exc)
        return []


# ---------------------------------------------------------------------------
# Tool: discover_leads
# ---------------------------------------------------------------------------


@mcp.tool()
def discover_leads(limit: int = 20) -> list[dict[str, object]]:
    """Dangling-citation discovery leads — papers cited but not yet ingested.

    Returns ``[{cited_id, pull}]`` ordered by citation count descending.
    """
    try:
        from codex.discover import discovery_leads as _leads

        raw = _leads(limit=limit)
        return [{"cited_id": row["cited_id"], "pull": int(row["pull"])} for row in raw]
    except Exception as exc:  # noqa: BLE001
        logger.warning("discover_leads tool error: %s", exc)
        return [{"error": str(exc)}]


# ---------------------------------------------------------------------------
# Tool: provenance_verify
# ---------------------------------------------------------------------------


@mcp.tool()
def provenance_verify(lib_path: str) -> list[dict[str, object]]:
    """Scan C++ sources for @cite tags and return code-to-paper links.

    Read-only: only scans, never writes.
    Returns ``[{file, line, symbol, bibkey}]`` or ``[{"error": ...}]``.
    """
    try:
        from codex.provenance import scan_cite_tags

        hits = scan_cite_tags(lib_path)
        return [dict(h) for h in hits]
    except Exception as exc:  # noqa: BLE001
        logger.warning("provenance_verify tool error: %s", exc)
        return [{"error": "feature not available"}]


# ---------------------------------------------------------------------------
# Tool: synthesis_browse
# ---------------------------------------------------------------------------


@mcp.tool()
def synthesis_browse(kind: str = "all") -> list[dict[str, object]]:
    """Leads / conjectures from the ``leads/`` directory (F-13).

    Reads ``.json`` files from ``leads/`` and filters by ``kind``
    (``"all"`` returns everything).  Gracefully returns
    ``[{"error": "feature not available"}]`` when the directory is absent.
    """
    try:
        import json
        from pathlib import Path

        from codex.config import get_settings

        settings = get_settings()
        leads_dir = Path(settings.leads_dir)
        if not leads_dir.exists():
            return [{"error": "feature not available"}]

        items: list[dict[str, object]] = []
        for f in sorted(leads_dir.rglob("*.json")):
            try:
                data: dict[str, object] = json.loads(f.read_text(encoding="utf-8"))
                item_kind = str(data.get("kind", "unknown"))
                if kind == "all" or item_kind == kind:
                    items.append(data)
            except (json.JSONDecodeError, OSError):
                continue

        return items
    except Exception as exc:  # noqa: BLE001
        logger.warning("synthesis_browse tool error: %s", exc)
        return [{"error": "feature not available"}]


# ---------------------------------------------------------------------------
# Server entry-point helpers
# ---------------------------------------------------------------------------


def _require_http_token() -> str:
    """Return the configured auth token or raise :exc:`RuntimeError`.

    Called before starting HTTP transport to enforce the security invariant:
    no unauthenticated HTTP endpoints are ever exposed.
    """
    from codex.config import get_settings

    settings = get_settings()
    if not settings.mcp_auth_token:
        raise RuntimeError(
            "HTTP transport requires MCP_AUTH_TOKEN to be set. "
            "Set the environment variable or add it to .env."
        )
    return settings.mcp_auth_token


# ---------------------------------------------------------------------------
# Server entry-point
# ---------------------------------------------------------------------------


def main() -> None:
    """Start the MCP server (stdio or HTTP depending on config).

    HTTP transport requires ``MCP_AUTH_TOKEN`` to be set; the server will
    raise :exc:`RuntimeError` and refuse to start if the token is missing.
    """
    from codex.config import get_settings

    settings = get_settings()
    transport = settings.mcp_transport.lower()

    if transport == "http":
        import uvicorn
        from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
        from starlette.requests import Request
        from starlette.responses import Response

        token = _require_http_token()

        class _TokenAuth(BaseHTTPMiddleware):
            """Reject requests that do not carry a valid Bearer token."""

            async def dispatch(
                self, request: Request, call_next: RequestResponseEndpoint
            ) -> Response:
                auth = request.headers.get("Authorization", "")
                if auth != f"Bearer {token}":
                    return Response("Unauthorized", status_code=401)
                return await call_next(request)

        base_app = mcp.streamable_http_app()
        base_app.add_middleware(_TokenAuth)

        uvicorn.run(
            base_app,
            host=settings.mcp_host,
            port=settings.mcp_port,
            log_level="info",
        )
    else:
        # Default: stdio transport
        mcp.run(transport="stdio")


if __name__ == "__main__":
    main()