toolgraph 0.0.1__py3-none-any.whl

toolgraph/crawler/client.py

@@ -0,0 +1,120 @@
+"""Enumerate one MCP server's factual capabilities into a CrawlResult."""
+
+from __future__ import annotations
+
+from toolgraph.crawler.transports import open_session
+from toolgraph.models import CrawlResult, ResourceRecord, ServerSpec, ToolRecord
+from toolgraph.redaction import endpoint_label
+from toolgraph.uris import normalize_resource_uri
+
+# A crawled server controls its own cursors, so an unbounded drain loop lets a
+# buggy or malicious server that keeps returning a cursor pin the crawler and
+# grow memory without limit. The per-server wall-clock timeout only applies on
+# the CLI path (crawl_all), not to library callers of crawl_server. Real
+# catalogs are a handful of pages; hitting this cap means the server is broken.
+MAX_LIST_PAGES = 1000
+
+
+class PaginationLimitError(RuntimeError):
+    """A server kept returning continuation cursors past MAX_LIST_PAGES."""
+
+
+def _annotation_fields(tool) -> dict:
+    """Tool annotations as ToolRecord fields (ADR-0006).
+
+    Annotations are optional in the protocol; absent ones stay None so the
+    loader can clear node properties the server stopped sending (SET to null
+    removes the property — same reconciliation lifecycle as description).
+    """
+    ann = getattr(tool, "annotations", None)
+    if ann is None:
+        return {}
+    return {
+        "read_only_hint": ann.readOnlyHint,
+        "destructive_hint": ann.destructiveHint,
+        "idempotent_hint": ann.idempotentHint,
+        "open_world_hint": ann.openWorldHint,
+    }
+
+
+async def _list_all_tools(session) -> list[ToolRecord]:
+    """Drain ``list_tools`` pages until ``nextCursor`` is exhausted.
+
+    Without pagination, a server with more than one page silently drops
+    tools on later pages — governance refs to them would then reject at
+    ingest and reachability analysis would miss the tools entirely (Codex
+    PR #1 review caught this).
+    """
+    out: list[ToolRecord] = []
+    cursor: str | None = None
+    for _ in range(MAX_LIST_PAGES):
+        listed = await session.list_tools(cursor=cursor)
+        out.extend(
+            ToolRecord(
+                name=t.name,
+                description=t.description,
+                input_schema=t.inputSchema,
+                **_annotation_fields(t),
+            )
+            for t in listed.tools
+        )
+        cursor = getattr(listed, "nextCursor", None)
+        if cursor is None:
+            # Stop ONLY on None: an empty-string cursor is a valid opaque
+            # continuation token in MCP and the SDK distinguishes "" from
+            # None when sending the next request. `if not cursor` would
+            # silently truncate after page 1 against such a server.
+            return out
+    raise PaginationLimitError(
+        f"list_tools returned a continuation cursor on {MAX_LIST_PAGES}"
+        " consecutive pages — refusing to drain further"
+    )
+
+
+async def _list_all_resources(session) -> list[ResourceRecord]:
+    """Drain ``list_resources`` pages until ``nextCursor`` is exhausted."""
+    out: list[ResourceRecord] = []
+    cursor: str | None = None
+    for _ in range(MAX_LIST_PAGES):
+        listed = await session.list_resources(cursor=cursor)
+        out.extend(
+            ResourceRecord(
+                uri=normalize_resource_uri(str(r.uri)),
+                name=r.name,
+                mime_type=r.mimeType,
+                description=r.description,
+            )
+            for r in listed.resources
+        )
+        cursor = getattr(listed, "nextCursor", None)
+        if cursor is None:
+            # Stop ONLY on None: an empty-string cursor is a valid opaque
+            # continuation token in MCP and the SDK distinguishes "" from
+            # None when sending the next request. `if not cursor` would
+            # silently truncate after page 1 against such a server.
+            return out
+    raise PaginationLimitError(
+        f"list_resources returned a continuation cursor on {MAX_LIST_PAGES}"
+        " consecutive pages — refusing to drain further"
+    )
+
+
+async def crawl_server(spec: ServerSpec) -> CrawlResult:
+    """Connect, initialize, and list tools/resources the server actually advertises."""
+    async with open_session(spec) as session:
+        init = await session.initialize()
+        caps = init.capabilities
+
+        tools = await _list_all_tools(session) if caps.tools is not None else []
+        resources = (
+            await _list_all_resources(session) if caps.resources is not None else []
+        )
+
+        return CrawlResult(
+            server_name=spec.name or init.serverInfo.name,
+            server_version=init.serverInfo.version,
+            transport=spec.transport,
+            endpoint=endpoint_label(spec),
+            tools=tools,
+            resources=resources,
+        )

toolgraph/crawler/crawl.py

@@ -0,0 +1,103 @@
+"""Crawl many servers; each server fails independently (partial success)."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from toolgraph.crawler.client import crawl_server
+from toolgraph.models import CrawlResult, ServerSpec, ServersConfig
+from toolgraph.redaction import endpoint_label, redact_text
+from toolgraph.yaml_loader import load_yaml_mapping
+
+
+def load_servers_config(path: Path) -> list[ServerSpec]:
+    data = load_yaml_mapping(path)
+    return ServersConfig.model_validate(data).servers
+
+
+def spec_label(spec: ServerSpec) -> str:
+    """Display label for a spec — NOT a graph identity (see fleet_keep_names)."""
+    return spec.name or endpoint_label(spec)
+
+
+DEFAULT_TIMEOUT = 30.0
+DEFAULT_CONCURRENCY = 5
+
+# A failure carries its ServerSpec, not a display label: fleet reconciliation
+# must know whether the failed spec had a ``name`` override, and inferring
+# that from a label is unsound (an unnamed spec's command can collide with
+# another spec's name — Codex review of this branch caught it).
+CrawlFailure = tuple[ServerSpec, str]
+
+
+async def _crawl_one(
+    spec: ServerSpec, timeout: float | None, sem: asyncio.Semaphore
+) -> tuple[CrawlResult | None, CrawlFailure | None]:
+    async with sem:
+        try:
+            return await asyncio.wait_for(crawl_server(spec), timeout), None
+        except TimeoutError:  # a hanging server must not block the crawl
+            return None, (spec, f"timeout after {timeout}s")
+        except Exception as exc:  # noqa: BLE001 - one bad server must not abort the crawl
+            return None, (spec, f"{type(exc).__name__}: {redact_text(exc)}")
+
+
+async def crawl_all_async(
+    specs: list[ServerSpec],
+    timeout: float | None = DEFAULT_TIMEOUT,
+    concurrency: int = DEFAULT_CONCURRENCY,
+) -> tuple[list[CrawlResult], list[CrawlFailure]]:
+    sem = asyncio.Semaphore(concurrency)
+    outcomes = await asyncio.gather(*(_crawl_one(s, timeout, sem) for s in specs))
+    results = [r for r, _ in outcomes if r is not None]
+    failures = [f for _, f in outcomes if f is not None]
+    return results, failures
+
+
+def crawl_all(
+    specs: list[ServerSpec],
+    timeout: float | None = DEFAULT_TIMEOUT,
+    concurrency: int = DEFAULT_CONCURRENCY,
+) -> tuple[list[CrawlResult], list[CrawlFailure]]:
+    return asyncio.run(crawl_all_async(specs, timeout, concurrency))
+
+
+def fleet_keep_names(
+    specs: list[ServerSpec],
+    results: list[CrawlResult],
+    failures: list[CrawlFailure],
+) -> tuple[set[str], list[str]]:
+    """Graph names that must survive fleet reconciliation, plus blockers.
+
+    A failed spec with a ``name`` override is protected by that name. A failed
+    spec WITHOUT one cannot be mapped to a graph identity (node names come
+    from ``serverInfo.name`` at crawl time, which a failed crawl never saw) —
+    its label is returned as a blocker and the caller must skip the prune:
+    a server that failed to crawl is not a server that was removed.
+    """
+    keep = {s.name for s in specs if s.name} | {r.server_name for r in results}
+    blockers = [spec_label(spec) for spec, _ in failures if not spec.name]
+    return keep, blockers
+
+
+def duplicate_name_overrides(specs: list[ServerSpec]) -> set[str]:
+    """``name:`` overrides claimed by more than one spec — a config error."""
+    names = [s.name for s in specs if s.name]
+    return {n for n in names if names.count(n) > 1}
+
+
+def duplicate_identities(
+    results: list[CrawlResult], failures: list[CrawlFailure]
+) -> set[str]:
+    """Graph names claimed by more than one server this pass.
+
+    Two specs resolving to the same ``server_name`` would silently MERGE into
+    one MCPServer node and last-load-wins reconcile each other's tools away —
+    and give fleet reconciliation an ambiguous identity. The caller must
+    reject the pass before any graph mutation. Failed named specs count too:
+    an unnamed server self-reporting a failed spec's name is the same
+    collision, one crawl later.
+    """
+    names = [r.server_name for r in results] + [s.name for s, _ in failures if s.name]
+    return {n for n in names if names.count(n) > 1}

toolgraph/crawler/transports.py

@@ -0,0 +1,60 @@
+"""Open an MCP ClientSession over the right transport for a ServerSpec.
+
+Normalizes the transport differences the SDK exposes: stdio yields a 2-tuple
+``(read, write)`` while streamable-http yields a 3-tuple
+``(read, write, get_session_id)``. SSE is supported only as a deprecated fallback.
+
+Uses the preferred ``streamable_http_client``; since it takes headers via a
+custom httpx client rather than a ``headers=`` kwarg, we build one with the
+SDK's ``create_mcp_http_client`` when auth headers are supplied.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import AsyncExitStack, asynccontextmanager
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.sse import sse_client
+from mcp.client.stdio import get_default_environment, stdio_client
+from mcp.client.streamable_http import create_mcp_http_client, streamable_http_client
+
+from toolgraph.models import ServerSpec
+
+
+@asynccontextmanager
+async def open_session(spec: ServerSpec) -> AsyncIterator[ClientSession]:
+    async with AsyncExitStack() as stack:
+        if spec.transport == "stdio":
+            if not spec.command:
+                raise ValueError(f"stdio server {spec.name!r} needs a 'command'")
+            # Merge over the SDK's safe default env so PATH etc. survive when the
+            # user supplies only secrets (e.g. API keys).
+            env = {**get_default_environment(), **spec.env} if spec.env else None
+            params = StdioServerParameters(command=spec.command, args=spec.args, env=env)
+            read, write = await stack.enter_async_context(stdio_client(params))
+
+        elif spec.transport == "streamable-http":
+            if not spec.url:
+                raise ValueError(f"streamable-http server {spec.name!r} needs a 'url'")
+            http_client = None
+            if spec.headers:
+                http_client = await stack.enter_async_context(
+                    create_mcp_http_client(headers=spec.headers)
+                )
+            read, write, _get_session_id = await stack.enter_async_context(
+                streamable_http_client(spec.url, http_client=http_client)
+            )
+
+        elif spec.transport == "sse":
+            if not spec.url:
+                raise ValueError(f"sse server {spec.name!r} needs a 'url'")
+            read, write = await stack.enter_async_context(
+                sse_client(spec.url, headers=spec.headers or None)
+            )
+
+        else:  # pragma: no cover - guarded by the Literal type
+            raise ValueError(f"unknown transport {spec.transport!r}")
+
+        session = await stack.enter_async_context(ClientSession(read, write))
+        yield session

toolgraph/graph/driver.py

@@ -0,0 +1,275 @@
+"""Backend lifecycle and transaction/session compatibility helpers.
+
+The public module seam is intentionally small while graph services migrate to
+domain-level ports. Neo4j remains the compatibility default; Ladybug is an
+explicit embedded backend and never runs in the gateway call path.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+import re
+import threading
+from typing import Any
+
+from neo4j import Driver, GraphDatabase, Session
+from neo4j.exceptions import (
+    ConnectionAcquisitionTimeoutError,
+    DatabaseUnavailable,
+    ServiceUnavailable,
+    SessionExpired,
+)
+
+from toolgraph import config
+
+_driver: Driver | None = None
+_ladybug_database: Any | None = None
+# Serializes lazy backend construction. The streamable-HTTP MCP server can run
+# two first-requests concurrently; an unguarded check-then-set would leak one
+# Neo4j driver, or open a second Ladybug Database that trips the file lock and
+# misreports a spurious BACKEND_LOCKED.
+_backend_lock = threading.Lock()
+
+
+class BackendConfigurationError(RuntimeError):
+    pass
+
+
+class BackendUnavailableError(RuntimeError):
+    """A known transient backend outage, raised at the backend boundary.
+
+    Callers above the driver classify on this type alone — the native
+    neo4j/Ladybug taxonomy stays private to this module.
+    """
+
+
+class BackendLockedError(BackendUnavailableError):
+    pass
+
+
+_NATIVE_UNAVAILABLE_ERRORS = (
+    ServiceUnavailable,
+    SessionExpired,
+    ConnectionAcquisitionTimeoutError,
+    DatabaseUnavailable,
+)
+
+
+@contextmanager
+def _translating_backend_errors() -> Iterator[None]:
+    """Re-raise the narrow native outage taxonomy as ``BackendUnavailableError``.
+
+    Deliberately narrower than "any backend exception": invalid configuration,
+    authentication, and Cypher/client errors must stay loud contract errors.
+    """
+    try:
+        yield
+    except _NATIVE_UNAVAILABLE_ERRORS as exc:
+        raise BackendUnavailableError(str(exc)) from exc
+
+
+def is_backend_unavailable(exc: BaseException) -> bool:
+    """Whether *exc* was caused by a backend outage typed at the driver seam.
+
+    FastMCP wraps tool exceptions, so walk the chain. ``__context__`` is
+    followed only when Python did not suppress it: a deliberate
+    ``raise ... from None`` while handling an outage is the author saying the
+    outer error is the real one, and must not be reclassified as retryable.
+    """
+    current: BaseException | None = exc
+    seen: set[int] = set()
+    while current is not None and id(current) not in seen:
+        if isinstance(current, BackendUnavailableError):
+            return True
+        seen.add(id(current))
+        if current.__cause__ is not None:
+            current = current.__cause__
+        elif not current.__suppress_context__:
+            current = current.__context__
+        else:
+            current = None
+    return False
+
+
+class _LadybugResult:
+    def __init__(self, result: Any) -> None:
+        self._rows: list[dict[str, Any]] = result.rows_as_dict().get_all()
+
+    def __iter__(self):
+        return iter(self._rows)
+
+    def single(self) -> dict[str, Any] | None:
+        if not self._rows:
+            return None
+        if len(self._rows) > 1:
+            raise RuntimeError(f"expected at most one row, received {len(self._rows)}")
+        return self._rows[0]
+
+
+_IDENTIFIER = r"[A-Za-z_][A-Za-z0-9_]*"
+_DATETIME_CALL = re.compile(r"(?<![A-Za-z0-9_])datetime\s*\(\s*\)")
+_TYPE_CALL = re.compile(rf"(?<![A-Za-z0-9_])type\s*\(\s*({_IDENTIFIER})\s*\)")
+_LABELS_HEAD = re.compile(
+    rf"(?<![A-Za-z0-9_])labels\s*\(\s*({_IDENTIFIER})\s*\)\s*\[\s*0\s*\]"
+)
+_UNSUPPORTED_NEO4J_EXPRESSION = re.compile(
+    r"(?<![A-Za-z0-9_])(?:datetime|type|labels)\s*\("
+)
+
+
+def _ladybug_query(query: str) -> str:
+    """Translate only the explicitly supported Neo4j expression shapes.
+
+    This is intentionally a narrow dialect contract, not a Cypher parser.
+    Identifier boundaries prevent incidental substrings from being rewritten;
+    a remaining Neo4j-only function fails fast instead of reaching Ladybug as
+    a subtly different query.
+    """
+    translated = _DATETIME_CALL.sub("current_timestamp()", query)
+    translated = _TYPE_CALL.sub(r"label(\1)", translated)
+    translated = _LABELS_HEAD.sub(r"label(\1)", translated)
+    for field in (
+        "read_only_hint",
+        "destructive_hint",
+        "idempotent_hint",
+        "open_world_hint",
+    ):
+        translated = re.sub(
+            rf"\btool\.{field}\s*=\s*t\.{field}\b",
+            f"tool.{field}=CAST(t.{field} AS BOOLEAN)",
+            translated,
+        )
+    if match := _UNSUPPORTED_NEO4J_EXPRESSION.search(translated):
+        raise BackendConfigurationError(
+            f"unsupported Neo4j expression for Ladybug near {match.group(0)!r}"
+        )
+    return translated
+
+
+class LadybugSession:
+    """Minimal transaction facade used during the domain-port migration."""
+
+    def __init__(self, database: Any) -> None:
+        import ladybug as lb
+
+        self._connection = lb.Connection(database)
+
+    def run(self, query: str, **parameters: Any) -> _LadybugResult:
+        result = self._connection.execute(_ladybug_query(query), parameters)
+        return _LadybugResult(result)
+
+    def execute_write(self, fn, *args):
+        self._connection.execute("BEGIN TRANSACTION")
+        try:
+            value = fn(self, *args)
+            self._connection.execute("COMMIT")
+            return value
+        except BaseException:
+            # Ladybug aborts the explicit transaction immediately on some
+            # binder/runtime errors; in that case there is nothing left to
+            # roll back. Preserve the original exception.
+            try:
+                self._connection.execute("ROLLBACK")
+            except RuntimeError as rollback_error:
+                if "No active transaction" not in str(rollback_error):
+                    raise
+            raise
+
+    def close(self) -> None:
+        self._connection.close()
+
+
+def backend_name() -> str:
+    backend = config.settings.backend.lower().strip()
+    if backend not in {"neo4j", "ladybug"}:
+        raise BackendConfigurationError(
+            f"unknown TOOLGRAPH_BACKEND {backend!r}; expected 'neo4j' or 'ladybug'"
+        )
+    return backend
+
+
+def get_driver() -> Driver:
+    """Return the process-wide Neo4j driver."""
+    if backend_name() != "neo4j":
+        raise BackendConfigurationError("Neo4j driver requested while backend=ladybug")
+    global _driver
+    if _driver is None:
+        with _backend_lock:
+            if _driver is None:
+                _driver = GraphDatabase.driver(
+                    config.settings.neo4j_uri,
+                    auth=(config.settings.neo4j_user, config.settings.neo4j_password),
+                )
+    return _driver
+
+
+def _get_ladybug_database() -> Any:
+    global _ladybug_database
+    if _ladybug_database is not None:
+        return _ladybug_database
+    try:
+        import ladybug as lb
+    except ImportError as exc:  # pragma: no cover - environment dependent
+        raise BackendConfigurationError(
+            "backend=ladybug requires the 'ladybug' extra: "
+            "install with `uv sync --extra ladybug` or `pip install toolgraph[ladybug]`"
+        ) from exc
+    with _backend_lock:
+        if _ladybug_database is not None:
+            return _ladybug_database
+        path = Path(config.settings.db_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            _ladybug_database = lb.Database(path)
+        except RuntimeError as exc:
+            message = str(exc)
+            if "lock" in message.lower() or "another process" in message.lower():
+                raise BackendLockedError(
+                    f"BACKEND_LOCKED: Ladybug database {path} already has an owner"
+                ) from exc
+            raise
+        return _ladybug_database
+
+
+def verify_connectivity() -> None:
+    if backend_name() == "neo4j":
+        with _translating_backend_errors():
+            get_driver().verify_connectivity()
+        return
+    with session() as current:
+        current.run("RETURN 1 AS ok").single()
+
+
+@contextmanager
+def session() -> Iterator[Session | LadybugSession]:
+    """Yield a backend session, typing outage failures at this seam.
+
+    The translation wraps the ``yield`` too: a driver outage that surfaces
+    mid-query is thrown back in here by the caller's ``with`` block, so one
+    boundary covers acquisition and execution alike.
+    """
+    with _translating_backend_errors():
+        if backend_name() == "neo4j":
+            with get_driver().session(
+                database=config.settings.neo4j_database
+            ) as current:
+                yield current
+            return
+        current = LadybugSession(_get_ladybug_database())
+        try:
+            yield current
+        finally:
+            current.close()
+
+
+def close_driver() -> None:
+    global _driver, _ladybug_database
+    with _backend_lock:
+        if _driver is not None:
+            _driver.close()
+            _driver = None
+        if _ladybug_database is not None:
+            _ladybug_database.close()
+            _ladybug_database = None