"""The Zyte MCP server: the graph_* tool surface over the retrieval engine. mcp-graph-server HTTP on 127.0.0.1:$PORT - /mcp, GET /prime, POST /log mcp-graph-server --stdio stdio transport for clients that spawn their server Environment: ``ZYTE_PROJECT_ROOT`` (project directory, default cwd), ``ZYTE_DATA_DIR`` (default ``/.zyte``, used for the usage log), ``ZYTE_TURN_READ_BUDGET_CHARS``, ``ZYTE_HARD_MAX_READ_CHARS``, ``ZYTE_FALLBACK_MAX_CALLS_PER_TURN``, ``ZYTE_RETRIEVE_CACHE_TTL_SEC``. Tool logic lives in plain functions so it needs nothing beyond the stdlib; ``main()`` is the only place the mcp/starlette/uvicorn extra is imported. """ from __future__ import annotations import json import os import shutil import sqlite3 import subprocess import sys import time from pathlib import Path from typing import Any from . import ( change_analyze, claim_check, fact, knowledge_events, knowledge_query, neighbors, obligation_list, obligation_resolve, prime, probe_get, probe_plan, probe_run, read, register_edits, register_reads, remember, retrieve, scan, snapshot_capture, status, ) # Supplementary-work caps the policy documents promise per confidence level: # (fallback_rg calls, extra graph_read calls). CONFIDENCE_CAPS = {"high": (0, 0), "medium": (1, 2), "low": (2, 3)} MIN_GRAPH_FILES = 5 RG_TIMEOUT_SECONDS = 5 RG_MAX_COUNT = 50 RG_MAX_LINES = 100 _rg_calls = 0 _retrieve_cache: dict[tuple[str, str, int, str], tuple[float, dict[str, Any]]] = {} def _root() -> Path: value = os.environ.get("ZYTE_PROJECT_ROOT", "").strip() return Path(value).resolve() if value else Path.cwd() def _data_dir() -> Path: value = os.environ.get("ZYTE_DATA_DIR", "").strip() return Path(value) if value else _root() / ".zyte" def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, "") or default) except ValueError: return default def _budget_chars() -> int: return max(0, _env_int("ZYTE_TURN_READ_BUDGET_CHARS", 18000)) def _max_read_chars() -> int: return max(0, _env_int("ZYTE_HARD_MAX_READ_CHARS", 4000)) def _rg_cap() -> int: return max(0, _env_int("ZYTE_FALLBACK_MAX_CALLS_PER_TURN", 1)) def _cached_retrieve(query: str, limit: int) -> dict[str, Any]: ttl = max(0, _env_int("ZYTE_RETRIEVE_CACHE_TTL_SEC", 900)) state = status(_root()) generation = str(state.get("generated_at") or "missing") key = (str(_root()), query, limit, generation) now = time.time() if ttl: hit = _retrieve_cache.get(key) if hit and now - hit[0] < ttl: return hit[1] result = retrieve(_root(), query, budget_chars=_budget_chars(), max_read_chars=_max_read_chars(), limit=limit) if ttl: if len(_retrieve_cache) >= 32: _retrieve_cache.clear() _retrieve_cache[key] = (now, result) return result def _refs(items: list[dict[str, Any]]) -> list[str]: out: list[str] = [] for item in items: path = item.get("path") if not isinstance(path, str) or not path: continue symbol = item.get("symbol") ref = f"{path}::{symbol}" if isinstance(symbol, str) and symbol else path if ref not in out: out.append(ref) return out def continue_impl(task: str = "") -> dict[str, Any]: """graph_continue: ranked starting set plus the supplementary-work caps.""" global _rg_calls _rg_calls = 0 root = _root() state = status(root) if not state["has_graph"]: if not root.is_dir(): return {"needs_project": True, "skip": False, "recommended_files": [], "confidence": "low"} return {"needs_project": True, "skip": False, "recommended_files": [], "confidence": "low"} if state["files"] < MIN_GRAPH_FILES: return {"needs_project": False, "skip": True, "recommended_files": [], "confidence": "low"} result = _cached_retrieve(str(task or ""), limit=12) confidence = result.get("confidence", "low") greps, files = CONFIDENCE_CAPS.get(confidence, CONFIDENCE_CAPS["low"]) return { "needs_project": False, "skip": False, "recommended_files": _refs(result.get("items") or []), "confidence": confidence, "stale": bool(result.get("stale")), "max_supplementary_greps": min(greps, _rg_cap()), "max_supplementary_files": files, } def scan_impl(project: str = "", force: bool = False) -> dict[str, Any]: """graph_scan: build or refresh the graph for ``project`` (default: the server root).""" _retrieve_cache.clear() target = Path(project).resolve() if project else _root() return scan(target, force=bool(force)) def fact_impl(query: str) -> dict[str, Any]: """graph_fact: extractive answer for a single-fact question, before any pack.""" return fact(_root(), query) def retrieve_impl(query: str, limit: int = 12) -> dict[str, Any]: """graph_retrieve: ranked excerpts for ``query`` within the read budget.""" return _cached_retrieve(str(query or ""), limit=max(0, int(limit))) def read_impl(file: str) -> dict[str, Any]: """graph_read: one file or one ``path::symbol`` slice, within the hard read cap.""" result = read(_root(), str(file or ""), max_chars=_max_read_chars()) if result.get("path"): register_reads(_root(), [str(file)]) return result def neighbors_impl(file: str, depth: int = 1) -> dict[str, Any]: """graph_neighbors: import and containment edges around a file or symbol.""" return neighbors(_root(), str(file or ""), depth=max(1, int(depth))) def register_edit_impl(files: list[str]) -> dict[str, Any]: """graph_register_edit: feed edited ``path`` / ``path::symbol`` refs into session weighting.""" _retrieve_cache.clear() entries: list[dict[str, Any]] = [] for ref in files or []: if not isinstance(ref, str) or not ref: continue path, _, symbol = ref.partition("::") entries.append({"path": path, "symbol": symbol or None}) return register_edits(_root(), entries) def add_memory_impl( type: str, # noqa: A002 - part of the published contract content: str, tags: list[str] | None = None, files: list[str] | None = None, ) -> dict[str, Any]: """graph_add_memory: append a durable project note.""" return remember(_root(), type, content, tags, files) def snapshot_capture_impl(source: str = "working_tree") -> dict[str, object]: """byte_snapshot_capture: freeze the current repository/config/lock state.""" return snapshot_capture(_root(), source) def knowledge_query_impl( question: str, snapshot_id: str = "", max_claims: int = 20, ) -> dict[str, object]: """byte_knowledge_query: retrieve and recheck prior claims without running code.""" return knowledge_query( _root(), str(question), snapshot_id=snapshot_id or None, budget={"max_claims": max(1, int(max_claims))}, ) def claim_check_impl(claim_ids: list[str], target_snapshot_id: str = "") -> dict[str, object]: """byte_claim_check: revalidate named claims against one snapshot.""" return claim_check(_root(), claim_ids, target_snapshot_id or None) def change_analyze_impl( base_snapshot_id: str, target_snapshot_id: str, intent_ref: str = "" ) -> dict[str, object]: """byte_change_analyze: open obligations only for invalidated claim dependencies.""" return change_analyze(_root(), base_snapshot_id, target_snapshot_id, intent_ref) def knowledge_events_impl(after_cursor: int = 0, limit: int = 100) -> dict[str, object]: """byte_knowledge_events: incremental claim, obligation, and probe events.""" return knowledge_events(_root(), after_cursor, limit) def obligation_list_impl(status_filter: str = "open", limit: int = 100) -> dict[str, object]: """byte_obligation_list: pending or resolved repository obligations.""" return obligation_list(_root(), status_filter, limit) def obligation_resolve_impl( obligation_id: str, resolution: str, rationale: str ) -> dict[str, object]: """byte_obligation_resolve: close an obligation with an auditable rationale.""" return obligation_resolve(_root(), obligation_id, resolution, rationale) def probe_plan_impl( snapshot_id: str, hypothesis: str, command: list[str], oracle_ref: dict[str, object], policy_id: str, claim_id: str = "", working_directory: str = ".", fixture: dict[str, object] | None = None, intervention: dict[str, object] | None = None, ) -> dict[str, object]: """byte_probe_plan: persist a challenge and independent oracle without executing it.""" return probe_plan( _root(), snapshot_id, hypothesis, command, oracle_ref, policy_id=policy_id, claim_id=claim_id or None, working_directory=working_directory, fixture=fixture, intervention=intervention, ) def probe_run_impl(probe_id: str, policy_id: str = "") -> dict[str, object]: """byte_probe_run: execute under a configured sandbox policy or refuse safely.""" return probe_run(_root(), probe_id, policy_id or None) def probe_get_impl(probe_id: str) -> dict[str, object]: """byte_probe_get: read one immutable probe plan.""" return probe_get(_root(), probe_id) def fallback_rg_impl(pattern: str, glob: str = "") -> dict[str, Any]: """fallback_rg: capped ripgrep escape hatch for when the graph answer is thin.""" global _rg_calls if _rg_calls >= _rg_cap(): return {"matches": [], "error": f"fallback_rg cap reached ({_rg_cap()} per turn)"} _rg_calls += 1 binary = shutil.which("rg") if not binary: return {"matches": [], "error": "ripgrep (rg) is not installed"} command = [binary, "-n", "--max-count", str(RG_MAX_COUNT), "--no-messages"] if glob: command += ["-g", glob] command += ["--", str(pattern or ""), "."] try: completed = subprocess.run( command, cwd=_root(), capture_output=True, text=True, timeout=RG_TIMEOUT_SECONDS, ) except (OSError, subprocess.TimeoutExpired) as error: return {"matches": [], "error": f"rg failed: {error}"} lines = (completed.stdout or "").splitlines() return {"matches": lines[:RG_MAX_LINES], "truncated": len(lines) > RG_MAX_LINES} def prime_text() -> str: """Plain-text session primer served at GET /prime.""" root = _root() state = status(root) lines = [] if state["has_graph"]: freshness = "stale" if state["stale"] else "fresh" lines.append( f"Byte repository graph: {state['files']} files, {state['symbols']} symbols ({freshness}). " "Call graph_continue before exploring." ) else: lines.append("Byte repository graph: not built yet. Call graph_scan, then graph_continue.") items = prime(root).get("items") or [] if items: lines.append("") lines.append("Recent project notes:") for item in items: kind = item.get("type", "fact") content = str(item.get("content", "")).strip() files = item.get("files") or [] suffix = f" [{', '.join(files)}]" if files else "" lines.append(f"- ({kind}) {content}{suffix}") try: obligations = obligation_list(root, "open", 10).get("items") or [] except (FileNotFoundError, RuntimeError, sqlite3.Error): obligations = [] if obligations: lines.append("") lines.append("Open Byte repository obligations:") for item in obligations: lines.append(f"- {item.get('reason', 'recheck required')} [claim {item.get('claim_id', '')}]") return "\n".join(lines) + "\n" def append_log(payload: Any) -> bool: """Append one usage record to ``/usage_log.jsonl``. Never raises.""" try: data_dir = _data_dir() data_dir.mkdir(parents=True, exist_ok=True) record = payload if isinstance(payload, dict) else {"raw": payload} record = {**record, "ts": time.time()} with open(data_dir / "usage_log.jsonl", "a", encoding="utf-8") as handle: handle.write(json.dumps(record) + "\n") return True except Exception: return False def _build_mcp(): # pragma: no cover - exercised via the stdio handshake smoke from mcp.server.fastmcp import FastMCP mcp = FastMCP("byte") mcp.tool(name="graph_continue", description=( "Call this FIRST, before any file exploration. Returns recommended_files ranked for the task, " "a confidence verdict, and hard caps for supplementary greps/reads." ))(continue_impl) mcp.tool(name="graph_scan", description="Build or refresh the repository graph.")(scan_impl) mcp.tool(name="graph_fact", description=( "Extractive answer for a single-fact question (a constant's value, a definition site). " "Try before graph_retrieve; a miss costs nothing." ))(fact_impl) mcp.tool(name="graph_retrieve", description="Ranked code excerpts for a query, within the read budget.")( retrieve_impl ) mcp.tool(name="graph_read", description=( "Read one file, or one symbol via 'path::symbol'. One call per file; never batch." ))(read_impl) mcp.tool(name="graph_neighbors", description="Import/containment neighbors of a file or symbol.")(neighbors_impl) mcp.tool(name="graph_register_edit", description=( "After edits, register changed 'path' or 'path::symbol' refs so ranking learns this session." ))(register_edit_impl) mcp.tool(name="graph_add_memory", description=( "Store a durable note (decision, task, next, fact, blocker) for future sessions." ))(add_memory_impl) mcp.tool(name="fallback_rg", description=( "Capped ripgrep escape hatch. Only after graph_continue, only within max_supplementary_greps." ))(fallback_rg_impl) mcp.tool(name="byte_snapshot_capture", description=( "Capture an immutable source, configuration, lockfile, and analyzer snapshot before making claims." ))(snapshot_capture_impl) mcp.tool(name="byte_knowledge_query", description=( "Find related repository claims and recheck whether their evidence still applies. Never runs code." ))(knowledge_query_impl) mcp.tool(name="byte_claim_check", description=( "Recheck named claim IDs against an immutable snapshot; outcome and applicability stay separate." ))(claim_check_impl) mcp.tool(name="byte_change_analyze", description=( "Compare snapshots and open scoped obligations for changed or unknown claim dependencies." ))(change_analyze_impl) mcp.tool(name="byte_knowledge_events", description=( "Read claim invalidations, obligations, and probe completions after a monotonic event cursor." ))(knowledge_events_impl) mcp.tool(name="byte_obligation_list", description=( "List repository obligations; affected behavior does not automatically mean a file must be edited." ))(obligation_list_impl) mcp.tool(name="byte_obligation_resolve", description=( "Resolve an obligation as verified, fixed, accepted, or not applicable with a rationale." ))(obligation_resolve_impl) mcp.tool(name="byte_probe_plan", description=( "Persist a reproducible counterexample probe and its oracle. Planning never executes code." ))(probe_plan_impl) mcp.tool(name="byte_probe_run", description=( "Run a planned probe under its exact user-configured sandbox policy. Policy bodies cannot be supplied here." ))(probe_run_impl) mcp.tool(name="byte_probe_get", description="Read one immutable Byte probe plan.")(probe_get_impl) return mcp def _http_app(mcp): # pragma: no cover - exercised via the live server smoke from starlette.responses import JSONResponse, PlainTextResponse async def prime_route(request): return PlainTextResponse(prime_text()) async def log_route(request): try: payload = json.loads((await request.body()).decode("utf-8", errors="replace") or "{}") except ValueError: payload = {} return JSONResponse({"ok": append_log(payload)}) try: # Modern mcp SDKs register extra routes on the FastMCP app directly. mcp.custom_route("/prime", methods=["GET"])(prime_route) mcp.custom_route("/log", methods=["POST"])(log_route) return mcp.streamable_http_app() except AttributeError: # Older SDKs: wrap the MCP app and drive its session manager ourselves. import contextlib from starlette.applications import Starlette from starlette.routing import Mount, Route inner = mcp.streamable_http_app() @contextlib.asynccontextmanager async def lifespan(app): async with mcp.session_manager.run(): yield return Starlette( routes=[ Route("/prime", prime_route, methods=["GET"]), Route("/log", log_route, methods=["POST"]), Mount("/", app=inner), ], lifespan=lifespan, ) def main(argv: list[str] | None = None) -> int: arguments = sys.argv[1:] if argv is None else argv mcp = _build_mcp() if "--stdio" in arguments: mcp.run(transport="stdio") return 0 import uvicorn port = _env_int("PORT", 8080) uvicorn.run(_http_app(mcp), host="127.0.0.1", port=port, log_level="warning") return 0 if __name__ == "__main__": raise SystemExit(main())