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()