pretensor 0.1.0__py3-none-any.whl

pretensor/mcp/service_context.py

@@ -0,0 +1,142 @@
+"""Per-process MCP server context (visibility filter bound to ``graph_dir``).
+
+The active context is stored in a :class:`contextvars.ContextVar` rather than a
+bare mutable module global. In the current stdio server it is set once at
+startup, so every async task observes the same value; the ContextVar's benefit
+is removing ``global`` mutation and giving each thread / async context an
+independent binding, which future-proofs the context for a concurrent transport
+(HTTP/SSE) without reintroducing a shared-mutable-global hazard.
+"""
+
+from __future__ import annotations
+
+import contextvars
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from pretensor.config import GraphConfig, PretensorConfig
+from pretensor.mcp.store_cache import StoreCache
+from pretensor.search.base import BaseSearchIndex
+from pretensor.search.index import KeywordSearchIndex
+from pretensor.visibility.filter import VisibilityFilter
+from pretensor.visibility.runtime import load_visibility_filter_for_graph_dir
+
+__all__ = [
+    "DEFAULT_GRAPH_CONFIG",
+    "ServerContext",
+    "build_server_context",
+    "get_effective_graph_config",
+    "get_effective_search_index_cls",
+    "get_effective_visibility_filter",
+    "get_server_context",
+    "reset_server_context",
+    "set_server_context",
+]
+
+DEFAULT_GRAPH_CONFIG = GraphConfig()
+
+
+@dataclass(frozen=True, slots=True)
+class ServerContext:
+    """Holds graph workspace path and merged visibility rules for tool handlers."""
+
+    graph_dir: Path
+    visibility_filter: VisibilityFilter
+    config: PretensorConfig = field(default_factory=PretensorConfig)
+    store_cache: StoreCache = field(default_factory=StoreCache)
+
+    @property
+    def graph_config(self) -> GraphConfig:
+        """Backward-compatible access to graph tuning config."""
+        return self.config.graph
+
+    @property
+    def search_index_cls(self) -> type[BaseSearchIndex]:
+        """Backward-compatible access to configured search index class."""
+        return self.config.search_index_cls
+
+
+_ctx: contextvars.ContextVar[ServerContext | None] = contextvars.ContextVar(
+    "pretensor_mcp_server_context", default=None
+)
+
+
+def set_server_context(ctx: ServerContext) -> None:
+    """Install context for the current MCP server process (stdio)."""
+    _ctx.set(ctx)
+
+
+def reset_server_context() -> None:
+    """Clear the active server context.
+
+    Closes any open stores in the cache before clearing. Intended for test
+    teardown only — production servers never need to call this.
+    """
+    ctx = _ctx.get()
+    if ctx is not None:
+        ctx.store_cache.close()
+    _ctx.set(None)
+
+
+def get_server_context() -> ServerContext:
+    """Return the active server context (must be set before handling tools)."""
+    val = _ctx.get()
+    if val is None:
+        raise RuntimeError("MCP server context not initialized")
+    return val
+
+
+def build_server_context(
+    graph_dir: Path,
+    *,
+    visibility_path: Path | None = None,
+    profile: str | None = None,
+    config: PretensorConfig | None = None,
+) -> ServerContext:
+    """Resolve ``graph_dir`` and load visibility for one server lifetime."""
+    resolved = graph_dir.resolve()
+    vf = load_visibility_filter_for_graph_dir(
+        resolved,
+        visibility_path=visibility_path,
+        profile=profile,
+    )
+    effective_config = config or PretensorConfig()
+    return ServerContext(
+        graph_dir=resolved,
+        visibility_filter=vf,
+        config=effective_config,
+    )
+
+
+def get_effective_visibility_filter(
+    explicit: VisibilityFilter | None = None,
+) -> VisibilityFilter | None:
+    """Return an explicit filter, the server context filter, or None (no MCP server)."""
+    if explicit is not None:
+        return explicit
+    try:
+        return get_server_context().visibility_filter
+    except RuntimeError:
+        return None
+
+
+def get_effective_graph_config(explicit: GraphConfig | None = None) -> GraphConfig:
+    """Return explicit graph config or server context fallback."""
+    if explicit is not None:
+        return explicit
+    try:
+        return get_server_context().config.graph
+    except RuntimeError:
+        return DEFAULT_GRAPH_CONFIG
+
+
+def get_effective_search_index_cls(
+    explicit: type[BaseSearchIndex] | None = None,
+) -> type[BaseSearchIndex]:
+    """Return an explicit class, context class, or the OSS default."""
+    if explicit is not None:
+        return explicit
+    try:
+        return get_server_context().config.search_index_cls
+    except RuntimeError:
+        return KeywordSearchIndex

pretensor/mcp/service_registry.py

@@ -0,0 +1,294 @@
+"""Registry, keystore, and Kuzu open helpers for MCP tools."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Literal, NamedTuple
+
+from pretensor.core.dsn_crypto import DSNEncryptor
+from pretensor.core.registry import GraphRegistry, RegistryEntry
+from pretensor.core.store import KuzuStore
+from pretensor.mcp.service_context import get_server_context
+from pretensor.visibility.filter import VisibilityFilter
+
+CapabilityState = Literal["not_attempted", "empty", "present"]
+
+
+class GraphCounts(NamedTuple):
+    """Per-database counts and capability flags derived from the graph in one pass."""
+
+    table_count: int
+    column_count: int
+    row_count: int
+    schemas: list[str]
+    has_dbt_manifest: CapabilityState
+    has_llm_enrichment: CapabilityState
+    has_external_consumers: CapabilityState
+
+
+logger = logging.getLogger(__name__)
+
+
+def load_registry(graph_dir: Path) -> GraphRegistry:
+    path = graph_dir / "registry.json"
+    if not path.exists():
+        return GraphRegistry(path)
+    try:
+        return GraphRegistry(path).load()
+    except (json.JSONDecodeError, KeyError, TypeError) as exc:
+        logger.error("Registry file %s is corrupt: %s", path, exc)
+        return GraphRegistry(path)
+    except Exception as exc:
+        logger.warning("Failed to load registry from %s: %s", path, exc)
+        return GraphRegistry(path)
+
+
+def resolve_registry_entry(
+    registry: GraphRegistry, db: str | None
+) -> RegistryEntry | None:
+    """Resolve ``db`` by ``connection_name`` or logical ``database``; None picks a unique entry."""
+    entries = registry.list_entries()
+    if not entries:
+        return None
+    if db is None:
+        if len(entries) == 1:
+            return entries[0]
+        return None
+    for entry in entries:
+        if entry.connection_name == db or entry.database == db:
+            return entry
+    return None
+
+
+def keystore_path(graph_dir: Path) -> Path:
+    return graph_dir / "keystore"
+
+
+def encryptor_for_dir(graph_dir: Path) -> DSNEncryptor:
+    return DSNEncryptor(keystore_path(graph_dir))
+
+
+def entry_plaintext_dsn(entry: RegistryEntry, graph_dir: Path) -> str:
+    """Decrypt registry DSN when stored encrypted."""
+    return entry.plaintext_dsn(
+        encryptor_for_dir(graph_dir) if entry.dsn_encrypted else None
+    )
+
+
+def graph_path_for_entry(entry: RegistryEntry) -> Path:
+    return Path(entry.unified_graph_path or entry.graph_path)
+
+
+def _active_store_cache():
+    try:
+        return get_server_context().store_cache
+    except RuntimeError:
+        return None
+
+
+def _open_uncaught(graph_path: Path) -> KuzuStore:
+    """Open (or fetch from the active cache) a store, letting any error propagate raw.
+
+    The public ``open_store_*`` helpers wrap the raw error in a ``RuntimeError``
+    with a caller-appropriate message (path-based vs connection-based).
+    """
+    cache = _active_store_cache()
+    if cache is not None:
+        return cache.get(graph_path)
+    store = KuzuStore(graph_path)
+    store.ensure_schema()
+    return store
+
+
+def open_store_for_path(graph_path: Path) -> KuzuStore:
+    """Return a cached store for this path when a server context is active, else a fresh one."""
+    try:
+        return _open_uncaught(graph_path)
+    except Exception as exc:
+        raise RuntimeError(
+            f"Cannot open graph file ({graph_path}): {exc}. "
+            f"The file may be corrupt — re-run `pretensor index`."
+        ) from exc
+
+
+def open_store_for_entry(entry: RegistryEntry) -> KuzuStore:
+    gp = graph_path_for_entry(entry)
+    try:
+        return _open_uncaught(gp)
+    except Exception as exc:
+        # Reference the connection by name only — the on-disk graph path and the
+        # underlying engine error are kept off the message so they cannot leak
+        # to a semi-trusted MCP client. The original exception is preserved via
+        # ``from exc`` for server-side logs.
+        raise RuntimeError(
+            f"Cannot open graph file for connection {entry.connection_name!r}. "
+            "The graph may be missing or corrupt — re-run indexing."
+        ) from exc
+
+
+def release_store(store: KuzuStore) -> None:
+    """Close a store unless it is owned by the active server-context cache (pooled)."""
+    cache = _active_store_cache()
+    if cache is not None and cache.owns(store):
+        return
+    store.close()
+
+
+def _column_count_for_connection(
+    store: KuzuStore,
+    *,
+    connection_name: str | None = None,
+) -> int:
+    """Return the number of ``SchemaColumn`` nodes, optionally scoped to a connection."""
+    if connection_name is None:
+        r = store.query_all_rows("MATCH (c:SchemaColumn) RETURN count(*)")
+    else:
+        r = store.query_all_rows(
+            "MATCH (c:SchemaColumn) WHERE c.connection_name = $cn RETURN count(*)",
+            {"cn": connection_name},
+        )
+    if not r:
+        return 0
+    val = r[0][0]
+    return int(val) if val is not None else 0
+
+
+_DBT_MODEL_DEPENDENCY_LINEAGE_TYPE = "model_dependency"
+
+
+def _has_dbt_manifest(store: KuzuStore, *, connection_name: str | None) -> bool:
+    """True when any ``LINEAGE`` edge with ``lineage_type='model_dependency'`` exists.
+
+    dbt's ``manifest.json`` is not stored as a graph node; its ``parent_map`` is
+    projected as ``LINEAGE`` edges tagged with this ``lineage_type`` (see
+    ``src/pretensor/enrichment/dbt/lineage.py``).
+    """
+    if connection_name is None:
+        rows = store.query_all_rows(
+            """
+            MATCH ()-[r:LINEAGE]->()
+            WHERE r.lineage_type = $lt
+            RETURN 1 LIMIT 1
+            """,
+            {"lt": _DBT_MODEL_DEPENDENCY_LINEAGE_TYPE},
+        )
+    else:
+        rows = store.query_all_rows(
+            """
+            MATCH (s:SchemaTable)-[r:LINEAGE]->()
+            WHERE s.connection_name = $cn AND r.lineage_type = $lt
+            RETURN 1 LIMIT 1
+            """,
+            {"cn": connection_name, "lt": _DBT_MODEL_DEPENDENCY_LINEAGE_TYPE},
+        )
+    return len(rows) > 0
+
+
+def _has_llm_enrichment(store: KuzuStore, *, connection_name: str | None) -> bool:
+    """True when any ``Entity`` node exists (LLM enrichment output)."""
+    if connection_name is None:
+        rows = store.query_all_rows("MATCH (e:Entity) RETURN 1 LIMIT 1")
+    else:
+        rows = store.query_all_rows(
+            "MATCH (e:Entity) WHERE e.connection_name = $cn RETURN 1 LIMIT 1",
+            {"cn": connection_name},
+        )
+    return len(rows) > 0
+
+
+def counts_for_graph(
+    store: KuzuStore,
+    *,
+    connection_name: str | None = None,
+    visibility_filter: VisibilityFilter | None = None,
+    entry: RegistryEntry | None = None,
+) -> GraphCounts:
+    """Return per-database counts and capability flags in a single store pass.
+
+    Optionally scoped to one connection. Visibility filtering applies to table
+    counts, ``schemas``, and ``has_external_consumers`` (hidden tables do not
+    contribute); it does not affect the dbt/Entity probes, which are graph-wide.
+    """
+    params = {"cn": connection_name} if connection_name is not None else {}
+    cn_clause = "WHERE t.connection_name = $cn" if connection_name is not None else ""
+    rows = store.query_all_rows(
+        f"""
+        MATCH (t:SchemaTable)
+        {cn_clause}
+        RETURN t.connection_name, t.schema_name, t.table_name, t.row_count,
+               t.has_external_consumers
+        """,
+        params,
+    )
+    total = 0
+    row_sum = 0
+    schemas: set[str] = set()
+    has_external = False
+    for cn, sn, tn, rc, hec in rows:
+        cns = str(cn or "")
+        sns = str(sn or "")
+        tns = str(tn or "")
+        if visibility_filter is not None and not visibility_filter.is_table_visible(
+            cns, sns, tns
+        ):
+            continue
+        total += 1
+        if rc is not None:
+            row_sum += int(rc)
+        if sns:
+            schemas.add(sns)
+        if bool(hec):
+            has_external = True
+    cc = _column_count_for_connection(store, connection_name=connection_name)
+    dbt_present = _has_dbt_manifest(store, connection_name=connection_name)
+    llm_present = _has_llm_enrichment(store, connection_name=connection_name)
+    # Trinary states use entry provenance when available: "not_attempted" means
+    # the user never ran that enrichment pass (so callers shouldn't infer absence
+    # of signal). An indexer that did attempt it returns "empty" vs "present".
+    dbt_attempted = entry is not None and entry.dbt_manifest_path is not None
+    if not dbt_attempted and entry is not None:
+        has_dbt: CapabilityState = "not_attempted"
+    else:
+        has_dbt = (
+            "present"
+            if dbt_present
+            else ("empty" if dbt_attempted else "not_attempted")
+        )
+    llm_attempted = entry is not None and entry.llm_enrichment_ran
+    if entry is None:
+        has_llm: CapabilityState = "present" if llm_present else "empty"
+    else:
+        has_llm = (
+            "present"
+            if llm_present
+            else ("empty" if llm_attempted else "not_attempted")
+        )
+    # External-consumer scanner is not shipped yet; unify the "never checked"
+    # default with the dbt/LLM flags so callers don't have to special-case it.
+    has_ext: CapabilityState = "present" if has_external else "not_attempted"
+    return GraphCounts(
+        table_count=total,
+        column_count=cc,
+        row_count=row_sum,
+        schemas=sorted(schemas),
+        has_dbt_manifest=has_dbt,
+        has_llm_enrichment=has_llm,
+        has_external_consumers=has_ext,
+    )
+
+
+__all__ = [
+    "GraphCounts",
+    "counts_for_graph",
+    "encryptor_for_dir",
+    "entry_plaintext_dsn",
+    "graph_path_for_entry",
+    "keystore_path",
+    "load_registry",
+    "open_store_for_entry",
+    "open_store_for_path",
+    "release_store",
+    "resolve_registry_entry",
+]

pretensor/mcp/store_cache.py

@@ -0,0 +1,43 @@
+"""Per-server-lifetime cache of open KuzuStore handles, keyed by resolved graph path.
+
+The MCP server is single-threaded (stdio + one asyncio loop), so a single reused
+KuzuStore (one kuzu.Connection) per graph path is safe and avoids re-opening the
+database + re-running ensure_schema() on every tool call. Cross-process staleness
+caveat: a separate `pretensor reindex` process writing the same file is not
+reflected in an already-open cached read handle until the server restarts — this
+matches the previous behavior's freshness boundary for the server's lifetime.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from pretensor.core.store import KuzuStore
+
+logger = logging.getLogger(__name__)
+
+
+class StoreCache:
+    def __init__(self) -> None:
+        self._stores: dict[Path, KuzuStore] = {}
+
+    def get(self, graph_path: Path) -> KuzuStore:
+        key = graph_path.resolve()
+        store = self._stores.get(key)
+        if store is None:
+            store = KuzuStore(key)
+            store.ensure_schema()
+            self._stores[key] = store
+        return store
+
+    def owns(self, store: KuzuStore) -> bool:
+        return any(s is store for s in self._stores.values())
+
+    def close(self) -> None:
+        for path, store in self._stores.items():
+            try:
+                store.close()
+            except Exception:
+                logger.warning("Error closing cached store for %s", path, exc_info=True)
+        self._stores.clear()

pretensor/mcp/tool_registry.py

@@ -0,0 +1,136 @@
+"""MCP tool plugin registry — McpTool and McpToolRegistry.
+
+Allows dynamic, pluggable registration of MCP tools so that extensions
+can add tools (e.g. ``semantic_search``) without modifying
+OSS server code.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+import mcp.types as types
+
+__all__ = ["McpTool", "McpToolRegistry"]
+
+logger = logging.getLogger(__name__)
+
+ToolHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
+
+
+@dataclass
+class McpTool:
+    """Descriptor and handler for a single MCP tool.
+
+    Args:
+        name: Unique tool name exposed to the MCP client.
+        description: Human-readable description shown to the client.
+        input_schema: JSON Schema object describing the tool's input.
+        handler: Async callable that receives ``arguments`` and returns a
+            JSON-serialisable ``dict``.
+    """
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+    handler: ToolHandler
+
+    def to_mcp_type(self) -> types.Tool:
+        """Return the :class:`mcp.types.Tool` representation."""
+        return types.Tool(
+            name=self.name,
+            description=self.description,
+            inputSchema=self.input_schema,
+        )
+
+
+class McpToolRegistry:
+    """Registry that owns the canonical set of MCP tools for a server.
+
+    Usage::
+
+        registry = McpToolRegistry()
+        registry.register(McpTool(name="my_tool", ...))
+
+        # In the MCP list_tools handler:
+        tools = registry.list_tools()
+
+        # In the MCP call_tool handler:
+        result = await registry.call_tool("my_tool", {"arg": "value"})
+    """
+
+    def __init__(self) -> None:
+        self._tools: dict[str, McpTool] = {}
+
+    def register(self, tool: McpTool) -> None:
+        """Register a tool; raises :exc:`ValueError` if the name is already taken.
+
+        Args:
+            tool: The tool descriptor and handler to register.
+
+        Raises:
+            ValueError: If a tool with the same name is already registered.
+        """
+        if tool.name in self._tools:
+            raise ValueError(f"Tool '{tool.name}' is already registered")
+        self._tools[tool.name] = tool
+        logger.debug("Registered MCP tool: %s", tool.name)
+
+    def list_tools(self) -> list[types.Tool]:
+        """Return all registered tools as :class:`mcp.types.Tool` objects.
+
+        Returns:
+            List of MCP tool descriptors in registration order.
+        """
+        return [t.to_mcp_type() for t in self._tools.values()]
+
+    async def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """Dispatch a tool call to its registered handler.
+
+        All handler exceptions are caught and returned as an opaque error
+        envelope so the MCP server never crashes on a bad tool invocation and
+        never discloses internals (tracebacks, filesystem paths, exception
+        text) to the semi-trusted client. The full stack is kept server-side
+        via ``logger.exception`` tagged with the same ``correlation_id`` so an
+        operator can still join a client-reported id to the server log.
+
+        Args:
+            name: Tool name as sent by the MCP client.
+            arguments: Raw argument dict (may be ``None`` or empty).
+
+        Returns:
+            JSON-serialisable result dict from the handler, or an opaque
+            ``{"error": "Internal tool error", "tool": ..., "correlation_id":
+            ...}`` envelope when the handler raises.
+        """
+        tool = self._tools.get(name)
+        if tool is None:
+            return {"error": f"Unknown tool: {name}"}
+        args = arguments or {}
+        try:
+            return await tool.handler(args)
+        except Exception:
+            correlation_id = uuid.uuid4().hex
+            logger.exception(
+                "Unhandled exception in MCP tool %r [correlation_id=%s]",
+                name,
+                correlation_id,
+            )
+            return {
+                "error": "Internal tool error",
+                "tool": name,
+                "correlation_id": correlation_id,
+            }
+
+    def __len__(self) -> int:
+        return len(self._tools)
+
+    def __contains__(self, name: object) -> bool:
+        return name in self._tools

pretensor/mcp/tools/__init__.py

@@ -0,0 +1 @@
+"""MCP tool payload implementations (list, query, context, traverse, impact, detect_changes)."""