197 lines
7.6 KiB
Python
197 lines
7.6 KiB
Python
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()
|