|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""MCP server that exposes OpenMC's RAG semantic search to AI coding agents. |
| 3 | +
|
| 4 | +This is the entry point for the MCP (Model Context Protocol) server registered |
| 5 | +in .mcp.json at the repo root. When an MCP-capable agent (e.g. Claude Code) |
| 6 | +opens a session in this repository, it launches this server as a subprocess |
| 7 | +(via start_server.sh) and the tools defined here appear in the agent's tool |
| 8 | +list automatically. |
| 9 | +
|
| 10 | +The server is long-lived — it stays running for the duration of the agent |
| 11 | +session. This matters for session state: the first RAG search call returns |
| 12 | +an index status message instead of results, prompting the agent to ask the |
| 13 | +user whether to rebuild the index. That first-call flag resets each session. |
| 14 | +
|
| 15 | +Tools exposed: |
| 16 | + openmc_rag_search — semantic search across the codebase and docs |
| 17 | + openmc_rag_rebuild — rebuild the RAG vector index |
| 18 | +
|
| 19 | +The actual search/indexing logic lives in the rag/ subdirectory (openmc_search.py, |
| 20 | +indexer.py, chunker.py, embeddings.py). This file is just the MCP interface |
| 21 | +layer and session state management. |
| 22 | +""" |
| 23 | + |
| 24 | +from mcp.server.fastmcp import FastMCP |
| 25 | +import json |
| 26 | +import logging |
| 27 | +import subprocess |
| 28 | +import sys |
| 29 | +from datetime import datetime |
| 30 | +from pathlib import Path |
| 31 | + |
| 32 | +# MCP communicates over stdin/stdout with JSON-RPC framing. Several libraries |
| 33 | +# (httpx, huggingface_hub, sentence_transformers) emit log messages and |
| 34 | +# progress bars to stderr by default. While stderr isn't part of the MCP |
| 35 | +# transport, noisy output there can confuse agent tooling, so we silence it. |
| 36 | +logging.getLogger("httpx").setLevel(logging.WARNING) |
| 37 | +logging.getLogger("huggingface_hub").setLevel(logging.ERROR) |
| 38 | +logging.getLogger("sentence_transformers").setLevel(logging.WARNING) |
| 39 | + |
| 40 | +# Path constants. This file lives at .claude/tools/openmc_mcp_server.py, |
| 41 | +# so parents[2] is the OpenMC repo root. |
| 42 | +OPENMC_ROOT = Path(__file__).resolve().parents[2] |
| 43 | +CACHE_DIR = OPENMC_ROOT / ".claude" / "cache" |
| 44 | +INDEX_DIR = CACHE_DIR / "rag_index" |
| 45 | +METADATA_FILE = INDEX_DIR / "metadata.json" |
| 46 | + |
| 47 | +# The RAG modules (openmc_search, indexer, etc.) live in .claude/tools/rag/. |
| 48 | +# We add that directory to sys.path so we can import them directly. |
| 49 | +TOOLS_DIR = Path(__file__).resolve().parent |
| 50 | +sys.path.insert(0, str(TOOLS_DIR / "rag")) |
| 51 | + |
| 52 | +mcp = FastMCP("openmc-code-tools") |
| 53 | + |
| 54 | +# First-call flag: the first openmc_rag_search call of each session returns |
| 55 | +# index status info instead of search results, so the agent can ask the user |
| 56 | +# whether to rebuild. This resets when the server process restarts (i.e. each |
| 57 | +# new agent session). |
| 58 | +_rag_first_call = True |
| 59 | + |
| 60 | + |
| 61 | +# --------------------------------------------------------------------------- |
| 62 | +# Helpers |
| 63 | +# --------------------------------------------------------------------------- |
| 64 | + |
| 65 | +def _get_current_branch(): |
| 66 | + """Get the current git branch name.""" |
| 67 | + try: |
| 68 | + result = subprocess.run( |
| 69 | + ["git", "rev-parse", "--abbrev-ref", "HEAD"], |
| 70 | + capture_output=True, text=True, cwd=str(OPENMC_ROOT), |
| 71 | + ) |
| 72 | + if result.returncode != 0 or not result.stdout.strip(): |
| 73 | + return "unknown" |
| 74 | + return result.stdout.strip() |
| 75 | + except Exception: |
| 76 | + return "unknown" |
| 77 | + |
| 78 | + |
| 79 | +def _get_index_metadata(): |
| 80 | + """Read index build metadata, or None if unavailable.""" |
| 81 | + if not METADATA_FILE.exists(): |
| 82 | + return None |
| 83 | + try: |
| 84 | + return json.loads(METADATA_FILE.read_text()) |
| 85 | + except Exception: |
| 86 | + return None |
| 87 | + |
| 88 | + |
| 89 | +def _save_index_metadata(): |
| 90 | + """Save index build metadata alongside the index.""" |
| 91 | + metadata = { |
| 92 | + "built_at": datetime.now().strftime("%Y-%m-%d %H:%M"), |
| 93 | + "branch": _get_current_branch(), |
| 94 | + } |
| 95 | + METADATA_FILE.write_text(json.dumps(metadata, indent=2)) |
| 96 | + |
| 97 | + |
| 98 | +def _check_index_first_call(): |
| 99 | + """On the first RAG call of the session, return a status message for the |
| 100 | + agent to relay to the user. Returns None if no prompt is needed (should |
| 101 | + not happen — we always prompt on first call).""" |
| 102 | + current_branch = _get_current_branch() |
| 103 | + |
| 104 | + if not INDEX_DIR.exists(): |
| 105 | + return ( |
| 106 | + "No RAG index found. Building one takes ~5 minutes but greatly " |
| 107 | + "improves code navigation by enabling semantic search across the " |
| 108 | + "entire OpenMC codebase (C++, Python, and docs).\n\n" |
| 109 | + "IMPORTANT: Use the AskUserQuestion tool to ask the user whether " |
| 110 | + "to build the index now (you would then call openmc_rag_rebuild) " |
| 111 | + "or proceed without it." |
| 112 | + ) |
| 113 | + |
| 114 | + meta = _get_index_metadata() |
| 115 | + if meta: |
| 116 | + built_at = meta.get("built_at", "unknown time") |
| 117 | + built_branch = meta.get("branch", "unknown") |
| 118 | + return ( |
| 119 | + f"Existing RAG index found — built at {built_at} on branch " |
| 120 | + f"'{built_branch}'. Current branch is '{current_branch}'.\n\n" |
| 121 | + f"REQUIRED: You must use the AskUserQuestion tool now to ask the " |
| 122 | + f"user whether to rebuild the index (you would then call " |
| 123 | + f"openmc_rag_rebuild) or use the existing one. Do not skip this " |
| 124 | + f"step — the user may have uncommitted changes. Do not decide " |
| 125 | + f"on their behalf." |
| 126 | + ) |
| 127 | + |
| 128 | + return ( |
| 129 | + f"RAG index found but has no build metadata. " |
| 130 | + f"Current branch is '{current_branch}'.\n\n" |
| 131 | + f"REQUIRED: You must use the AskUserQuestion tool now to ask the " |
| 132 | + f"user whether to rebuild the index (you would then call " |
| 133 | + f"openmc_rag_rebuild) or use the existing one. Do not skip this " |
| 134 | + f"step. Do not decide on their behalf." |
| 135 | + ) |
| 136 | + |
| 137 | + |
| 138 | +# --------------------------------------------------------------------------- |
| 139 | +# Tools |
| 140 | +# --------------------------------------------------------------------------- |
| 141 | + |
| 142 | +@mcp.tool() |
| 143 | +def openmc_rag_search( |
| 144 | + query: str = "", |
| 145 | + related_file: str = "", |
| 146 | + scope: str = "code", |
| 147 | + top_k: int = 10, |
| 148 | +) -> str: |
| 149 | + """Semantic search across the OpenMC codebase and documentation. |
| 150 | +
|
| 151 | + Finds code by meaning, not just text match — surfaces related code across |
| 152 | + subsystems even when naming differs. Use for discovery and exploration |
| 153 | + before reaching for grep. Covers C++, Python, and RST docs. |
| 154 | +
|
| 155 | + Args: |
| 156 | + query: Search query (e.g. "particle weight adjustment variance reduction") |
| 157 | + related_file: Instead of a text query, find code related to this file |
| 158 | + scope: "code" (default), "docs", or "all" |
| 159 | + top_k: Number of results to return (default 10) |
| 160 | + """ |
| 161 | + global _rag_first_call |
| 162 | + |
| 163 | + # First call of the session — prompt the agent to check with the user |
| 164 | + if _rag_first_call: |
| 165 | + _rag_first_call = False |
| 166 | + status = _check_index_first_call() |
| 167 | + if status: |
| 168 | + return status |
| 169 | + |
| 170 | + # No index available |
| 171 | + if not INDEX_DIR.exists(): |
| 172 | + return ( |
| 173 | + "No RAG index available. Call openmc_rag_rebuild() to build one " |
| 174 | + "(takes ~5 minutes)." |
| 175 | + ) |
| 176 | + |
| 177 | + if not query and not related_file: |
| 178 | + return "Error: provide either 'query' or 'related_file'." |
| 179 | + |
| 180 | + if query and related_file: |
| 181 | + return "Error: provide 'query' or 'related_file', not both." |
| 182 | + |
| 183 | + if scope not in ("code", "docs", "all"): |
| 184 | + return f"Error: scope must be 'code', 'docs', or 'all' (got '{scope}')." |
| 185 | + |
| 186 | + if top_k < 1: |
| 187 | + return f"Error: top_k must be at least 1 (got {top_k})." |
| 188 | + |
| 189 | + try: |
| 190 | + from openmc_search import ( |
| 191 | + get_db_and_embedder, search_table, format_results, search_related, |
| 192 | + ) |
| 193 | + |
| 194 | + db, embedder = get_db_and_embedder() |
| 195 | + |
| 196 | + if related_file: |
| 197 | + results = search_related(db, embedder, related_file, top_k) |
| 198 | + return format_results(results, f"Code related to {related_file}") |
| 199 | + elif scope == "all": |
| 200 | + code_results = search_table(db, embedder, "code", query, top_k) |
| 201 | + doc_results = search_table(db, embedder, "docs", query, top_k) |
| 202 | + return (format_results(code_results, "Code") + "\n" |
| 203 | + + format_results(doc_results, "Documentation")) |
| 204 | + elif scope == "docs": |
| 205 | + results = search_table(db, embedder, "docs", query, top_k) |
| 206 | + return format_results(results, "Documentation") |
| 207 | + else: |
| 208 | + results = search_table(db, embedder, "code", query, top_k) |
| 209 | + return format_results(results, "Code") |
| 210 | + except Exception as e: |
| 211 | + return f"Error during search: {e}" |
| 212 | + |
| 213 | + |
| 214 | +@mcp.tool() |
| 215 | +def openmc_rag_rebuild() -> str: |
| 216 | + """Rebuild the RAG semantic search index from the current codebase. |
| 217 | +
|
| 218 | + Chunks all C++, Python, and RST files, embeds them with a local |
| 219 | + sentence-transformers model, and stores in a LanceDB vector index. |
| 220 | + Takes ~5 minutes on 10 CPU cores. Call this after pulling new code |
| 221 | + or switching branches. |
| 222 | + """ |
| 223 | + global _rag_first_call |
| 224 | + _rag_first_call = False # no need to prompt after an explicit rebuild |
| 225 | + |
| 226 | + try: |
| 227 | + import io |
| 228 | + from indexer import build_index |
| 229 | + |
| 230 | + old_stdout = sys.stdout |
| 231 | + sys.stdout = captured = io.StringIO() |
| 232 | + try: |
| 233 | + build_index() |
| 234 | + finally: |
| 235 | + sys.stdout = old_stdout |
| 236 | + |
| 237 | + _save_index_metadata() |
| 238 | + |
| 239 | + branch = _get_current_branch() |
| 240 | + build_output = captured.getvalue() |
| 241 | + return ( |
| 242 | + f"Index rebuilt successfully on branch '{branch}'.\n\n" |
| 243 | + f"{build_output}" |
| 244 | + ) |
| 245 | + except Exception as e: |
| 246 | + return f"Error rebuilding index: {e}" |
| 247 | + |
| 248 | + |
| 249 | +if __name__ == "__main__": |
| 250 | + mcp.run() |
0 commit comments