"""Project-bound, read-only graph visualization over a validated DocForge index.""" from __future__ import annotations import atexit import json import os import secrets import socket import sqlite3 import subprocess import sys import threading import urllib.parse from collections.abc import Generator from contextlib import contextmanager from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from time import monotonic from .errors import DocForgeError from .index import APPLICATION_ID, INDEX_SCHEMA_VERSION, ProjectIndex, re_tokenize VISUALIZATION_TEMPLATE = "graph-browser@5" DEFAULT_EDGE_LIMIT = 100 MAX_EDGE_LIMIT = 400 DEFAULT_INITIAL_GRACE_SECONDS = 120.0 DEFAULT_LEASE_SECONDS = 180.0 LEASE_MONITOR_INTERVAL_SECONDS = 1.0 class _VisualizationHttpServer(ThreadingHTTPServer): daemon_threads = True allow_reuse_address = False class VisualizationIndexSnapshot: """Fast read model pinned to one index file validated by ``ProjectIndex.check``.""" _IDENTITY_KEYS = ( "project_id", "project_root_fingerprint", "revision", "source_hash", "adapter", "node_count", "edge_count", ) def __init__(self, index: ProjectIndex, checked: dict[str, object]) -> None: self.path = index.path self.title = index.project.descriptor.title self.max_query_chars = index.project.descriptor.limits.max_query_chars self.max_results = index.project.descriptor.limits.max_results self.max_depth = index.project.descriptor.limits.max_traversal_depth self.identity = {key: checked[key] for key in self._IDENTITY_KEYS} self._stat = self._safe_stat() @classmethod def from_spec(cls, spec: dict[str, object]) -> VisualizationIndexSnapshot: snapshot = cls.__new__(cls) snapshot.path = Path(str(spec["path"])) snapshot.title = str(spec["title"]) snapshot.max_query_chars = int(spec["max_query_chars"]) snapshot.max_results = int(spec["max_results"]) snapshot.max_depth = int(spec["max_depth"]) identity = spec["identity"] if not isinstance(identity, dict): raise DocForgeError("invalid_index", "Visualization identity is invalid") snapshot.identity = {key: identity[key] for key in cls._IDENTITY_KEYS} snapshot._stat = snapshot._safe_stat() return snapshot def spec(self) -> dict[str, object]: return { "path": str(self.path), "title": self.title, "max_query_chars": self.max_query_chars, "max_results": self.max_results, "max_depth": self.max_depth, "identity": dict(self.identity), } def overview(self) -> dict[str, object]: with self._connection() as connection: return self._result( title=self.title, node_count=connection.execute("SELECT COUNT(*) FROM nodes").fetchone()[0], edge_count=connection.execute("SELECT COUNT(*) FROM edges").fetchone()[0], families=_facet_rows(connection, "nodes", "family"), authorities=_facet_rows(connection, "nodes", "authority"), statuses=_facet_rows(connection, "nodes", "status"), relations=_facet_rows(connection, "edges", "relation"), max_results=self.max_results, snapshot=True, ) def search( self, *, query: str, family: str | None, limit: int, ) -> dict[str, object]: bounded = self._bounded_limit(limit) with self._connection() as connection: if query: if len(query) > self.max_query_chars: raise DocForgeError( "invalid_query", "Search query exceeds the configured limit" ) terms = re_tokenize(query) if not terms: raise DocForgeError("invalid_query", "Search query contains no searchable text") expression = " AND ".join( f'"{term.replace(chr(34), chr(34) * 2)}"' for term in terms ) family_clause = "AND nodes.family = ?" if family else "" values: tuple[object, ...] = ( (expression, family, bounded) if family else (expression, bounded) ) rows = connection.execute( """ SELECT nodes.*, bm25(node_fts) AS rank, snippet(node_fts, 3, '[', ']', ' … ', 18) AS snippet FROM node_fts JOIN nodes USING(node_id) WHERE node_fts MATCH ? """ + family_clause + " ORDER BY rank, nodes.node_id LIMIT ?", values, ).fetchall() results = [] for row in rows: item = _node_dict(row, include_content=False) item.update({"rank": row["rank"], "snippet": row["snippet"]}) results.append(item) else: family_clause = "WHERE family = ?" if family else "" values = (family, bounded) if family else (bounded,) rows = connection.execute( f"SELECT * FROM nodes {family_clause} ORDER BY node_id LIMIT ?", values, ).fetchall() results = [_node_dict(row, include_content=False) for row in rows] return self._result( query=query, family=family, count=len(results), results=results, snapshot=True, ) def node(self, node_id: str, *, depth: int, limit: int) -> dict[str, object]: if type(depth) is not int or depth < 1 or depth > self.max_depth: raise DocForgeError("invalid_depth", "Traversal depth is outside the configured limit") if type(limit) is not int or limit < 1 or limit > MAX_EDGE_LIMIT: raise DocForgeError( "invalid_limit", "Visualization edge limit exceeds the fixed safety boundary", maximum=MAX_EDGE_LIMIT, ) with self._connection() as connection: root_row = connection.execute( "SELECT * FROM nodes WHERE node_id = ?", (node_id,) ).fetchone() if root_row is None: raise DocForgeError( "missing_node", "No node has the requested stable ID", node_id=node_id, ) visited = {node_id} frontier = {node_id} selected: list[dict[str, str]] = [] selected_keys: set[tuple[str, str, str]] = set() truncated = False for _ in range(depth): if not frontier or len(selected) >= limit: break placeholders = ",".join("?" for _ in frontier) values = tuple(sorted(frontier)) rows = connection.execute( "SELECT source_id, relation, target_id FROM edges " f"WHERE source_id IN ({placeholders}) OR target_id IN ({placeholders}) " "ORDER BY source_id, relation, target_id LIMIT ?", (*values, *values, limit + 1), ).fetchall() next_frontier: set[str] = set() for row in rows: key = (row["source_id"], row["relation"], row["target_id"]) if key in selected_keys: continue if len(selected) >= limit: truncated = True break selected_keys.add(key) edge = { "source_id": row["source_id"], "relation": row["relation"], "target_id": row["target_id"], } selected.append(edge) for candidate in (edge["source_id"], edge["target_id"]): if candidate not in visited: visited.add(candidate) next_frontier.add(candidate) frontier = next_frontier placeholders = ",".join("?" for _ in visited) node_rows = connection.execute( f"SELECT * FROM nodes WHERE node_id IN ({placeholders}) ORDER BY node_id", tuple(sorted(visited)), ).fetchall() return self._result( root=node_id, depth=depth, edge_limit=limit, truncated=truncated, node=_node_dict(root_row), nodes=[_node_dict(row, include_content=False) for row in node_rows], edges=selected, snapshot=True, ) def require_node(self, node_id: str) -> None: with self._connection() as connection: row = connection.execute("SELECT 1 FROM nodes WHERE node_id = ?", (node_id,)).fetchone() if row is None: raise DocForgeError( "missing_node", "No node has the requested stable ID", node_id=node_id, ) def _bounded_limit(self, value: int) -> int: if type(value) is not int or value < 1 or value > self.max_results: raise DocForgeError("invalid_limit", "Result limit is outside the configured range") return value def _safe_stat(self) -> tuple[int, int, int, int]: if ( self.path.is_symlink() or not self.path.is_file() or self.path.resolve(strict=True) != self.path ): raise DocForgeError("missing_index", "Validated visualization index is unavailable") stat = self.path.stat() return (stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns) @contextmanager def _connection(self) -> Generator[sqlite3.Connection, None, None]: if self._safe_stat() != self._stat: raise DocForgeError( "visualization_stale", "The validated index changed; invoke docforge_visualize again", ) connection: sqlite3.Connection | None = None try: connection = sqlite3.connect(f"file:{self.path}?mode=ro", uri=True) connection.row_factory = sqlite3.Row application_id = connection.execute("PRAGMA application_id").fetchone()[0] schema_version = connection.execute("PRAGMA user_version").fetchone()[0] if application_id != APPLICATION_ID or schema_version != INDEX_SCHEMA_VERSION: raise DocForgeError( "invalid_index", "Visualization index has an unsupported schema" ) metadata = dict(connection.execute("SELECT key, value FROM metadata")) if any(metadata.get(key) != str(value) for key, value in self.identity.items()): raise DocForgeError( "visualization_stale", "The validated index identity changed; invoke docforge_visualize again", ) yield connection if self._safe_stat() != self._stat: raise DocForgeError( "visualization_stale", "The validated index changed during the request", ) except DocForgeError: raise except (json.JSONDecodeError, KeyError, sqlite3.Error, TypeError) as error: raise DocForgeError( "invalid_index", "Visualization index is corrupt or unreadable" ) from error finally: if connection is not None: connection.close() def _result(self, **payload: object) -> dict[str, object]: return { "status": "ok", **self.identity, **payload, } class VisualizationRunner: """Start one token-protected loopback reader for one immutable project binding.""" def __init__( self, index: ProjectIndex | None, *, snapshot_spec: dict[str, object] | None = None, token: str | None = None, register_atexit: bool = True, initial_grace_seconds: float = DEFAULT_INITIAL_GRACE_SECONDS, lease_seconds: float = DEFAULT_LEASE_SECONDS, monitor_interval_seconds: float = LEASE_MONITOR_INTERVAL_SECONDS, ) -> None: if (index is None) == (snapshot_spec is None): raise ValueError("Provide exactly one visualization index or snapshot") if initial_grace_seconds <= 0 or lease_seconds <= 0 or monitor_interval_seconds <= 0: raise ValueError("Visualization lease durations must be positive") self.index = index self._snapshot_spec = snapshot_spec self._register_atexit = register_atexit self.initial_grace_seconds = initial_grace_seconds self.lease_seconds = lease_seconds self.monitor_interval_seconds = monitor_interval_seconds self._lock = threading.Lock() self._token = token or secrets.token_urlsafe(24) self._server: _VisualizationHttpServer | None = None self._thread: threading.Thread | None = None self._lease_thread: threading.Thread | None = None self._lease_stop = threading.Event() self._lease_last_activity = monotonic() self._lease_connected = False self._atexit_registered = False self._reader: VisualizationIndexSnapshot | None = None def start( self, *, node_id: str | None = None, query: str | None = None, depth: int = 1, ) -> dict[str, object]: if node_id is not None and query is not None: raise DocForgeError( "invalid_visualization_target", "Choose either one exact node ID or one search query", ) reader = ( VisualizationIndexSnapshot(self.index, self.index.check()) if self.index is not None else VisualizationIndexSnapshot.from_spec(self._snapshot_spec or {}) ) maximum_depth = reader.max_depth if type(depth) is not int or depth < 1 or depth > maximum_depth: raise DocForgeError( "invalid_depth", "Visualization depth is outside the configured traversal limit", ) if node_id is not None: reader.require_node(node_id) elif query is not None: reader.search(query=query, family=None, limit=1) with self._lock: self._reader = reader if self._server is None: runner = self class Handler(BaseHTTPRequestHandler): def do_GET(self) -> None: # noqa: N802 runner._handle_get(self) def do_HEAD(self) -> None: # noqa: N802 runner._handle_get(self, include_body=False) def do_POST(self) -> None: # noqa: N802 runner._respond_error( self, HTTPStatus.METHOD_NOT_ALLOWED, "method_not_allowed", "The visualization service is read-only", ) def do_PUT(self) -> None: # noqa: N802 self.do_POST() def do_PATCH(self) -> None: # noqa: N802 self.do_POST() def do_DELETE(self) -> None: # noqa: N802 self.do_POST() def log_message(self, _format: str, *_args: object) -> None: return self._server = _VisualizationHttpServer(("127.0.0.1", 0), Handler) self._lease_last_activity = monotonic() self._lease_connected = False self._lease_stop.clear() self._thread = threading.Thread( target=self._server.serve_forever, name="docforge-visualization", daemon=False, ) self._thread.start() self._lease_thread = threading.Thread( target=self._monitor_lease, name="docforge-visualization-lease", daemon=True, ) self._lease_thread.start() if self._register_atexit and not self._atexit_registered: atexit.register(self.stop) self._atexit_registered = True server = self._server assert server is not None port = int(server.server_address[1]) parameters: dict[str, str] = {"depth": str(depth)} if node_id is not None: parameters["node"] = node_id if query is not None: parameters["q"] = query query_string = urllib.parse.urlencode(parameters) url = f"http://127.0.0.1:{port}/{self._token}/" if query_string: url = f"{url}?{query_string}" return { "state": "running", "url": url, "bind": "127.0.0.1", "port": port, "template": VISUALIZATION_TEMPLATE, "read_only": True, "project_bound": True, "lifetime": { "policy": "browser_lease", "initial_grace_seconds": self.initial_grace_seconds, "lease_seconds": self.lease_seconds, }, "target": { "node_id": node_id, "query": query, "depth": depth, }, "snapshot": dict(reader.identity), } def stop(self) -> None: with self._lock: server = self._server thread = self._thread self._server = None self._thread = None self._lease_thread = None self._reader = None self._lease_stop.set() if server is None: return server.shutdown() server.server_close() if thread is not None and thread is not threading.current_thread(): thread.join(timeout=2) def is_running(self) -> bool: with self._lock: return self._server is not None def _touch_lease(self) -> None: with self._lock: if self._server is None: return self._lease_last_activity = monotonic() self._lease_connected = True def _monitor_lease(self) -> None: while not self._lease_stop.wait(self.monitor_interval_seconds): with self._lock: if self._server is None: return timeout = ( self.lease_seconds if self._lease_connected else self.initial_grace_seconds ) expired = monotonic() - self._lease_last_activity >= timeout if expired: self.stop() return def _handle_get(self, handler: BaseHTTPRequestHandler, *, include_body: bool = True) -> None: parsed = urllib.parse.urlparse(handler.path) prefix = f"/{self._token}" if parsed.path not in {prefix, f"{prefix}/"} and not parsed.path.startswith( f"{prefix}/api/" ): self._respond_error( handler, HTTPStatus.NOT_FOUND, "not_found", "Not found", include_body=include_body, ) return self._touch_lease() if parsed.path in {prefix, f"{prefix}/"}: self._respond( handler, _GRAPH_BROWSER_HTML.encode("utf-8"), "text/html; charset=utf-8", include_body=include_body, ) return try: params = urllib.parse.parse_qs(parsed.query, keep_blank_values=True) reader = self._current_reader() if parsed.path == f"{prefix}/api/overview": payload = reader.overview() elif parsed.path == f"{prefix}/api/heartbeat": payload = reader._result( viewer="alive", lease_seconds=self.lease_seconds, ) elif parsed.path == f"{prefix}/api/search": payload = self._search(reader, params) elif parsed.path == f"{prefix}/api/node": payload = self._node(reader, params) else: self._respond_error( handler, HTTPStatus.NOT_FOUND, "not_found", "Not found", include_body=include_body, ) return except DocForgeError as error: status = { "missing_node": HTTPStatus.NOT_FOUND, "stale_index": HTTPStatus.CONFLICT, "source_changed": HTTPStatus.CONFLICT, "stale_adapter_source": HTTPStatus.CONFLICT, "visualization_stale": HTTPStatus.CONFLICT, "invalid_query": HTTPStatus.BAD_REQUEST, "invalid_depth": HTTPStatus.BAD_REQUEST, "invalid_limit": HTTPStatus.BAD_REQUEST, }.get(error.code, HTTPStatus.SERVICE_UNAVAILABLE) self._respond_json( handler, {"status": "error", "error": error.as_dict()}, status=status, include_body=include_body, ) return except (TypeError, ValueError): self._respond_error( handler, HTTPStatus.BAD_REQUEST, "invalid_request", "Request parameters are invalid", include_body=include_body, ) return self._respond_json(handler, payload, include_body=include_body) def _search( self, reader: VisualizationIndexSnapshot, params: dict[str, list[str]], ) -> dict[str, object]: query = _one(params, "q").strip() family = _one(params, "family").strip() or None limit = _integer(_one(params, "limit") or "50") return reader.search(query=query, family=family, limit=limit) def _node( self, reader: VisualizationIndexSnapshot, params: dict[str, list[str]], ) -> dict[str, object]: node_id = _one(params, "id").strip() if not node_id: raise DocForgeError("missing_node", "One exact node ID is required") depth = _integer(_one(params, "depth") or "1") limit = _integer(_one(params, "limit") or str(DEFAULT_EDGE_LIMIT)) if limit > MAX_EDGE_LIMIT: raise DocForgeError( "invalid_limit", "Visualization edge limit exceeds the fixed safety boundary", maximum=MAX_EDGE_LIMIT, ) return reader.node(node_id, depth=depth, limit=limit) def _current_reader(self) -> VisualizationIndexSnapshot: with self._lock: reader = self._reader if reader is None: raise DocForgeError( "visualization_unavailable", "The visualization snapshot is unavailable", ) return reader @staticmethod def _respond( handler: BaseHTTPRequestHandler, payload: bytes, content_type: str, *, status: HTTPStatus = HTTPStatus.OK, include_body: bool = True, ) -> None: handler.send_response(status) handler.send_header("Content-Type", content_type) handler.send_header("Content-Length", str(len(payload))) handler.send_header("Cache-Control", "no-store") handler.send_header( "Content-Security-Policy", "default-src 'none'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; " "connect-src 'self'; img-src 'self' data:; base-uri 'none'; form-action 'none'; " "frame-ancestors 'none'", ) handler.send_header("X-Content-Type-Options", "nosniff") handler.send_header("X-Frame-Options", "DENY") handler.send_header("Referrer-Policy", "no-referrer") handler.end_headers() if include_body: handler.wfile.write(payload) @classmethod def _respond_json( cls, handler: BaseHTTPRequestHandler, body: dict[str, object], *, status: HTTPStatus = HTTPStatus.OK, include_body: bool = True, ) -> None: cls._respond( handler, json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8"), "application/json; charset=utf-8", status=status, include_body=include_body, ) @classmethod def _respond_error( cls, handler: BaseHTTPRequestHandler, status: HTTPStatus, code: str, message: str, *, include_body: bool = True, ) -> None: cls._respond_json( handler, { "status": "error", "error": {"code": code, "message": message, "details": {}}, }, status=status, include_body=include_body, ) class DetachedVisualizationRunner: """Launch a viewer worker that survives a short-lived MCP transport process.""" def __init__( self, index: ProjectIndex, *, initial_grace_seconds: float = DEFAULT_INITIAL_GRACE_SECONDS, lease_seconds: float = DEFAULT_LEASE_SECONDS, monitor_interval_seconds: float = LEASE_MONITOR_INTERVAL_SECONDS, ) -> None: self.index = index self.initial_grace_seconds = initial_grace_seconds self.lease_seconds = lease_seconds self.monitor_interval_seconds = monitor_interval_seconds self._process: subprocess.Popen[bytes] | None = None def start( self, *, node_id: str | None = None, query: str | None = None, depth: int = 1, ) -> dict[str, object]: if node_id is not None and query is not None: raise DocForgeError( "invalid_visualization_target", "Choose either one exact node ID or one search query", ) snapshot = VisualizationIndexSnapshot(self.index, self.index.check()) if type(depth) is not int or depth < 1 or depth > snapshot.max_depth: raise DocForgeError( "invalid_depth", "Visualization depth is outside the configured traversal limit", ) if node_id is not None: snapshot.require_node(node_id) elif query is not None: snapshot.search(query=query, family=None, limit=1) self.stop() parent_socket, child_socket = socket.socketpair() try: command = ( sys.executable, "-m", "docforge.visualization_worker", "--control-fd", str(child_socket.fileno()), ) process = subprocess.Popen( command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, close_fds=True, pass_fds=(child_socket.fileno(),), start_new_session=True, ) self._process = process child_socket.close() request = { "snapshot": snapshot.spec(), "token": secrets.token_urlsafe(24), "initial_grace_seconds": self.initial_grace_seconds, "lease_seconds": self.lease_seconds, "monitor_interval_seconds": self.monitor_interval_seconds, "owner_pid": os.getppid(), "target": { "node_id": node_id, "query": query, "depth": depth, }, } parent_socket.settimeout(10) parent_socket.sendall( json.dumps(request, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n" ) response = _receive_worker_response(parent_socket) except (OSError, subprocess.SubprocessError, ValueError) as error: self.stop() raise DocForgeError( "visualization_unavailable", "The detached visualization worker failed to start", ) from error finally: child_socket.close() parent_socket.close() if response.get("status") != "ok" or not isinstance(response.get("visualization"), dict): self.stop() raise DocForgeError( "visualization_unavailable", "The detached visualization worker rejected the snapshot", ) return response["visualization"] def stop(self) -> None: process = self._process self._process = None if process is None or process.poll() is not None: return process.terminate() try: process.wait(timeout=2) except subprocess.TimeoutExpired: process.kill() process.wait(timeout=2) def _receive_worker_response(control: socket.socket) -> dict[str, object]: chunks: list[bytes] = [] size = 0 while True: chunk = control.recv(65536) if not chunk: break chunks.append(chunk) size += len(chunk) if size > 1_000_000: raise ValueError("Visualization worker response exceeded its fixed boundary") if b"\n" in chunk: break payload = b"".join(chunks).split(b"\n", 1)[0] result = json.loads(payload) if not isinstance(result, dict): raise ValueError("Visualization worker returned an invalid response") return result def _one(params: dict[str, list[str]], name: str) -> str: values = params.get(name) or [""] return values[0] def _integer(value: str) -> int: return int(value) def _node_dict(row: sqlite3.Row, *, include_content: bool = True) -> dict[str, object]: result = { "node_id": row["node_id"], "title": row["title"], "family": row["family"], "authority": row["authority"], "status": row["status"], "tags": json.loads(row["tags_json"]), "summary": row["summary"], "source_path": row["source_path"], "source_anchor": row["source_anchor"], "content_hash": row["content_hash"], } if include_content: result["content"] = row["content"] return result def _facet_rows(connection: sqlite3.Connection, table: str, column: str) -> list[dict[str, object]]: allowed = { ("nodes", "family"), ("nodes", "authority"), ("nodes", "status"), ("edges", "relation"), } if (table, column) not in allowed: raise DocForgeError("invalid_index", "Unsupported visualization facet") rows = connection.execute( f"SELECT {column}, COUNT(*) AS count FROM {table} " f"GROUP BY {column} ORDER BY count DESC, {column}" ).fetchall() return [{"value": row[0], "count": row[1]} for row in rows] _GRAPH_BROWSER_HTML = r""" DocForge graph

DocForge graph

nodes edges
Loading validated index snapshot…
100%
Search for a node to inspect its neighborhood.
Click node to inspect · mouse wheel to zoom · left-drag to pan
Inspect node
"""