from __future__ import annotations import json import os import re from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from datetime import UTC, datetime from pathlib import Path from typing import Any from uuid import UUID import yaml # type: ignore[import-untyped] from fastapi import FastAPI, Request, Response from sqlalchemy import select from packages.hub_policy import Governance, PolicyEngine, PolicyRule, SlidingWindowRateLimiter from packages.hub_runtime import read_secret_file from services.control_api.auth import OIDCProviderClient, OIDCSettings, PrincipalMapper from services.control_api.auth.errors import InvalidToken, UnknownMachinePrincipal from services.control_api.auth.machine import MachineClientBinding from services.control_api.contextforge import ContextForgeClient, ContextForgeSettings from services.control_api.persistence import ( SqlAlchemyAliasSource, create_engine_from_url, session_factory, ) from services.control_api.persistence.demo_client import DEMO_CLIENT_ID from services.control_api.persistence.models import ClientRow, WorkspaceRow from services.dataplane_router import ( AliasRegistry, DataplaneProxy, InvocationRecorder, RouterSettings, build_router, ) from services.dataplane_router.models import InvocationMetadata from .audit import SqlAlchemyAuditLog from .common import ( hub_database_url, identity_assertion, load_connector_identity_key, run_asgi, ) from .demo_token import TOKEN_PREFIX, DemoTokenAuthenticator, load_demo_token_signing_key from .mcp_skills import StaticSkills from .oauth_metadata import OAuthMetadataSettings, bearer_challenge class JsonLineInvocationRecorder(InvocationRecorder): """Emit bounded response-contract metadata without request data or credentials.""" def record(self, metadata: InvocationMetadata) -> None: print( json.dumps( { "event": "mcp.response.completed", "traceId": metadata.trace_id, "project": metadata.project_slug, "mcpMethod": metadata.mcp_method, "mcpName": metadata.mcp_name, "protocolRevision": metadata.protocol_revision, "responseStatus": metadata.response_status, "responseBytes": metadata.response_bytes, "responseContentType": metadata.response_content_type, "jsonrpcResponse": metadata.jsonrpc_response, "responseIdMatches": metadata.response_id_matches, "responseShape": metadata.response_shape, "durationMs": metadata.duration_ms, }, separators=(",", ":"), sort_keys=True, ), flush=True, ) def _policy_engine(catalog_root: Path) -> PolicyEngine: policies: list[PolicyRule] = [] for path in sorted((catalog_root / "policies").glob("*.yaml")): document = yaml.safe_load(path.read_text(encoding="utf-8")) if not isinstance(document, dict): raise ValueError("policy catalog document is invalid") policies.append(PolicyRule.from_document(document)) if not policies or not any(policy.slug == "global-deny" for policy in policies): raise ValueError("policy catalog must include global-deny") return PolicyEngine(policies) def _tool_governance_metadata(catalog_root: Path) -> dict[str, tuple[str, str, str]]: metadata: dict[str, tuple[str, str, str]] = {} for path in sorted((catalog_root / "tools").glob("*.yaml")): document = yaml.safe_load(path.read_text(encoding="utf-8")) if not isinstance(document, dict): raise ValueError("tool catalog document is invalid") spec = document.get("spec") if not isinstance(spec, dict): raise ValueError("tool catalog spec is invalid") name = spec.get("canonicalName") connector = spec.get("connectorType") risk_class = spec.get("riskClass") if not isinstance(name, str) or not name: raise ValueError("tool governance metadata is incomplete") if not isinstance(connector, str) or not connector: raise ValueError("tool governance metadata is incomplete") if not isinstance(risk_class, str) or not risk_class: raise ValueError("tool governance metadata is incomplete") published_name = f"itworx-{connector}-{name.replace('.', '-')}" configured_names = spec.get("exposedNames", []) if not isinstance(configured_names, list) or not all( isinstance(item, str) for item in configured_names ): raise ValueError("tool governance exposed names are invalid") for exposed_name in (name, published_name, *configured_names): if exposed_name in metadata: raise ValueError(f"duplicate tool governance metadata: {exposed_name}") metadata[exposed_name] = (connector, risk_class, name) if not metadata: raise ValueError("tool governance metadata is empty") return metadata class DatabaseAliasSource: def __init__(self, sessions: Any, gateway_url: str) -> None: self.sessions = sessions self.gateway_url = gateway_url def aliases(self) -> tuple[Any, ...]: with self.sessions() as session: workspace_id = session.scalar( select(WorkspaceRow.id).where(WorkspaceRow.slug == "itworx") ) if workspace_id is None: return () return SqlAlchemyAliasSource( session, workspace_id=workspace_id, gateway_base_url=self.gateway_url, ).aliases() class DatabaseMachineAuthenticator: _bearer = re.compile(r"^Bearer ([A-Za-z0-9._~+/=-]{1,65536})$") def __init__( self, providers: OIDCProviderClient | tuple[OIDCProviderClient, ...], sessions: Any, ) -> None: if not isinstance(providers, tuple): providers = (providers,) if not providers: raise ValueError("at least one OIDC provider is required") self.providers = providers self.sessions = sessions def __call__(self, authorization_header: str | None) -> Any: if authorization_header is None: raise InvalidToken() match = self._bearer.fullmatch(authorization_header.strip()) if match is None: raise InvalidToken() claims = None for provider in self.providers: try: claims = provider.validate_token(match.group(1)) break except InvalidToken: continue if claims is None: raise InvalidToken() with self.sessions() as session: client = session.scalar( select(ClientRow).where( ClientRow.issuer == claims.issuer, ClientRow.subject == claims.subject, ClientRow.status == "active", ) ) if client is None: raise UnknownMachinePrincipal() client.last_seen_at = datetime.now(UTC) session.commit() binding = MachineClientBinding( subject=client.subject, client_id=UUID(str(client.id)), enabled=True, ) return PrincipalMapper(machine_subjects={binding.subject: binding}).map_claims(claims) class CompositeAuthenticator: """Routes demo-token-shaped bearer credentials to the demo path, else OIDC.""" def __init__( self, demo: Callable[[str | None], Any], oidc: Callable[[str | None], Any] ) -> None: self.demo = demo self.oidc = oidc def __call__(self, authorization_header: str | None) -> Any: if authorization_header is not None and authorization_header.startswith( f"Bearer {TOKEN_PREFIX}" ): return self.demo(authorization_header) return self.oidc(authorization_header) class DemoAwareRateLimiter(SlidingWindowRateLimiter): """Applies a stricter shared ceiling to the one fixed demo client identity. Every other client_id behaves exactly as the base SlidingWindowRateLimiter it subclasses — this only special-cases the demo client so Governance's `rate_limiter: SlidingWindowRateLimiter | None` type still holds. """ def __init__( self, *, default_limit: int, demo_limit: int, demo_client_id: UUID, window_seconds: float = 60.0, ) -> None: super().__init__(limit=default_limit, window_seconds=window_seconds) self._demo_limiter = SlidingWindowRateLimiter( limit=demo_limit, window_seconds=window_seconds ) self._demo_client_id = demo_client_id def allow(self, *, client_id: UUID | str, project: str, tool: str) -> bool: if str(client_id) == str(self._demo_client_id): return self._demo_limiter.allow(client_id=client_id, project=project, tool=tool) return super().allow(client_id=client_id, project=project, tool=tool) def _oidc_providers() -> tuple[OIDCProviderClient, ...]: providers = [OIDCProviderClient(OIDCSettings.from_env())] additional = { "issuer": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_ISSUER", "").strip(), "audience": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_AUDIENCE", "").strip(), "client_id": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_CLIENT_ID", "").strip(), "redirect_uri": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_REDIRECT_URI", "").strip(), } configured = [bool(value) for value in additional.values()] if any(configured) and not all(configured): raise ValueError("additional MCP OIDC provider configuration is incomplete") if all(configured): settings = OIDCSettings.from_env( { "HUB_OIDC_ISSUER": additional["issuer"], "HUB_OIDC_AUDIENCE": additional["audience"], "HUB_OIDC_CLIENT_ID": additional["client_id"], "HUB_OIDC_REDIRECT_URI": additional["redirect_uri"], "HUB_OIDC_ALLOWED_ALGORITHMS": os.environ.get( "HUB_OIDC_ALLOWED_ALGORITHMS", "RS256" ), } ) providers.append(OIDCProviderClient(settings)) return tuple(providers) def build_app() -> FastAPI: oauth_metadata = OAuthMetadataSettings.from_environment() engine = create_engine_from_url(hub_database_url()) sessions = session_factory(engine) oidc_providers = _oidc_providers() oidc_provider = oidc_providers[0] contextforge = ContextForgeClient(ContextForgeSettings.from_env()) catalog_root = Path(os.environ.get("HUB_CATALOG_ROOT", "/opt/hub/catalog")).resolve() rate_limit = int(os.environ.get("HUB_RATE_LIMIT_PER_MINUTE", "600")) if not 1 <= rate_limit <= 10_000: raise ValueError("HUB_RATE_LIMIT_PER_MINUTE must be between 1 and 10000") demo_rate_limit = int(os.environ.get("DEMO_TOKEN_RATE_LIMIT_PER_MINUTE", "10")) if not 1 <= demo_rate_limit <= rate_limit: raise ValueError( "DEMO_TOKEN_RATE_LIMIT_PER_MINUTE must be between 1 and HUB_RATE_LIMIT_PER_MINUTE" ) governance = Governance( _policy_engine(catalog_root), SqlAlchemyAuditLog(sessions, issuer=str(oidc_provider.settings.issuer)), rate_limiter=DemoAwareRateLimiter( default_limit=rate_limit, demo_limit=demo_rate_limit, demo_client_id=DEMO_CLIENT_ID ), ) connector_identity_key = load_connector_identity_key() intake_key_path = os.environ.get("KNOWLEDGE_HUB_INGESTION_IDENTITY_KEY_FILE") intake_identity_key: bytes | None = None if intake_key_path: intake_identity_key = read_secret_file(Path(intake_key_path)).encode("utf-8") if len(intake_identity_key) < 32: raise ValueError("knowledge intake identity key must contain at least 32 bytes") if ( os.environ.get("INGESTION_CONNECTOR_ENABLED", "false").lower() == "true" and not intake_identity_key ): raise ValueError("knowledge intake identity key is required when intake is enabled") alias_cache_ttl = int(os.environ.get("HUB_ALIAS_CACHE_TTL_SECONDS", "5")) if not 1 <= alias_cache_ttl <= 60: raise ValueError("HUB_ALIAS_CACHE_TTL_SECONDS must be between 1 and 60") router_settings = RouterSettings( alias_cache_ttl_seconds=alias_cache_ttl, max_upstream_concurrency=int(os.environ.get("HUB_UPSTREAM_MAX_CONCURRENCY", "24")), upstream_queue_timeout_seconds=float( os.environ.get("HUB_UPSTREAM_QUEUE_TIMEOUT_SECONDS", "5") ), ) # Use ContextForge's unique control-network alias. The generic alias is # present on several Docker networks and can resolve outside the service # token's source-IP allowlist, causing intermittent 403 responses. source = DatabaseAliasSource(sessions, "http://contextforge-control:4444/mcp") registry = AliasRegistry(source, settings=router_settings) demo_revocations_path = Path( os.environ.get("DEMO_TOKEN_REVOCATIONS_FILE", "/etc/itworx/demo-token-revocations.json") ) authenticator = CompositeAuthenticator( demo=DemoTokenAuthenticator( None, demo_revocations_path, key_loader=load_demo_token_signing_key ), oidc=DatabaseMachineAuthenticator(oidc_providers, sessions), ) proxy = DataplaneProxy( registry, authenticator, recorder=JsonLineInvocationRecorder(), settings=router_settings, governance=governance, tool_metadata=_tool_governance_metadata(catalog_root), upstream_authorization=contextforge.authorization_header, identity_assertion=lambda project, environment, client_id, trace_id: identity_assertion( (intake_identity_key or b"") if project == "knowledge-hub" else connector_identity_key, project=project, environment=environment, client_id=client_id, trace_id=trace_id, ), mcp_extension=StaticSkills(serve_public_widgets=True), ) @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: try: yield finally: await proxy.close() for provider in oidc_providers: provider.close() contextforge.close() engine.dispose() app = FastAPI( title="ITWorx MCP Hub dataplane", docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan, ) @app.middleware("http") async def advertise_chatgpt_oauth( request: Request, call_next: Callable[[Request], Awaitable[Response]] ) -> Response: response = await call_next(request) path_parts = request.url.path.removeprefix("/mcp/").split("/") if ( len(path_parts) == 2 and all(part.replace("-", "").isalnum() for part in path_parts) and response.status_code == 401 ): resource_path = "/".join(path_parts) response.headers["WWW-Authenticate"] = bearer_challenge(resource_path, oauth_metadata) return response app.include_router(build_router(proxy)) return app def run() -> None: run_asgi(build_app())