agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/serve/engine.py

@@ -0,0 +1,204 @@
+"""Process-lifetime engine holder behind the MCP tools.
+
+Opened lazily on first tool call so ``code_graph_tools(".")`` can be built
+synchronously and passed straight to ``Agent(tools=…)`` / the MCP server.
+This module (and the whole ``serve`` package) is the framework-facing layer —
+it may import ``agentforge`` (ADR-0001 exception); the engine packages do not.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from agentforge_graph.config import (
+    EmbedConfig,
+    RepoMapConfig,
+    RetrieveConfig,
+    ServeConfig,
+    StoreConfig,
+)
+from agentforge_graph.core import GraphQuery
+from agentforge_graph.embed import embedder_from_config
+from agentforge_graph.ingest import CodeGraph
+from agentforge_graph.repomap import RepoMap
+from agentforge_graph.retrieve import Retriever
+
+TOOL_API_VERSION = "1.0"
+_ALL = 10_000_000
+
+
+def _git_head(repo_path: str | Path) -> str:
+    try:
+        out = subprocess.run(
+            ["git", "-C", str(repo_path), "rev-parse", "HEAD"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return out.stdout.strip()
+    except (subprocess.SubprocessError, OSError):
+        return ""
+
+
+class _Engine:
+    def __init__(self, repo_path: str | Path = ".", config: str | Path | None = None) -> None:
+        self.repo_path = repo_path
+        self.config = config
+        self.serve = ServeConfig.load(config)
+        self._cg: CodeGraph | None = None
+        self._retriever: Retriever | None = None
+        self._repomap: RepoMap | None = None
+
+    async def code_graph(self) -> CodeGraph:
+        if self._cg is None:
+            self._cg = await CodeGraph.open(self.repo_path, self.config)
+        return self._cg
+
+    async def retriever(self) -> Retriever:
+        if self._retriever is None:
+            from agentforge_graph.retrieve.rerank import reranker_from_config
+
+            cg = await self.code_graph()
+            embedder = embedder_from_config(EmbedConfig.load(self.config))
+            rcfg = RetrieveConfig.load(self.config)
+            # ENH-009: honor retrieve.rerank over MCP too (previously ignored here).
+            self._retriever = Retriever(
+                cg.store,
+                embedder,
+                rcfg,
+                reranker=reranker_from_config(rcfg.rerank, rcfg.rerank_weight, rcfg.rerank_model),
+            )
+        return self._retriever
+
+    async def repomap(self) -> RepoMap:
+        if self._repomap is None:
+            cg = await self.code_graph()
+            self._repomap = RepoMap(cg.store, RepoMapConfig.load(self.config))
+        return self._repomap
+
+    def _meta(self) -> Any:
+        from agentforge_graph.ingest.incremental import IndexMeta
+
+        root = Path(self.repo_path) / StoreConfig.load(self.config).path
+        return IndexMeta.load(root)
+
+    async def staleness(self) -> dict[str, Any]:
+        """Cheap envelope: the indexed commit + dirty flag, read from the
+        persisted manifest (feat-004) rather than probed from a node."""
+        meta = self._meta()
+        head = _git_head(self.repo_path)
+        return {
+            "indexed_commit": meta.indexed_commit,
+            "dirty": bool(head) and bool(meta.indexed_commit) and head != meta.indexed_commit,
+        }
+
+    async def routes(self, method: str = "", path: str = "") -> dict[str, Any]:
+        """Extracted endpoints (feat-011), optionally filtered by HTTP method
+        and/or path prefix, wrapped in the staleness envelope."""
+        cg = await self.code_graph()
+        items = [r.to_dict() for r in await cg.routes()]
+        if method:
+            items = [r for r in items if str(r["method"]).upper() == method.upper()]
+        if path:
+            items = [r for r in items if str(r["path"]).startswith(path)]
+        return {
+            "routes": items,
+            "count": len(items),
+            **(await self.staleness()),
+            "tool_api_version": TOOL_API_VERSION,
+        }
+
+    async def decisions(self, scope: str = "", status: str = "") -> dict[str, Any]:
+        """Architecture decisions (feat-010), optionally filtered by governed
+        path ``scope`` and ``status``, wrapped in the staleness envelope."""
+        cg = await self.code_graph()
+        items = [
+            d.to_dict() for d in await cg.decisions(scope=scope or None, status=status or None)
+        ]
+        return {
+            "decisions": items,
+            "count": len(items),
+            **(await self.staleness()),
+            "tool_api_version": TOOL_API_VERSION,
+        }
+
+    async def explain(self, symbol_id: str) -> dict[str, Any]:
+        """A symbol's LLM summary + design-pattern tags (feat-012) + its 1-hop
+        typed facts — the reserved ckg_explain."""
+        from agentforge_graph.core import EdgeKind, NodeKind, SymbolID
+
+        cg = await self.code_graph()
+        node = await cg.store.graph.get(symbol_id)
+        tags: list[dict[str, Any]] = []
+        facts: list[dict[str, str]] = []
+        if node is not None:
+            for e in await cg.store.graph.adjacent(symbol_id, [EdgeKind.TAGGED], "out"):
+                target = await cg.store.graph.get(e.dst)
+                tags.append(
+                    {
+                        "pattern": target.name if target else "",
+                        "confidence": e.attrs.get("confidence", 0.0),
+                        "rationale": e.attrs.get("rationale", ""),
+                    }
+                )
+            for e in await cg.store.graph.adjacent(symbol_id, None, "both"):
+                if e.kind is not EdgeKind.TAGGED:
+                    facts.append({"src": e.src, "dst": e.dst, "kind": e.kind.value})
+        # the owning file's summary, if one exists (feat-012)
+        summary = ""
+        if node is not None:
+            path = SymbolID.parse(symbol_id).path
+            for n in (
+                await cg.store.graph.query(GraphQuery(kinds=[NodeKind.SUMMARY], limit=10**9))
+            ).nodes:
+                if str(n.attrs.get("level")) == "file" and str(n.attrs.get("path")) == path:
+                    summary = str(n.attrs.get("text", ""))
+                    break
+        return {
+            "symbol_id": symbol_id,
+            "name": node.name if node else "",
+            "kind": node.kind.value if node else "",
+            "summary": summary,
+            "tags": tags,
+            "facts": facts,
+            **(await self.staleness()),
+            "tool_api_version": TOOL_API_VERSION,
+        }
+
+    async def history(self, symbol_id: str) -> dict[str, Any]:
+        """A symbol's git evolution (feat-009): introduced / last-changed /
+        churn / authors / lifecycle events, wrapped in the staleness envelope."""
+        cg = await self.code_graph()
+        hist = await cg.history(symbol_id)
+        body: dict[str, Any] = (
+            hist.model_dump() if hist is not None else {"symbol_id": symbol_id, "available": False}
+        )
+        return {**body, **(await self.staleness()), "tool_api_version": TOOL_API_VERSION}
+
+    async def status(self) -> dict[str, Any]:
+        meta = self._meta()
+        cg = await self.code_graph()
+        nodes = (await cg.store.graph.query(GraphQuery(limit=_ALL))).nodes
+        head = _git_head(self.repo_path)
+        dirty = bool(head) and bool(meta.indexed_commit) and head != meta.indexed_commit
+        by_kind: dict[str, int] = {}
+        for n in nodes:
+            by_kind[n.kind.value] = by_kind.get(n.kind.value, 0) + 1
+        store_root = Path(self.repo_path) / StoreConfig.load(self.config).path
+        return {
+            "indexed_commit": meta.indexed_commit,
+            "head_commit": head,
+            "dirty": dirty,
+            "files_indexed": len(meta.files),
+            "nodes": len(nodes),
+            "by_kind": by_kind,
+            "temporal": await cg.temporal_status(),
+            "store_path": str(store_root),
+            "tool_api_version": TOOL_API_VERSION,
+        }
+
+    async def close(self) -> None:
+        if self._cg is not None:
+            await self._cg.close()

agentforge_graph/serve/http_runner.py

@@ -0,0 +1,133 @@
+"""Bearer-token auth for the HTTP MCP transport (ENH-005).
+
+The framework's ``MCPServer.from_http`` serves streamable-HTTP with **no auth** —
+fine for the localhost/trusted-container default, but a wide-open surface on any
+exposed port. ``agentforge-mcp`` exposes no auth hook (see the framework wishlist
+note), but ``from_http(runner=…)`` lets us inject a custom ``MCPServerRunner``.
+
+So: the **no-auth** path stays 100% framework (unchanged default). When a token is
+configured, ``serve`` uses :class:`CkgHttpRunner` here — it wires the MCP tools
+exactly like the framework runner, but wraps the Starlette app in
+:class:`BearerAuthMiddleware`, which rejects any request lacking a matching
+``Authorization: Bearer …`` with ``401`` (constant-time compare; the token is
+never logged). Drop this once the framework grows an auth hook.
+"""
+
+from __future__ import annotations
+
+import hmac
+from typing import Any
+
+# Hosts that don't need auth-by-default (the client owns the loopback surface).
+LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost", "::1", "0:0:0:0:0:0:0:1"})
+
+
+def is_loopback(host: str) -> bool:
+    return host in LOOPBACK_HOSTS
+
+
+class BearerAuthMiddleware:
+    """Pure-ASGI middleware: require ``Authorization: Bearer <token>`` on every
+    HTTP request, else ``401``. Non-HTTP scopes (lifespan) pass through."""
+
+    def __init__(self, app: Any, token: str) -> None:
+        self._app = app
+        self._expected = f"Bearer {token}"
+
+    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+        if scope["type"] != "http":
+            await self._app(scope, receive, send)
+            return
+        provided = ""
+        for name, value in scope.get("headers", []):
+            if name == b"authorization":
+                provided = value.decode("latin-1")
+                break
+        # constant-time compare so a wrong token can't be timed out char by char.
+        if not (provided and hmac.compare_digest(provided, self._expected)):
+            await _send_401(send)
+            return
+        await self._app(scope, receive, send)
+
+
+async def _send_401(send: Any) -> None:
+    body = b'{"error":"unauthorized"}'
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 401,
+            "headers": [
+                (b"content-type", b"application/json"),
+                (b"www-authenticate", b"Bearer"),
+                (b"content-length", str(len(body)).encode()),
+            ],
+        }
+    )
+    await send({"type": "http.response.body", "body": body})
+
+
+class CkgHttpRunner:
+    """An ``MCPServerRunner`` that serves the streamable-HTTP MCP transport with
+    a bearer-token gate. Mirrors the framework's HTTP runner wiring (list/call
+    tool dispatch + ``StreamableHTTPSessionManager`` under uvicorn) and wraps the
+    app in :class:`BearerAuthMiddleware`."""
+
+    def __init__(self, *, host: str, port: int, token: str, server_name: str = "ckg") -> None:
+        from mcp.server import Server
+
+        self._server = Server(server_name)
+        self._host = host
+        self._port = port
+        self._token = token
+        self._tools: dict[str, tuple[str, str, dict[str, Any], Any]] = {}
+        self._uv_server: Any = None
+
+    def register_tool(
+        self, name: str, description: str, input_schema: dict[str, Any], handler: Any
+    ) -> None:
+        self._tools[name] = (name, description, dict(input_schema or {}), handler)
+
+    async def serve(self) -> None:
+        import contextlib
+
+        import uvicorn
+        from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+        from mcp.types import TextContent
+        from mcp.types import Tool as MCPTool
+        from starlette.applications import Starlette
+        from starlette.routing import Mount
+
+        registered = self._tools
+
+        @self._server.list_tools()  # type: ignore[no-untyped-call, untyped-decorator]
+        async def _list() -> list[MCPTool]:
+            return [
+                MCPTool(name=n, description=d, inputSchema=s) for n, d, s, _ in registered.values()
+            ]
+
+        @self._server.call_tool()  # type: ignore[untyped-decorator]
+        async def _call(name: str, arguments: dict[str, Any]) -> list[TextContent]:
+            tool = registered.get(name)
+            if tool is None:
+                raise ValueError(f"MCP server: unknown tool {name!r}")
+            return [TextContent(type="text", text=await tool[3](arguments))]
+
+        manager = StreamableHTTPSessionManager(app=self._server, json_response=True, stateless=True)
+
+        async def _asgi(scope: Any, receive: Any, send: Any) -> None:
+            await manager.handle_request(scope, receive, send)
+
+        @contextlib.asynccontextmanager
+        async def _lifespan(_app: Any) -> Any:
+            async with manager.run():
+                yield
+
+        app = Starlette(routes=[Mount("/mcp", app=_asgi)], lifespan=_lifespan)
+        guarded = BearerAuthMiddleware(app, self._token)
+        config = uvicorn.Config(guarded, host=self._host, port=self._port, log_level="warning")
+        self._uv_server = uvicorn.Server(config)
+        await self._uv_server.serve()
+
+    async def stop(self) -> None:
+        if self._uv_server is not None:
+            self._uv_server.should_exit = True

agentforge_graph/serve/server.py

@@ -0,0 +1,110 @@
+"""Dual binding: the same nine ``Tool`` instances power both
+``code_graph_tools`` (for ``Agent(tools=…)``) and ``serve_mcp`` (an MCP server).
+One definition, every call site — no toolset drift.
+
+The MCP server runs over two transports (feat-008):
+
+- **stdio** (default) — the client launches ``ckg serve-mcp`` as a subprocess
+  (``command``/``args`` in an ``mcpServers`` block; ``claude mcp add``).
+- **http** — a long-running streamable-HTTP server (mounted at ``/mcp`` under
+  uvicorn) that clients connect to by ``url``. Same tools, same guardrails.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Literal
+
+from agentforge_core.contracts.tool import Tool
+from agentforge_mcp import MCPServer
+
+from .engine import _Engine
+from .http_runner import CkgHttpRunner, is_loopback
+from .tools import ALL_TOOLS
+
+Transport = Literal["stdio", "http"]
+
+
+def code_graph_tools(repo_path: str | Path = ".", config: str | Path | None = None) -> list[Tool]:
+    """The CKG toolset as native AgentForge ``Tool`` instances, sharing one
+    lazily-opened engine. Pass straight to ``Agent(tools=code_graph_tools("."))``."""
+    engine = _Engine(repo_path, config)
+    # ALL_TOOLS holds concrete Tool subclasses; mypy joins them to the abstract
+    # base and can't see that, hence the abstract ignore.
+    return [tool_cls(engine) for tool_cls in ALL_TOOLS]  # type: ignore[abstract]
+
+
+def build_mcp_server(
+    repo_path: str | Path = ".",
+    config: str | Path | None = None,
+    *,
+    transport: Transport = "stdio",
+    host: str = "127.0.0.1",
+    port: int = 8765,
+    refresh_on_call: bool = False,
+    auth_token: str = "",
+    allow_unauthenticated: bool = False,
+) -> MCPServer:
+    """Build (but don't serve) the MCP server with the CKG tools, over the
+    chosen ``transport`` (``stdio`` default, or ``http`` at ``host:port``).
+
+    HTTP auth (ENH-005): ``auth_token`` (or ``$CKG_HTTP_AUTH_TOKEN``) requires a
+    matching ``Authorization: Bearer …`` on every request — off by default to
+    preserve the localhost loop. Binding a **non-loopback** host with no token is
+    refused unless ``allow_unauthenticated`` (a loud, deliberate opt-in), so an
+    exposed port is never silently wide open. ``refresh_on_call`` is a no-op at
+    0.1 (tools are read-only; cheap refresh is feat-004)."""
+    if transport not in ("stdio", "http"):
+        msg = f"unknown MCP transport {transport!r}; use 'stdio' or 'http'"
+        raise ValueError(msg)
+    # Resolve + validate auth before opening the engine, so a misconfigured bind
+    # fails fast (not after the index is loaded).
+    token = auth_token or os.environ.get("CKG_HTTP_AUTH_TOKEN", "")
+    if transport == "http" and not token and not is_loopback(host) and not allow_unauthenticated:
+        msg = (
+            f"refusing to serve the HTTP MCP transport on non-loopback host {host!r} "
+            "with no auth: set serve.http_auth_token / $CKG_HTTP_AUTH_TOKEN, or pass "
+            "--allow-unauthenticated to bind it open on purpose"
+        )
+        raise ValueError(msg)
+    tools = code_graph_tools(repo_path, config)
+    allowed = tuple(t.name for t in tools)
+    if transport == "stdio":
+        # no auth: the client owns the subprocess (stdin/stdout).
+        return MCPServer.from_stdio(tools=tools, allowed=allowed, server_name="ckg")
+    if token:
+        runner = CkgHttpRunner(host=host, port=port, token=token)
+        return MCPServer.from_http(
+            tools=tools, host=host, port=port, allowed=allowed, server_name="ckg", runner=runner
+        )
+    return MCPServer.from_http(
+        tools=tools, host=host, port=port, allowed=allowed, server_name="ckg"
+    )
+
+
+async def serve_mcp(
+    repo_path: str | Path = ".",
+    config: str | Path | None = None,
+    *,
+    transport: Transport = "stdio",
+    host: str = "127.0.0.1",
+    port: int = 8765,
+    refresh_on_call: bool = False,
+    auth_token: str = "",
+    allow_unauthenticated: bool = False,
+) -> None:
+    """Run the CKG MCP server (blocks until stopped). ``transport='http'`` serves
+    streamable-HTTP at ``http://{host}:{port}/mcp``; the default ``stdio`` serves
+    over stdin/stdout for a subprocess client. See ``build_mcp_server`` for the
+    ``auth_token`` / ``allow_unauthenticated`` HTTP-auth semantics (ENH-005)."""
+    await build_mcp_server(
+        repo_path,
+        config,
+        transport=transport,
+        host=host,
+        port=port,
+        refresh_on_call=refresh_on_call,
+        auth_token=auth_token,
+        allow_unauthenticated=allow_unauthenticated,
+    ).serve()