database_url, migration_database_url, mcp_auth_token and the mathpix app id/key are now pydantic SecretStr, so they render as '**********' in any repr/log/traceback instead of leaking the plaintext credential. Consumers (db.get_conn, cli.migrate, mcp._require_http_token, mathpix.extract_formulas) call .get_secret_value() at the point of use. Tests updated to the SecretStr type. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
359 lines
12 KiB
Python
359 lines
12 KiB
Python
"""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
|
|
import secrets
|
|
|
|
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:
|
|
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 the actual cited bibkeys from the page's inline [BibKey #loc]
|
|
# citations (audit C-3: this previously returned the concept slug itself,
|
|
# not the cited sources).
|
|
from codex.wiki import _parse_claims
|
|
|
|
sources = sorted({claim.bibkey for claim in _parse_claims(markdown)})
|
|
|
|
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()
|
|
token = settings.mcp_auth_token.get_secret_value() if settings.mcp_auth_token else ""
|
|
if not token:
|
|
raise RuntimeError(
|
|
"HTTP transport requires MCP_AUTH_TOKEN to be set. "
|
|
"Set the environment variable or add it to .env."
|
|
)
|
|
return 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", "")
|
|
# Constant-time comparison to avoid a token-length/equality timing
|
|
# side-channel (audit S-2).
|
|
if not secrets.compare_digest(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()
|