chore: import upstream snapshot with attribution

2026-07-13 12:12:13 +08:00
commit 0446c45d8e
898 changed files with 328024 additions and 0 deletions
@@ -0,0 +1,273 @@
+# deploy/docker/mcp_bridge.py
+
+from __future__ import annotations
+import inspect, json, re, anyio
+from contextlib import suppress
+from typing import Any, Callable, Dict, List, Tuple
+import httpx
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from starlette.routing import Route, Mount
+from mcp.server.sse import SseServerTransport
+
+import mcp.types as t
+from mcp.server.lowlevel.server import Server, NotificationOptions
+from mcp.server.models import InitializationOptions
+
+# ── opt‑in decorators ───────────────────────────────────────────
+def mcp_resource(name: str | None = None):
+    def deco(fn):
+        fn.__mcp_kind__, fn.__mcp_name__ = "resource", name
+        return fn
+    return deco
+
+def mcp_template(name: str | None = None):
+    def deco(fn):
+        fn.__mcp_kind__, fn.__mcp_name__ = "template", name
+        return fn
+    return deco
+
+def mcp_tool(name: str | None = None):
+    def deco(fn):
+        fn.__mcp_kind__, fn.__mcp_name__ = "tool", name
+        return fn
+    return deco
+
+# ── HTTP‑proxy helper for FastAPI endpoints ─────────────────────
+def _service_auth_headers() -> dict:
+    """Authenticate the internal loopback call to our own gated endpoints.
+
+    The MCP transport is now behind the AuthGateMiddleware, so by the time a
+    tool is invoked the MCP client is already authenticated. The loopback HTTP
+    call this proxy makes must therefore carry a credential too, or the gate
+    would 401 it. We mint a short-lived, data-scope service token: MCP exposes
+    only data-plane tools (no admin/monitor actions), so this grants no
+    privilege escalation. (Requires a shared SECRET_KEY across workers, which
+    the auth startup check already mandates for any real deployment.)
+    """
+    from auth import create_access_token
+    token = create_access_token({"sub": "mcp-service"}, scope="data")
+    return {"Authorization": f"Bearer {token}"}
+
+
+def _make_http_proxy(base_url: str, route, *, timeout: float | None = None):
+    method = list(route.methods - {"HEAD", "OPTIONS"})[0]
+    async def proxy(**kwargs):
+        # replace `/items/{id}` style params first
+        path = route.path
+        for k, v in list(kwargs.items()):
+            placeholder = "{" + k + "}"
+            if placeholder in path:
+                path = path.replace(placeholder, str(v))
+                kwargs.pop(k)
+        url = base_url.rstrip("/") + path
+
+        headers = _service_auth_headers()
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            try:
+                r = (
+                    await client.get(url, params=kwargs, headers=headers)
+                    if method == "GET"
+                    else await client.request(method, url, json=kwargs, headers=headers)
+                )
+                r.raise_for_status()
+                return r.text if method == "GET" else r.json()
+            except httpx.HTTPStatusError as e:
+                # surface FastAPI error details instead of plain 500
+                raise HTTPException(e.response.status_code, e.response.text)
+            except httpx.TimeoutException:
+                raise HTTPException(504, "upstream request timed out")
+    return proxy
+
+# ── main entry point ────────────────────────────────────────────
+def attach_mcp(
+    app: FastAPI,
+    *,                          # keyword‑only
+    base: str = "/mcp",
+    name: str | None = None,
+    base_url: str,              # eg. "http://127.0.0.1:8020"
+    timeout: float | None = None,  # httpx timeout in seconds; None = no limit
+) -> None:
+    """Call once after all routes are declared to expose WS+SSE MCP endpoints."""
+    server_name = name or app.title or "FastAPI-MCP"
+    mcp = Server(server_name)
+
+    # tools: Dict[str, Callable] = {}
+    tools: Dict[str, Tuple[Callable, Callable]] = {}
+    resources: Dict[str, Callable] = {}
+    templates: Dict[str, Callable] = {}
+
+    # register decorated FastAPI routes
+    for route in app.routes:
+        fn = getattr(route, "endpoint", None)
+        kind = getattr(fn, "__mcp_kind__", None)
+        if not kind:
+            continue
+
+        key = fn.__mcp_name__ or re.sub(r"[/{}}]", "_", route.path).strip("_")
+
+        # if kind == "tool":
+        #     tools[key] = _make_http_proxy(base_url, route)
+        if kind == "tool":
+            proxy = _make_http_proxy(base_url, route, timeout=timeout)
+            tools[key] = (proxy, fn)
+            continue
+        if kind == "resource":
+            resources[key] = fn
+        if kind == "template":
+            templates[key] = fn
+
+    # helpers for JSON‑Schema
+    def _schema(model: type[BaseModel] | None) -> dict:
+        return {"type": "object"} if model is None else model.model_json_schema()
+
+    def _body_model(fn: Callable) -> type[BaseModel] | None:
+        for p in inspect.signature(fn).parameters.values():
+            a = p.annotation
+            if inspect.isclass(a) and issubclass(a, BaseModel):
+                return a
+        return None
+
+    # MCP handlers
+    @mcp.list_tools()
+    async def _list_tools() -> List[t.Tool]:
+        out = []
+        for k, (proxy, orig_fn) in tools.items():
+            desc   = getattr(orig_fn, "__mcp_description__", None) or inspect.getdoc(orig_fn) or ""
+            schema = getattr(orig_fn, "__mcp_schema__", None) or _schema(_body_model(orig_fn))
+            out.append(
+                t.Tool(name=k, description=desc, inputSchema=schema)
+            )
+        return out
+             
+
+    @mcp.call_tool()
+    async def _call_tool(name: str, arguments: Dict | None) -> List[t.TextContent]:
+        if name not in tools:
+            raise HTTPException(404, "tool not found")
+        
+        proxy, _ = tools[name]
+        try:
+            res = await proxy(**(arguments or {}))
+        except HTTPException as exc:
+            # map server‑side errors into MCP "text/error" payloads
+            err = {"error": exc.status_code, "detail": exc.detail}
+            return [t.TextContent(type = "text", text=json.dumps(err, ensure_ascii=False))]
+        return [t.TextContent(type = "text", text=json.dumps(res, default=str, ensure_ascii=False))]
+
+    @mcp.list_resources()
+    async def _list_resources() -> List[t.Resource]:
+        return [
+            t.Resource(name=k, description=inspect.getdoc(f) or "", mime_type="application/json")
+            for k, f in resources.items()
+        ]
+
+    @mcp.read_resource()
+    async def _read_resource(name: str) -> List[t.TextContent]:
+        if name not in resources:
+            raise HTTPException(404, "resource not found")
+        res = resources[name]()
+        return [t.TextContent(type = "text", text=json.dumps(res, default=str, ensure_ascii=False))]
+
+    @mcp.list_resource_templates()
+    async def _list_templates() -> List[t.ResourceTemplate]:
+        return [
+            t.ResourceTemplate(
+                name=k,
+                description=inspect.getdoc(f) or "",
+                parameters={
+                    p: {"type": "string"} for p in _path_params(app, f)
+                },
+            )
+            for k, f in templates.items()
+        ]
+
+    init_opts = InitializationOptions(
+        server_name=server_name,
+        server_version="0.1.0",
+        capabilities=mcp.get_capabilities(
+            notification_options=NotificationOptions(),
+            experimental_capabilities={},
+        ),
+    )
+
+    # ── WebSocket transport ────────────────────────────────────
+    @app.websocket_route(f"{base}/ws")
+    async def _ws(ws: WebSocket):
+        await ws.accept()
+        c2s_send, c2s_recv = anyio.create_memory_object_stream(100)
+        s2c_send, s2c_recv = anyio.create_memory_object_stream(100)
+
+        from pydantic import TypeAdapter
+        from mcp.types import JSONRPCMessage
+        adapter = TypeAdapter(JSONRPCMessage)
+
+        init_done = anyio.Event()
+
+        async def srv_to_ws():
+            first = True 
+            try:
+                async for msg in s2c_recv:
+                    await ws.send_json(msg.model_dump())
+                    if first:
+                        init_done.set()
+                        first = False
+            finally:
+                # make sure cleanup survives TaskGroup cancellation
+                with anyio.CancelScope(shield=True):
+                    with suppress(RuntimeError):       # idempotent close
+                        await ws.close()
+
+        async def ws_to_srv():
+            try:
+                # 1st frame is always "initialize"
+                first = adapter.validate_python(await ws.receive_json())
+                await c2s_send.send(first)
+                await init_done.wait()          # block until server ready
+                while True:
+                    data = await ws.receive_json()
+                    await c2s_send.send(adapter.validate_python(data))
+            except WebSocketDisconnect:
+                await c2s_send.aclose()
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(mcp.run, c2s_recv, s2c_send, init_opts)
+            tg.start_soon(ws_to_srv)
+            tg.start_soon(srv_to_ws)
+
+    # ── SSE transport (raw ASGI — avoids Starlette middleware conflict) ──
+    sse = SseServerTransport(f"{base}/messages/")
+
+    # Starlette's Route wraps plain async functions in request_response(),
+    # which calls handler(request) instead of handler(scope, receive, send).
+    # Using a callable class bypasses this — Route passes classes through
+    # as raw ASGI apps.  See #1594, #1850.
+    class _MCPSseApp:
+        async def __call__(self, scope, receive, send):
+            async with sse.connect_sse(scope, receive, send) as (read_stream, write_stream):
+                await mcp.run(read_stream, write_stream, init_opts)
+
+    app.routes.append(Route(f"{base}/sse", endpoint=_MCPSseApp()))
+    app.routes.append(Mount(f"{base}/messages", app=sse.handle_post_message))
+
+    # ── schema endpoint ───────────────────────────────────────
+    @app.get(f"{base}/schema")
+    async def _schema_endpoint():
+        return JSONResponse({
+            "tools": [x.model_dump() for x in await _list_tools()],
+            "resources": [x.model_dump() for x in await _list_resources()],
+            "resource_templates": [x.model_dump() for x in await _list_templates()],
+        })
+
+
+# ── helpers ────────────────────────────────────────────────────
+def _route_name(path: str) -> str:
+    return re.sub(r"[/{}}]", "_", path).strip("_")
+
+def _path_params(app: FastAPI, fn: Callable) -> List[str]:
+    for r in app.routes:
+        if r.endpoint is fn:
+            return list(r.param_convertors.keys())
+    return []