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": "projectbrain.project.list", "list_project_repositories": "projectbrain.repository.list", "get_portfolio_control_plane": "projectbrain.portfolio.get", "list_portfolio_repositories": "projectbrain.portfolio.repository.list", "reconcile_portfolio_repositories": "projectbrain.portfolio.repository.reconcile", "set_portfolio_repository_exclusion": "projectbrain.portfolio.repository.exclusion.set", "search": "projectbrain.repository.search", "fetch": "projectbrain.repository.fetch", "get_project_context": "projectbrain.context.get", "render_project_workspace": "projectbrain.workspace.render", "get_project_modules": "projectbrain.modules.get", "save_brainstorm": "projectbrain.brainstorm.save", "save_project_knowledge": "projectbrain.knowledge.save", "save_documentation_draft": "projectbrain.documentation.save", "save_memory_rule": "projectbrain.memory.save", "create_project_context_pack": "projectbrain.contextpack.create", "record_execution_feedback": "projectbrain.feedback.record", "list_feature_deliveries": "projectbrain.delivery.list", "get_feature_delivery": "projectbrain.delivery.get", "create_feature_delivery": "projectbrain.delivery.create", "update_feature_delivery": "projectbrain.delivery.update", "list_repository_tree": "projectbrain.repository.tree", "get_file": "projectbrain.repository.file", "list_branches": "projectbrain.repository.branches", "list_pull_requests": "projectbrain.repository.pullrequests", "start_feature_delivery": "projectbrain.delivery.start", "write_feature_file": "projectbrain.delivery.file.write", "delete_feature_file": "projectbrain.delivery.file.delete", "apply_feature_changeset": "projectbrain.delivery.changeset.apply", "record_feature_validation": "projectbrain.delivery.validation.record", "get_feature_validation_runs": "projectbrain.delivery.validation.runs", "sync_feature_validation": "projectbrain.delivery.validation.sync", "complete_feature_delivery": "projectbrain.delivery.complete", "merge_feature_delivery": "projectbrain.delivery.merge", "start_repository_analysis": "projectbrain.analysis.start", "list_analysis_jobs": "projectbrain.analysis.list", "get_analysis_job": "projectbrain.analysis.status", "get_analysis_job_results": "projectbrain.analysis.results", "cancel_analysis_job": "projectbrain.analysis.cancel", "render_analysis_job": "projectbrain.analysis.render", "list_repository_findings": "projectbrain.finding.list", "get_repository_finding": "projectbrain.finding.get", "create_finding_issue": "projectbrain.finding.issue.sync", "start_finding_investigation": "projectbrain.finding.investigation.start", "get_finding_investigation": "projectbrain.finding.investigation.get", "save_finding_investigation": "projectbrain.finding.investigation.save", "suppress_repository_finding": "projectbrain.finding.suppress", "create_delivery_from_finding": "projectbrain.finding.delivery.create", "verify_finding_resolution": "projectbrain.finding.verify", "list_auto_mode_projects": "projectbrain.autonomy.projects.list", "heartbeat_auto_executor": "projectbrain.autonomy.executor.heartbeat", "get_auto_mode_state": "projectbrain.autonomy.state.get", "configure_auto_mode": "projectbrain.autonomy.configure", "import_auto_roadmap": "projectbrain.autonomy.roadmap.import", "resume_auto_mode": "projectbrain.autonomy.resume", "claim_auto_work": "projectbrain.autonomy.work.claim", "report_auto_work": "projectbrain.autonomy.work.report", } PUBLIC_TO_UPSTREAM = {value: key for key, value in UPSTREAM_TO_PUBLIC.items()} MAX_BODY_BYTES = 4 * 1024 * 1024 MAX_RESPONSE_BYTES = 2 * 1024 * 1024 def upstream_headers( request_headers: Mapping[str, str], token: str, *, has_body: bool ) -> dict[str, str]: """Build the bounded MCP transport headers forwarded to ProjectBrain.""" headers = { "Authorization": f"Bearer {token}", # Streamable HTTP POST clients must accept either valid response form. # Using the same explicit value for GET/DELETE is harmless and avoids # forwarding browser-specific wildcard Accept headers. "Accept": "application/json, text/event-stream", } if has_body: headers["Content-Type"] = request_headers.get("content-type", "application/json") for name in ( "mcp-session-id", "mcp-protocol-version", "last-event-id", "x-itworx-client-id", ): if value := request_headers.get(name): headers[name] = value return headers def _settings() -> tuple[str, str, set[str], bytes]: endpoint = os.environ["PROJECTBRAIN_ENDPOINT"].rstrip("/") parsed = urlparse(endpoint) if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password: raise ValueError("ProjectBrain endpoint must be an HTTPS origin") token = read_secret_file( Path( os.environ.get( "PROJECTBRAIN_SERVICE_TOKEN_FILE", "/run/secrets/projectbrain_hub_service_token" ) ) ) if len(token) < 32: raise ValueError("ProjectBrain service token must contain at least 32 characters") projects = { item.strip().lower() for item in os.environ.get("PROJECTBRAIN_PROJECTS", "gitea-all").split(",") if item.strip() } return endpoint, token, projects, load_connector_identity_key() 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): return document public_name = params["name"] upstream_name = PUBLIC_TO_UPSTREAM.get(public_name) if upstream_name is None: raise ValueError("ProjectBrain 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): return document tools = result.get("tools") if not isinstance(tools, list): return document rewritten = [] for tool in tools: if not isinstance(tool, dict) or not isinstance(tool.get("name"), str): continue public_name = UPSTREAM_TO_PUBLIC.get(tool["name"]) if public_name: rewritten.append({**tool, "name": public_name}) return {**document, "result": {**result, "tools": rewritten}} def projectbrain_readiness_status(payload: Any) -> str | None: """Return the readiness status from either the API envelope or legacy payload.""" if not isinstance(payload, dict): return None data = payload.get("data", payload) return data.get("status") if isinstance(data, dict) else None def build_application() -> Starlette: endpoint, token, projects, signing_key = _settings() client = httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=5.0), follow_redirects=False) async def health(_: Request) -> JSONResponse: return JSONResponse({"status": "healthy", "component": "projectbrain-connector"}) async def ready(_: Request) -> JSONResponse: try: response = await client.get(f"{endpoint}/api/system/readiness") payload = response.json() ready_state = response.status_code == 200 and projectbrain_readiness_status( payload ) in {"ready", "degraded"} except (httpx.HTTPError, ValueError, TypeError): ready_state, payload = False, {} return JSONResponse( { "status": "ready" if ready_state else "unavailable", "component": "projectbrain-connector", "upstream": payload, }, status_code=200 if ready_state else 503, ) async def mcp(request: Request) -> Response: body = await request.body() if len(body) > MAX_BODY_BYTES: return JSONResponse({"error": "PROJECTBRAIN_REQUEST_TOO_LARGE"}, status_code=413) document: Any = None 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(key).lower(): str(value) for key, value in request.headers.items()}, projects, signing_key=signing_key, ) document = rewrite_request(document) body = json.dumps(document, separators=(",", ":")).encode() except PermissionError: return JSONResponse( {"error": "PROJECTBRAIN_CONNECTOR_IDENTITY_DENIED"}, status_code=403 ) except (ValueError, TypeError, json.JSONDecodeError): return JSONResponse({"error": "PROJECTBRAIN_REQUEST_INVALID"}, status_code=400) headers = upstream_headers(request.headers, token, has_body=bool(body)) try: upstream = await client.request( request.method, f"{endpoint}/mcp", content=body or None, headers=headers ) except httpx.HTTPError: return JSONResponse({"error": "PROJECTBRAIN_UNAVAILABLE"}, status_code=502) content = upstream.content if len(content) > MAX_RESPONSE_BYTES: return JSONResponse({"error": "PROJECTBRAIN_RESPONSE_TOO_LARGE"}, status_code=502) content_type = upstream.headers.get("content-type", "application/json") 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": "PROJECTBRAIN_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())