141 lines
6.0 KiB
Python
141 lines
6.0 KiB
Python
"""Internal MCP aggregator — one ASGI app, N FastMCP namespaces + ``all``.
|
|
|
|
Each ``McpServer`` declared in the user's config becomes its own
|
|
``FastMCP`` instance (regular for ``python_tool``, ``FastMCPProxy`` for
|
|
``stdio``/``http``) and is mounted under ``/mcp/<name>`` on a single
|
|
Starlette app. This app runs on ``127.0.0.1:INTERNAL_MCP_PORT`` (not
|
|
EXPOSE'd in Docker) so the ClaudeCode subprocess can reach each
|
|
namespace via loopback as a distinct MCP server URL — preserving
|
|
per-domain framing while costing only one process worth of RAM
|
|
(PRD §6).
|
|
|
|
Phase 3 adds ``/mcp/all/``: a single FastMCP whose tools are the union
|
|
of every namespace's tools, prefixed by FastMCP's ``namespace_<tool>``
|
|
convention (e.g. ``time_current_time``). It's the escape-hatch for
|
|
clients that can only configure one MCP server — discouraged for tool-
|
|
heavy setups (PRD §6 cites the ~95%→~71% tool-selection drop on flat
|
|
namespaces) but real and reachable.
|
|
|
|
The aggregator returns both the app and a ``{name: url}`` map; Phase
|
|
``ClaudeSdkBackend`` plugs the map directly into
|
|
``BackendOptions.mcp_servers``. ``/mcp/all/`` is NOT included in that
|
|
map — claude-code-agents always get per-domain framing; only the
|
|
external MCP frontend (Phase 3.1) reverse-proxies the flat endpoint.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import AsyncExitStack, asynccontextmanager
|
|
from typing import TYPE_CHECKING
|
|
|
|
from fastmcp import FastMCP
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import Mount
|
|
|
|
from beaver_gateway.mcp.client_pool import build_http_proxy, build_stdio_proxy
|
|
from beaver_gateway.mcp.types import HttpMcp, PythonToolMcp, StdioMcp
|
|
from beaver_gateway.mcp.wrap import build_python_tool_server
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator, Iterable
|
|
|
|
from beaver_gateway.mcp.types import McpServerT
|
|
|
|
|
|
ALL_NAMESPACE = "all"
|
|
"""URL segment for the flat-namespace aggregator (``/mcp/all/``)."""
|
|
|
|
|
|
def build_internal_app(
|
|
mcps: Iterable[McpServerT], *, host: str, port: int
|
|
) -> tuple[Starlette, dict[str, str], dict[str, FastMCP]]:
|
|
"""Build the aggregator ``Starlette`` app, per-namespace URL map, and server map.
|
|
|
|
``host``/``port`` only flavour the URL strings handed back — actually
|
|
listening on them is the caller's job (``cli.main`` runs a uvicorn
|
|
server in a TaskGroup). We accept the address here so callers don't
|
|
have to format the URLs themselves and risk drifting from the
|
|
``/mcp/<name>`` convention.
|
|
|
|
Returns:
|
|
* Starlette app to serve via uvicorn.
|
|
* ``{namespace: url}`` for the per-domain endpoints (claude-code's
|
|
MCP routing expects per-domain framing). ``/mcp/all/`` is
|
|
omitted — only meaningful to external clients via the MCP
|
|
frontend, not to claude-code.
|
|
* ``{namespace: FastMCP}`` for backends that need in-process
|
|
access (Raycast doesn't natively understand MCP, so the
|
|
gateway calls ``list_tools``/``call_tool`` directly to splice
|
|
MCP tools into the Raycast wire).
|
|
"""
|
|
servers: dict[str, FastMCP] = {spec.name: _build_server(spec) for spec in mcps}
|
|
|
|
child_apps = {
|
|
name: s.http_app(transport="http", path="/") for name, s in servers.items()
|
|
}
|
|
routes = [Mount(f"/mcp/{name}", app=app) for name, app in child_apps.items()]
|
|
|
|
# /mcp/all — flat-namespace bundle. Skip when there's nothing to
|
|
# bundle so we don't pay for an empty session manager lifecycle.
|
|
all_app = None
|
|
if servers:
|
|
all_server = _build_all_server(servers)
|
|
all_app = all_server.http_app(transport="http", path="/")
|
|
routes.append(Mount(f"/mcp/{ALL_NAMESPACE}", app=all_app))
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_parent: Starlette) -> AsyncIterator[None]:
|
|
# Each FastMCP http_app stores its session manager init in its
|
|
# own lifespan. Without entering them the streamable-HTTP layer
|
|
# 500s on every request. AsyncExitStack composes them so all
|
|
# children come up together and unwind in reverse order on
|
|
# shutdown.
|
|
async with AsyncExitStack() as stack:
|
|
for child in child_apps.values():
|
|
await stack.enter_async_context(child.router.lifespan_context(child))
|
|
if all_app is not None:
|
|
await stack.enter_async_context(
|
|
all_app.router.lifespan_context(all_app)
|
|
)
|
|
yield
|
|
|
|
app = Starlette(routes=routes, lifespan=lifespan)
|
|
# Trailing slash on the published URL skips Starlette's
|
|
# 307 redirect from ``/mcp/<name>`` to ``/mcp/<name>/`` that
|
|
# ``Mount`` produces when a child route lives at ``/``.
|
|
urls = {name: f"http://{host}:{port}/mcp/{name}/" for name in servers}
|
|
return app, urls, servers
|
|
|
|
|
|
def _build_server(spec: McpServerT) -> FastMCP:
|
|
"""Dispatch on the discriminated union to the matching builder."""
|
|
if isinstance(spec, PythonToolMcp):
|
|
return build_python_tool_server(spec)
|
|
if isinstance(spec, StdioMcp):
|
|
return build_stdio_proxy(spec)
|
|
if isinstance(spec, HttpMcp):
|
|
return build_http_proxy(spec)
|
|
# `McpServerT` is a closed union; this is unreachable but keeps
|
|
# type-narrowing honest if a new variant lands without updates here.
|
|
msg = f"unsupported McpServer variant: {type(spec).__name__}"
|
|
raise TypeError(msg)
|
|
|
|
|
|
def _build_all_server(children: dict[str, FastMCP]) -> FastMCP:
|
|
"""Compose a single FastMCP whose tools union every namespace's tools.
|
|
|
|
FastMCP's ``mount(namespace=...)`` namespaces every tool as
|
|
``<namespace>_<original_name>``, so the flat endpoint becomes
|
|
``time_current_time``, ``calendar_event_create``, etc.
|
|
|
|
Mounted children are accessed in-memory via ``FastMCPProvider`` —
|
|
requests don't bounce through the child's own ``http_app``, so this
|
|
aggregator has its own independent streamable-HTTP session manager
|
|
and lifespan, and the per-namespace ``/mcp/<name>/`` mounts keep
|
|
working unchanged.
|
|
"""
|
|
parent = FastMCP(name="beaver-gateway-all")
|
|
for name, child in children.items():
|
|
parent.mount(child, namespace=name)
|
|
return parent
|