266 lines
12 KiB
Python
266 lines
12 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": "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())
|