administrator/DocForge2
1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions Pages

Detach graph viewer from MCP transport

Browse source
This commit is contained in:
Andraxion 2026-07-24 22:07:33 -04:00
parent 5e77cd2adb
commit f6b9816ffd
10 changed files with 369 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

185
src/docforge/visualization.py
View file

@ -4,14 +4,19 @@ 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
@ -52,6 +57,31 @@ class VisualizationIndexSnapshot:
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(
@ -270,20 +300,27 @@ class VisualizationRunner:
def __init__(
self,
index: ProjectIndex,
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 = secrets.token_urlsafe(24)
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
@ -305,13 +342,17 @@ class VisualizationRunner:
"invalid_visualization_target",
"Choose either one exact node ID or one search query",
)
maximum_depth = self.index.project.descriptor.limits.max_traversal_depth
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",
)
reader = VisualizationIndexSnapshot(self.index, self.index.check())
if node_id is not None:
reader.require_node(node_id)
elif query is not None:
@ -365,7 +406,7 @@ class VisualizationRunner:
daemon=True,
)
self._lease_thread.start()
if not self._atexit_registered:
if self._register_atexit and not self._atexit_registered:
atexit.register(self.stop)
self._atexit_registered = True
@ -419,6 +460,10 @@ class VisualizationRunner:
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:
@ -617,6 +662,136 @@ class VisualizationRunner:
)
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]