"""Integration with codebase-memory-mcp (https://github.com/DeusData/codebase-memory-mcp). The Code agent gains "codebase memory" — a persistent knowledge graph of the working directory (functions, classes, call chains, routes). We bridge to the single static binary through its documented CLI mode:: codebase-memory-mcp cli --raw '' so there is no long-lived MCP process to manage. Indexing and queries run locally; nothing leaves the machine. """ from __future__ import annotations import json import shutil import subprocess import sys from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from ..providers.base import ToolSpec BINARY_NAME = "codebase-memory-mcp" _QUERY_TIMEOUT = 90 _INDEX_TIMEOUT = 900 class CodebaseMemoryError(RuntimeError): pass def resolve_binary(configured: str = "") -> Optional[str]: """Return a usable binary path, or None if not installed.""" if configured: p = Path(configured).expanduser() if p.exists(): return str(p) return shutil.which(BINARY_NAME) def pip_install() -> Tuple[bool, str]: """Install the codebase-memory-mcp package into the current environment.""" try: proc = subprocess.run( [sys.executable, "-m", "pip", "install", "--upgrade", BINARY_NAME], capture_output=True, text=True, timeout=900, ) except Exception as exc: # noqa: BLE001 return False, str(exc) if proc.returncode == 0: return True, "Installed codebase-memory-mcp." return False, (proc.stderr or proc.stdout or "pip install failed")[-400:] def _extract_json(text: str): """Parse JSON from CLI output that may contain log lines before/after it.""" text = (text or "").strip() if not text: return None try: return json.loads(text) except json.JSONDecodeError: pass starts = [p for p in (text.find("{"), text.find("[")) if p >= 0] if not starts: return None start = min(starts) for end in (text.rfind("}"), text.rfind("]")): if end > start: try: return json.loads(text[start:end + 1]) except json.JSONDecodeError: continue return None class CodebaseMemory: def __init__(self, binary_path: str = ""): self.binary = resolve_binary(binary_path) @property def available(self) -> bool: return self.binary is not None def _run(self, tool: str, args: Dict[str, Any], timeout: int) -> Dict[str, Any]: if not self.binary: raise CodebaseMemoryError( "codebase-memory-mcp is not installed. See the instructions in Settings." ) # Note: no '--raw' flag — some binary versions reject it ("unknown tool: --raw"). cmd = [self.binary, "cli", tool, json.dumps(args, ensure_ascii=False)] try: proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) except subprocess.TimeoutExpired: raise CodebaseMemoryError(f"'{tool}' timed out ({timeout}s).") except OSError as exc: raise CodebaseMemoryError(f"Could not run the binary: {exc}") out = (proc.stdout or "").strip() # The binary may print download/init logs around the JSON — extract it. parsed = _extract_json(out) if parsed is not None: return parsed if proc.returncode != 0: detail = (proc.stderr or out or f"{tool} error (exit {proc.returncode})").strip() raise CodebaseMemoryError(detail[:300]) return {"raw": out} # ---- high level ops --------------------------------------------- def index_repository(self, repo_path: str) -> Dict[str, Any]: return self._run("index_repository", {"repo_path": str(repo_path)}, _INDEX_TIMEOUT) def list_projects(self) -> Dict[str, Any]: return self._run("list_projects", {}, _QUERY_TIMEOUT) def call(self, tool: str, args: Dict[str, Any]) -> Dict[str, Any]: timeout = _INDEX_TIMEOUT if tool == "index_repository" else _QUERY_TIMEOUT return self._run(tool, args, timeout) # Agent-facing tools (namespaced cmem_*). Read-only — never gated by permission. CMEM_TOOL_SPECS: List[ToolSpec] = [ ToolSpec( name="cmem_get_architecture", description="Architecture overview of the indexed codebase: languages, packages, routes, hotspots, clusters.", parameters={"type": "object", "properties": {}}, ), ToolSpec( name="cmem_search_graph", description="Find nodes by label/name pattern/file pattern in the knowledge graph (functions, classes, routes...).", parameters={ "type": "object", "properties": { "name_pattern": {"type": "string", "description": "Name regex, e.g. '.*Handler.*'"}, "label": {"type": "string", "description": "Function | Class | Route ..."}, "file_pattern": {"type": "string"}, "limit": {"type": "integer"}, }, }, ), ToolSpec( name="cmem_trace_path", description="Trace the call graph: who calls a function and what it calls (BFS, depth 1-5).", parameters={ "type": "object", "properties": { "function_name": {"type": "string"}, "direction": {"type": "string", "enum": ["inbound", "outbound", "both"]}, "depth": {"type": "integer"}, }, "required": ["function_name"], }, ), ToolSpec( name="cmem_get_code_snippet", description="Get the source code of a function by its qualified name.", parameters={ "type": "object", "properties": {"qualified_name": {"type": "string"}}, "required": ["qualified_name"], }, ), ToolSpec( name="cmem_search_code", description="Grep-like text search within the project's indexed files.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, ), ToolSpec( name="cmem_query_graph", description="Run a read-only Cypher-like query on the knowledge graph.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, ), ] CMEM_TOOL_NAMES = {t.name for t in CMEM_TOOL_SPECS} _CLI_NAME = {t.name: t.name[len("cmem_"):] for t in CMEM_TOOL_SPECS} def make_executor(mem: CodebaseMemory): """Return an executor(name, args) -> {ok, output} for cmem_* tools.""" def execute(name: str, args: Dict[str, Any]) -> Dict[str, Any]: cli_tool = _CLI_NAME.get(name) if not cli_tool: return {"ok": False, "output": f"Unsupported codebase-memory tool: {name}"} try: result = mem.call(cli_tool, args or {}) except CodebaseMemoryError as exc: return {"ok": False, "output": str(exc)} text = json.dumps(result, ensure_ascii=False, indent=2) if len(text) > 12000: text = text[:12000] + "\n… (truncated)" return {"ok": True, "output": text} return execute