Initial curated ITWorx_MCP_Hub source from 368fba691ea8d23403b4d56b53b5ec670d7c5020

services/runtime/connector_stitch.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import AsyncIterator, Mapping
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import httpx
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route
+
+from packages.hub_runtime import read_secret_file
+
+from .common import connector_context, load_connector_identity_key, run_asgi
+
+UPSTREAM_TO_PUBLIC = {
+    "list_projects": "stitch.project.list",
+    "get_project": "stitch.project.get",
+    "create_project": "stitch.project.create",
+    "list_screens": "stitch.screen.list",
+    "get_screen": "stitch.screen.get",
+    "generate_screen_from_text": "stitch.screen.generate",
+    "edit_screens": "stitch.screen.edit",
+    "generate_variants": "stitch.screen.variants",
+}
+PUBLIC_TO_UPSTREAM = {public: upstream for upstream, public in UPSTREAM_TO_PUBLIC.items()}
+MAX_BODY_BYTES = 512 * 1024
+MAX_RESPONSE_BYTES = 4 * 1024 * 1024
+DEFAULT_ENDPOINT = "https://stitch.googleapis.com/mcp"
+
+
+def validate_endpoint(value: str) -> str:
+    endpoint = value.rstrip("/")
+    parsed = urlparse(endpoint)
+    if (
+        parsed.scheme != "https"
+        or parsed.hostname != "stitch.googleapis.com"
+        or parsed.username
+        or parsed.password
+        or parsed.query
+        or parsed.fragment
+        or parsed.path != "/mcp"
+    ):
+        raise ValueError("Stitch endpoint must be the official HTTPS MCP endpoint")
+    return endpoint
+
+
+def upstream_headers(
+    request_headers: Mapping[str, str], api_key: str, *, has_body: bool
+) -> dict[str, str]:
+    headers = {
+        "X-Goog-Api-Key": api_key,
+        "Accept": "application/json, text/event-stream",
+    }
+    if has_body:
+        headers["Content-Type"] = "application/json"
+    for name in ("mcp-session-id", "mcp-protocol-version", "last-event-id"):
+        if value := request_headers.get(name):
+            headers[name] = value
+    return headers
+
+
+def rewrite_request(document: Any) -> Any:
+    if not isinstance(document, dict) or document.get("method") != "tools/call":
+        return document
+    params = document.get("params")
+    if not isinstance(params, dict) or not isinstance(params.get("name"), str):
+        raise ValueError("invalid Stitch tool call")
+    upstream_name = PUBLIC_TO_UPSTREAM.get(params["name"])
+    if upstream_name is None:
+        raise ValueError("Stitch tool is not registered")
+    return {**document, "params": {**params, "name": upstream_name}}
+
+
+def rewrite_response(document: Any, request_method: str | None) -> Any:
+    if request_method != "tools/list" or not isinstance(document, dict):
+        return document
+    result = document.get("result")
+    if not isinstance(result, dict) or not isinstance(result.get("tools"), list):
+        return document
+    tools = []
+    for tool in result["tools"]:
+        if not isinstance(tool, dict) or not isinstance(tool.get("name"), str):
+            continue
+        if public_name := UPSTREAM_TO_PUBLIC.get(tool["name"]):
+            tools.append({**tool, "name": public_name})
+    return {**document, "result": {**result, "tools": tools}}
+
+
+def _settings() -> tuple[str, str, set[str], bytes]:
+    endpoint = validate_endpoint(os.environ.get("STITCH_ENDPOINT", DEFAULT_ENDPOINT))
+    api_key = read_secret_file(
+        Path(os.environ.get("STITCH_API_KEY_FILE", "/run/secrets/stitch_api_key"))
+    )
+    if len(api_key) < 20:
+        raise ValueError("Stitch API key is invalid")
+    projects = {
+        item.strip().lower()
+        for item in os.environ.get("STITCH_HUB_PROJECTS", "stitch-design").split(",")
+        if item.strip()
+    }
+    return endpoint, api_key, projects, load_connector_identity_key()
+
+
+def build_application() -> Starlette:
+    endpoint, api_key, projects, signing_key = _settings()
+    client = httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=10.0), follow_redirects=False)
+
+    async def health(_: Request) -> JSONResponse:
+        return JSONResponse({"status": "healthy", "component": "stitch-connector"})
+
+    async def ready(_: Request) -> JSONResponse:
+        # Configuration and secret loading already fail closed at startup. Avoid a
+        # write-capable upstream call or leaking Google's authentication response.
+        return JSONResponse({"status": "ready", "component": "stitch-connector"})
+
+    async def mcp(request: Request) -> Response:
+        body = await request.body()
+        if len(body) > MAX_BODY_BYTES:
+            return JSONResponse({"error": "STITCH_REQUEST_TOO_LARGE"}, status_code=413)
+        method: str | None = None
+        if body:
+            try:
+                document = json.loads(body)
+                if not isinstance(document, dict):
+                    raise ValueError("JSON-RPC batches are not accepted")
+                method = document.get("method")
+                if method == "tools/call":
+                    connector_context(
+                        {str(k).lower(): str(v) for k, v in request.headers.items()},
+                        projects,
+                        signing_key=signing_key,
+                    )
+                document = rewrite_request(document)
+                body = json.dumps(document, separators=(",", ":")).encode()
+            except PermissionError:
+                return JSONResponse({"error": "STITCH_CONNECTOR_IDENTITY_DENIED"}, status_code=403)
+            except (ValueError, TypeError, json.JSONDecodeError):
+                return JSONResponse({"error": "STITCH_REQUEST_INVALID"}, status_code=400)
+        try:
+            upstream = await client.request(
+                request.method,
+                endpoint,
+                content=body or None,
+                headers=upstream_headers(request.headers, api_key, has_body=bool(body)),
+            )
+        except httpx.HTTPError:
+            return JSONResponse({"error": "STITCH_UNAVAILABLE"}, status_code=502)
+        content = upstream.content
+        if len(content) > MAX_RESPONSE_BYTES:
+            return JSONResponse({"error": "STITCH_RESPONSE_TOO_LARGE"}, status_code=502)
+        content_type = upstream.headers.get("content-type", "application/json")
+        if method == "tools/list" and "text/event-stream" in content_type:
+            # Discovery must be parsed so the reviewed allowlist can be applied.
+            # Never relay an opaque event stream that could contain new tools.
+            return JSONResponse({"error": "STITCH_DISCOVERY_RESPONSE_INVALID"}, status_code=502)
+        if content and "application/json" in content_type:
+            try:
+                content = json.dumps(
+                    rewrite_response(upstream.json(), method), separators=(",", ":")
+                ).encode()
+            except (ValueError, TypeError):
+                return JSONResponse({"error": "STITCH_RESPONSE_INVALID"}, status_code=502)
+        response_headers = {"content-type": content_type}
+        if session_id := upstream.headers.get("mcp-session-id"):
+            response_headers["mcp-session-id"] = session_id
+        return Response(content, status_code=upstream.status_code, headers=response_headers)
+
+    @asynccontextmanager
+    async def lifespan(_: Starlette) -> AsyncIterator[None]:
+        yield
+        await client.aclose()
+
+    return Starlette(
+        routes=[
+            Route("/health", health),
+            Route("/healthz", health),
+            Route("/ready", ready),
+            Route("/health/ready", ready),
+            Route("/mcp", mcp, methods=["GET", "POST", "DELETE"]),
+        ],
+        lifespan=lifespan,
+    )
+
+
+def run() -> None:
+    run_asgi(build_application())
+
+
+if __name__ == "__main__":
+    run()