| 1 | """Provider-neutral future external MCP registry and explicit fallback gate.""" |
| 2 | from __future__ import annotations |
| 3 | |
| 4 | import asyncio |
| 5 | from datetime import datetime, timezone |
| 6 | from enum import StrEnum |
| 7 | from time import monotonic |
| 8 | from typing import Any, Literal, Protocol |
| 9 | from uuid import UUID |
| 10 | |
| 11 | from pydantic import Field, field_validator |
| 12 | |
| 13 | from app.contracts import ( |
| 14 | McpErrorCode, |
| 15 | McpAuthContext, |
| 16 | McpGatewayError, |
| 17 | McpRiskClass, |
| 18 | McpToolKind, |
| 19 | StrictContract, |
| 20 | sanitize_response_data, |
| 21 | ) |
| 22 | from app.audit import McpAuditSink |
| 23 | |
| 24 | |
| 25 | class ExternalCapabilityState(StrEnum): |
| 26 | SUPPORTED = "SUPPORTED" |
| 27 | UNSUPPORTED = "UNSUPPORTED" |
| 28 | UNKNOWN = "UNKNOWN" |
| 29 | TEMPORARILY_UNAVAILABLE = "TEMPORARILY_UNAVAILABLE" |
| 30 | |
| 31 | |
| 32 | class ExternalProviderHealthState(StrEnum): |
| 33 | UNKNOWN = "UNKNOWN" |
| 34 | UP = "UP" |
| 35 | DEGRADED = "DEGRADED" |
| 36 | DOWN = "DOWN" |
| 37 | |
| 38 | |
| 39 | class ExternalMcpProviderCapability(StrictContract): |
| 40 | region: str |
| 41 | requirement_id: str = Field(alias="requirementId", min_length=1) |
| 42 | capability: str = Field(min_length=1) |
| 43 | tool: str | None = None |
| 44 | state: ExternalCapabilityState = ExternalCapabilityState.UNKNOWN |
| 45 | minimum_items: int = Field(default=1, ge=1, le=1000, alias="minimumItems") |
| 46 | max_age_seconds: int | None = Field(default=None, ge=1, le=31_536_000, alias="maxAgeSeconds") |
| 47 | required_fields: tuple[str, ...] = Field(default=(), alias="requiredFields") |
| 48 | tool_arguments: dict[str, Any] = Field(default_factory=dict, alias="toolArguments") |
| 49 | |
| 50 | @field_validator("region", "requirement_id", "capability") |
| 51 | @classmethod |
| 52 | def normalize_identifier(cls, value: str) -> str: |
| 53 | return value.strip().upper() |
| 54 | |
| 55 | @field_validator("tool") |
| 56 | @classmethod |
| 57 | def normalize_tool(cls, value: str | None) -> str | None: |
| 58 | return value.strip() if value and value.strip() else None |
| 59 | |
| 60 | @field_validator("required_fields", mode="before") |
| 61 | @classmethod |
| 62 | def normalize_fields(cls, value) -> tuple[str, ...]: |
| 63 | return tuple(str(item).strip() for item in (value or ()) if str(item).strip()) |
| 64 | |
| 65 | |
| 66 | class ExternalMcpProviderMetadata(StrictContract): |
| 67 | provider_id: str = Field(alias="providerId", min_length=1) |
| 68 | kind: Literal[McpToolKind.EXTERNAL] = McpToolKind.EXTERNAL |
| 69 | regions: tuple[str, ...] |
| 70 | supported_requirements: tuple[str, ...] = Field(alias="supportedRequirements") |
| 71 | supported_tools: tuple[str, ...] = Field(alias="supportedTools") |
| 72 | risk_class: McpRiskClass = Field(alias="riskClass") |
| 73 | auth_type: str = Field(alias="authType", min_length=1) |
| 74 | priority: int = Field(default=100, ge=1, le=1000) |
| 75 | source_tier: str = Field(default="APPROVED_EXTERNAL_TOOL", alias="sourceTier") |
| 76 | health_state: ExternalProviderHealthState = Field( |
| 77 | default=ExternalProviderHealthState.UNKNOWN, alias="healthState" |
| 78 | ) |
| 79 | timeout_seconds: float = Field(default=10.0, gt=0, le=120, alias="timeoutSeconds") |
| 80 | timeout_behavior: str = Field(default="BOUNDED_FAIL_CLOSED", alias="timeoutBehavior") |
| 81 | adapter_version: str = Field(default="UNSPECIFIED", alias="adapterVersion") |
| 82 | |
| 83 | @field_validator("provider_id", "auth_type", "source_tier", "timeout_behavior") |
| 84 | @classmethod |
| 85 | def upper_identifier(cls, value: str) -> str: |
| 86 | return value.strip().upper() |
| 87 | |
| 88 | @field_validator("regions", "supported_requirements", mode="before") |
| 89 | @classmethod |
| 90 | def upper_values(cls, value): |
| 91 | return tuple(str(item).strip().upper() for item in value) |
| 92 | |
| 93 | @field_validator("supported_tools", mode="before") |
| 94 | @classmethod |
| 95 | def exact_tool_names(cls, value): |
| 96 | return tuple(str(item).strip() for item in value) |
| 97 | |
| 98 | |
| 99 | class ProviderFallbackAuthorization(StrictContract): |
| 100 | authorized: bool |
| 101 | issued_by: str = Field(alias="issuedBy") |
| 102 | global_instrument_id: UUID = Field(alias="globalInstrumentId") |
| 103 | requirement_id: str = Field(alias="requirementId") |
| 104 | permitted_provider_ids: tuple[str, ...] = Field(alias="permittedProviderIds") |
| 105 | issued_at: datetime = Field(alias="issuedAt") |
| 106 | |
| 107 | @field_validator("issued_at") |
| 108 | @classmethod |
| 109 | def aware_time(cls, value: datetime) -> datetime: |
| 110 | if value.tzinfo is None or value.utcoffset() is None: |
| 111 | raise ValueError("issuedAt must be timezone-aware") |
| 112 | return value.astimezone(timezone.utc) |
| 113 | |
| 114 | |
| 115 | class ExternalMcpProvider(Protocol): |
| 116 | metadata: ExternalMcpProviderMetadata |
| 117 | |
| 118 | async def supports(self, *, region: str, requirement_id: str, tool: str) -> bool: ... |
| 119 | |
| 120 | def capability_for( |
| 121 | self, *, region: str, requirement_id: str |
| 122 | ) -> ExternalMcpProviderCapability | None: ... |
| 123 | |
| 124 | async def invoke( |
| 125 | self, *, tool: str, arguments: dict[str, Any], request_id: str |
| 126 | ) -> dict[str, Any]: ... |
| 127 | |
| 128 | async def health(self) -> dict[str, Any]: ... |
| 129 | |
| 130 | |
| 131 | class McpServerRegistry: |
| 132 | def __init__(self) -> None: |
| 133 | self._providers: dict[str, ExternalMcpProvider] = {} |
| 134 | |
| 135 | def register(self, provider: ExternalMcpProvider) -> None: |
| 136 | provider_id = provider.metadata.provider_id |
| 137 | if provider_id in self._providers: |
| 138 | raise ValueError(f"External MCP provider already registered: {provider_id}") |
| 139 | self._providers[provider_id] = provider |
| 140 | |
| 141 | def get(self, provider_id: str) -> ExternalMcpProvider | None: |
| 142 | return self._providers.get(provider_id.strip().upper()) |
| 143 | |
| 144 | @property |
| 145 | def count(self) -> int: |
| 146 | return len(self._providers) |
| 147 | |
| 148 | @property |
| 149 | def provider_ids(self) -> tuple[str, ...]: |
| 150 | return tuple(sorted(self._providers)) |
| 151 | |
| 152 | |
| 153 | class ExternalMcpGateway: |
| 154 | """Future routing seam; no provider is registered by the production container in 5A.""" |
| 155 | |
| 156 | def __init__( |
| 157 | self, |
| 158 | registry: McpServerRegistry, |
| 159 | *, |
| 160 | enabled: bool = False, |
| 161 | timeout_seconds: float = 10, |
| 162 | audit: McpAuditSink | None = None, |
| 163 | auth: McpAuthContext | None = None, |
| 164 | ) -> None: |
| 165 | if timeout_seconds <= 0: |
| 166 | raise ValueError("timeout_seconds must be positive") |
| 167 | self.registry = registry |
| 168 | self.enabled = enabled |
| 169 | self.timeout_seconds = timeout_seconds |
| 170 | self.audit = audit |
| 171 | self.auth = auth |
| 172 | |
| 173 | async def invoke_requirement( |
| 174 | self, |
| 175 | *, |
| 176 | provider_id: str, |
| 177 | region: str, |
| 178 | requirement_id: str, |
| 179 | global_instrument_id: UUID, |
| 180 | arguments: dict[str, Any], |
| 181 | request_id: str, |
| 182 | authorization: ProviderFallbackAuthorization | None, |
| 183 | ) -> dict[str, Any]: |
| 184 | """Resolve a configured capability so callers cannot choose arbitrary provider tools.""" |
| 185 | try: |
| 186 | self._authorize( |
| 187 | provider_id=provider_id, |
| 188 | requirement_id=requirement_id, |
| 189 | global_instrument_id=global_instrument_id, |
| 190 | authorization=authorization, |
| 191 | ) |
| 192 | except McpGatewayError as exc: |
| 193 | self._audit( |
| 194 | "MCP_TOOL_DENIED", |
| 195 | f"{provider_id}:{requirement_id}", |
| 196 | request_id, |
| 197 | code=exc.code, |
| 198 | ) |
| 199 | raise |
| 200 | provider = self.registry.get(provider_id) |
| 201 | if provider is None: |
| 202 | self._audit( |
| 203 | "MCP_TOOL_FAILED", |
| 204 | f"{provider_id}:{requirement_id}", |
| 205 | request_id, |
| 206 | code=McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE, |
| 207 | ) |
| 208 | raise McpGatewayError(McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE) |
| 209 | resolver = getattr(provider, "capability_for", None) |
| 210 | capability = resolver(region=region, requirement_id=requirement_id) if resolver else None |
| 211 | if ( |
| 212 | capability is None |
| 213 | or capability.state is not ExternalCapabilityState.SUPPORTED |
| 214 | or not capability.tool |
| 215 | ): |
| 216 | self._audit( |
| 217 | "MCP_TOOL_FAILED", |
| 218 | f"{provider_id}:{requirement_id}", |
| 219 | request_id, |
| 220 | code=McpErrorCode.EXTERNAL_CAPABILITY_UNSUPPORTED, |
| 221 | ) |
| 222 | raise McpGatewayError(McpErrorCode.EXTERNAL_CAPABILITY_UNSUPPORTED) |
| 223 | return await self.invoke( |
| 224 | provider_id=provider_id, |
| 225 | region=region, |
| 226 | requirement_id=requirement_id, |
| 227 | tool=capability.tool, |
| 228 | global_instrument_id=global_instrument_id, |
| 229 | arguments=arguments, |
| 230 | request_id=request_id, |
| 231 | authorization=authorization, |
| 232 | ) |
| 233 | |
| 234 | async def invoke( |
| 235 | self, |
| 236 | *, |
| 237 | provider_id: str, |
| 238 | region: str, |
| 239 | requirement_id: str, |
| 240 | tool: str, |
| 241 | global_instrument_id: UUID, |
| 242 | arguments: dict[str, Any], |
| 243 | request_id: str, |
| 244 | authorization: ProviderFallbackAuthorization | None, |
| 245 | ) -> dict[str, Any]: |
| 246 | started = monotonic() |
| 247 | try: |
| 248 | self._authorize( |
| 249 | provider_id=provider_id, |
| 250 | requirement_id=requirement_id, |
| 251 | global_instrument_id=global_instrument_id, |
| 252 | authorization=authorization, |
| 253 | ) |
| 254 | except McpGatewayError as exc: |
| 255 | self._audit("MCP_TOOL_DENIED", tool, request_id, code=exc.code) |
| 256 | raise |
| 257 | provider = self.registry.get(provider_id) |
| 258 | if not self.enabled or provider is None: |
| 259 | self._audit( |
| 260 | "MCP_TOOL_FAILED", |
| 261 | tool, |
| 262 | request_id, |
| 263 | code=McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE, |
| 264 | ) |
| 265 | raise McpGatewayError(McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE) |
| 266 | if provider.metadata.risk_class is not McpRiskClass.SAFE_READ: |
| 267 | self._audit("MCP_TOOL_DENIED", tool, request_id, code=McpErrorCode.MCP_TOOL_DENIED) |
| 268 | raise McpGatewayError(McpErrorCode.MCP_TOOL_DENIED) |
| 269 | self._audit("MCP_TOOL_INVOKED", tool, request_id) |
| 270 | try: |
| 271 | async with asyncio.timeout(self.timeout_seconds): |
| 272 | if not await provider.supports( |
| 273 | region=region, requirement_id=requirement_id, tool=tool |
| 274 | ): |
| 275 | raise McpGatewayError(McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE) |
| 276 | result = await provider.invoke( |
| 277 | tool=tool, |
| 278 | arguments={ |
| 279 | **arguments, |
| 280 | "globalInstrumentId": str(global_instrument_id), |
| 281 | "region": region.strip().upper(), |
| 282 | "requirementId": requirement_id.strip().upper(), |
| 283 | }, |
| 284 | request_id=request_id, |
| 285 | ) |
| 286 | except TimeoutError as exc: |
| 287 | self._audit( |
| 288 | "MCP_TOOL_FAILED", |
| 289 | tool, |
| 290 | request_id, |
| 291 | code=McpErrorCode.DOWNSTREAM_TIMEOUT, |
| 292 | duration_ms=round((monotonic() - started) * 1000), |
| 293 | ) |
| 294 | raise McpGatewayError(McpErrorCode.DOWNSTREAM_TIMEOUT) from exc |
| 295 | except McpGatewayError as exc: |
| 296 | self._audit( |
| 297 | "MCP_TOOL_FAILED", |
| 298 | tool, |
| 299 | request_id, |
| 300 | code=exc.code, |
| 301 | duration_ms=round((monotonic() - started) * 1000), |
| 302 | ) |
| 303 | raise |
| 304 | except Exception as exc: |
| 305 | self._audit( |
| 306 | "MCP_TOOL_FAILED", |
| 307 | tool, |
| 308 | request_id, |
| 309 | code=McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE, |
| 310 | duration_ms=round((monotonic() - started) * 1000), |
| 311 | ) |
| 312 | raise McpGatewayError(McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE) from exc |
| 313 | try: |
| 314 | self._enforce_identity_boundary(result, global_instrument_id) |
| 315 | except McpGatewayError as exc: |
| 316 | self._audit( |
| 317 | "MCP_TOOL_FAILED", |
| 318 | tool, |
| 319 | request_id, |
| 320 | code=exc.code, |
| 321 | duration_ms=round((monotonic() - started) * 1000), |
| 322 | ) |
| 323 | raise |
| 324 | normalized = sanitize_response_data(result) |
| 325 | self._audit( |
| 326 | "MCP_TOOL_SUCCEEDED", |
| 327 | tool, |
| 328 | request_id, |
| 329 | duration_ms=round((monotonic() - started) * 1000), |
| 330 | ) |
| 331 | return normalized |
| 332 | |
| 333 | def _audit( |
| 334 | self, |
| 335 | event: str, |
| 336 | tool: str, |
| 337 | request_id: str, |
| 338 | *, |
| 339 | code: McpErrorCode | None = None, |
| 340 | duration_ms: int | None = None, |
| 341 | ) -> None: |
| 342 | if self.audit is None or self.auth is None: |
| 343 | return |
| 344 | self.audit.emit( |
| 345 | event, |
| 346 | tool=tool, |
| 347 | request_id=request_id, |
| 348 | auth=self.auth, |
| 349 | risk_class=McpRiskClass.SAFE_READ, |
| 350 | kind=McpToolKind.EXTERNAL, |
| 351 | code=code.value if code else None, |
| 352 | duration_ms=duration_ms, |
| 353 | ) |
| 354 | |
| 355 | @staticmethod |
| 356 | def _authorize( |
| 357 | *, provider_id, requirement_id, global_instrument_id, authorization |
| 358 | ) -> None: |
| 359 | if ( |
| 360 | authorization is None |
| 361 | or not authorization.authorized |
| 362 | or authorization.issued_by != "ProviderFallbackPolicy" |
| 363 | or authorization.global_instrument_id != global_instrument_id |
| 364 | or authorization.requirement_id.upper() != requirement_id.upper() |
| 365 | or provider_id.upper() |
| 366 | not in {value.upper() for value in authorization.permitted_provider_ids} |
| 367 | ): |
| 368 | raise McpGatewayError(McpErrorCode.FORBIDDEN) |
| 369 | |
| 370 | @staticmethod |
| 371 | def _enforce_identity_boundary(result: dict[str, Any], global_instrument_id: UUID) -> None: |
| 372 | forbidden = {"canonicalidentity", "providermapping", "providermappings", "createidentity"} |
| 373 | |
| 374 | def inspect(value: Any) -> None: |
| 375 | if isinstance(value, dict): |
| 376 | for key, item in value.items(): |
| 377 | normalized = "".join(char for char in str(key).lower() if char.isalnum()) |
| 378 | if normalized in forbidden: |
| 379 | raise McpGatewayError(McpErrorCode.FORBIDDEN) |
| 380 | if normalized == "globalinstrumentid" and str(item) != str( |
| 381 | global_instrument_id |
| 382 | ): |
| 383 | raise McpGatewayError(McpErrorCode.FORBIDDEN) |
| 384 | inspect(item) |
| 385 | elif isinstance(value, (list, tuple)): |
| 386 | for item in value: |
| 387 | inspect(item) |
| 388 | |
| 389 | inspect(result) |