| 1 | """Provider-neutral MCP contracts and safe result envelopes.""" |
| 2 | from __future__ import annotations |
| 3 | |
| 4 | import re |
| 5 | from dataclasses import dataclass |
| 6 | from datetime import datetime, timezone |
| 7 | from enum import StrEnum |
| 8 | from typing import Any, Awaitable, Callable |
| 9 | from uuid import UUID, uuid4 |
| 10 | |
| 11 | from pydantic import BaseModel, ConfigDict, Field, field_validator |
| 12 | |
| 13 | |
| 14 | class StrictContract(BaseModel): |
| 15 | model_config = ConfigDict(extra="forbid", populate_by_name=True) |
| 16 | |
| 17 | |
| 18 | class McpRiskClass(StrEnum): |
| 19 | SAFE_READ = "SAFE_READ" |
| 20 | SENSITIVE_READ = "SENSITIVE_READ" |
| 21 | WRITE_NON_FINANCIAL = "WRITE_NON_FINANCIAL" |
| 22 | FINANCIAL_ACTION = "FINANCIAL_ACTION" |
| 23 | ADMIN_ACTION = "ADMIN_ACTION" |
| 24 | |
| 25 | |
| 26 | class McpToolKind(StrEnum): |
| 27 | INTERNAL = "INTERNAL" |
| 28 | EXTERNAL = "EXTERNAL" |
| 29 | |
| 30 | |
| 31 | class McpAuthenticationType(StrEnum): |
| 32 | LOCAL_SERVICE = "LOCAL_SERVICE" |
| 33 | WORKLOAD_IDENTITY = "WORKLOAD_IDENTITY" |
| 34 | FEDERATED_USER = "FEDERATED_USER" |
| 35 | TEST = "TEST" |
| 36 | |
| 37 | |
| 38 | class McpErrorCode(StrEnum): |
| 39 | MCP_TOOL_NOT_FOUND = "MCP_TOOL_NOT_FOUND" |
| 40 | MCP_TOOL_DENIED = "MCP_TOOL_DENIED" |
| 41 | INVALID_ARGUMENT = "INVALID_ARGUMENT" |
| 42 | COMPANY_NOT_RESOLVED = "COMPANY_NOT_RESOLVED" |
| 43 | READINESS_NOT_AVAILABLE = "READINESS_NOT_AVAILABLE" |
| 44 | ANALYSIS_NOT_AVAILABLE = "ANALYSIS_NOT_AVAILABLE" |
| 45 | DOWNSTREAM_TIMEOUT = "DOWNSTREAM_TIMEOUT" |
| 46 | DOWNSTREAM_UNAVAILABLE = "DOWNSTREAM_UNAVAILABLE" |
| 47 | UNAUTHORIZED = "UNAUTHORIZED" |
| 48 | FORBIDDEN = "FORBIDDEN" |
| 49 | EXTERNAL_PROVIDER_UNAVAILABLE = "EXTERNAL_PROVIDER_UNAVAILABLE" |
| 50 | EXTERNAL_CAPABILITY_UNSUPPORTED = "EXTERNAL_CAPABILITY_UNSUPPORTED" |
| 51 | EXTERNAL_SCHEMA_INVALID = "EXTERNAL_SCHEMA_INVALID" |
| 52 | EXTERNAL_RESULT_INCOMPLETE = "EXTERNAL_RESULT_INCOMPLETE" |
| 53 | EXTERNAL_IDENTITY_CONFLICT = "EXTERNAL_IDENTITY_CONFLICT" |
| 54 | EXTERNAL_RESULT_STALE = "EXTERNAL_RESULT_STALE" |
| 55 | EXTERNAL_PROVIDER_RATE_LIMITED = "EXTERNAL_PROVIDER_RATE_LIMITED" |
| 56 | |
| 57 | |
| 58 | SAFE_ERROR_MESSAGES: dict[McpErrorCode, str] = { |
| 59 | McpErrorCode.MCP_TOOL_NOT_FOUND: "The requested MCP tool is not registered.", |
| 60 | McpErrorCode.MCP_TOOL_DENIED: "The requested MCP tool is not permitted by policy.", |
| 61 | McpErrorCode.INVALID_ARGUMENT: "The tool arguments are invalid.", |
| 62 | McpErrorCode.COMPANY_NOT_RESOLVED: "The canonical global instrument was not resolved.", |
| 63 | McpErrorCode.READINESS_NOT_AVAILABLE: "Research readiness is not available.", |
| 64 | McpErrorCode.ANALYSIS_NOT_AVAILABLE: "Company analysis is not available.", |
| 65 | McpErrorCode.DOWNSTREAM_TIMEOUT: "The internal application request timed out.", |
| 66 | McpErrorCode.DOWNSTREAM_UNAVAILABLE: "The internal application capability is unavailable.", |
| 67 | McpErrorCode.UNAUTHORIZED: "An authenticated MCP identity is required.", |
| 68 | McpErrorCode.FORBIDDEN: "The MCP identity is not authorized for this operation.", |
| 69 | McpErrorCode.EXTERNAL_PROVIDER_UNAVAILABLE: "No approved external MCP provider is available.", |
| 70 | McpErrorCode.EXTERNAL_CAPABILITY_UNSUPPORTED: "The approved external provider does not support this capability.", |
| 71 | McpErrorCode.EXTERNAL_SCHEMA_INVALID: "The external provider returned an invalid schema.", |
| 72 | McpErrorCode.EXTERNAL_RESULT_INCOMPLETE: "The external provider result did not satisfy the requirement.", |
| 73 | McpErrorCode.EXTERNAL_IDENTITY_CONFLICT: "The external provider result conflicts with canonical identity.", |
| 74 | McpErrorCode.EXTERNAL_RESULT_STALE: "The external provider result is outside the allowed freshness window.", |
| 75 | McpErrorCode.EXTERNAL_PROVIDER_RATE_LIMITED: "The external provider is temporarily rate limited.", |
| 76 | } |
| 77 | |
| 78 | |
| 79 | class McpAuthContext(StrictContract): |
| 80 | user_id: UUID | None = Field(default=None, alias="userId") |
| 81 | service_identity: str = Field(min_length=1, max_length=200, alias="serviceIdentity") |
| 82 | roles: tuple[str, ...] = () |
| 83 | scopes: tuple[str, ...] = () |
| 84 | authentication_type: McpAuthenticationType = Field(alias="authenticationType") |
| 85 | |
| 86 | @field_validator("roles", "scopes", mode="before") |
| 87 | @classmethod |
| 88 | def normalize_authorities(cls, value: Any) -> tuple[str, ...]: |
| 89 | if value is None: |
| 90 | return () |
| 91 | values = value.split(",") if isinstance(value, str) else value |
| 92 | return tuple(sorted({str(item).strip() for item in values if str(item).strip()})) |
| 93 | |
| 94 | def identity_headers(self) -> dict[str, str]: |
| 95 | headers = { |
| 96 | "X-AIP-Service-Identity": self.service_identity, |
| 97 | "X-AIP-Authentication-Type": self.authentication_type.value, |
| 98 | } |
| 99 | if self.user_id is not None: |
| 100 | user_id = str(self.user_id) |
| 101 | headers.update( |
| 102 | { |
| 103 | "X-AIP-User-Id": user_id, |
| 104 | "X-AIP-User-Issuer": self.service_identity, |
| 105 | "X-AIP-User-Subject": user_id, |
| 106 | } |
| 107 | ) |
| 108 | if self.roles: |
| 109 | headers["X-AIP-User-Roles"] = ",".join(self.roles) |
| 110 | return headers |
| 111 | |
| 112 | |
| 113 | class McpProvenance(StrictContract): |
| 114 | source: str = "INTERNAL_APPLICATION" |
| 115 | generated_at: datetime = Field(alias="generatedAt") |
| 116 | rule_engine_version: str | None = Field(default=None, alias="ruleEngineVersion") |
| 117 | |
| 118 | |
| 119 | class McpError(StrictContract): |
| 120 | code: McpErrorCode |
| 121 | message: str |
| 122 | |
| 123 | |
| 124 | class McpResultEnvelope(StrictContract): |
| 125 | ok: bool |
| 126 | tool: str |
| 127 | request_id: str = Field(alias="requestId") |
| 128 | data: Any | None = None |
| 129 | provenance: McpProvenance | None = None |
| 130 | warnings: list[str] = Field(default_factory=list) |
| 131 | error: McpError | None = None |
| 132 | |
| 133 | @classmethod |
| 134 | def success( |
| 135 | cls, |
| 136 | *, |
| 137 | tool: str, |
| 138 | request_id: str, |
| 139 | data: Any, |
| 140 | generated_at: datetime | None = None, |
| 141 | rule_engine_version: str | None = None, |
| 142 | warnings: list[str] | None = None, |
| 143 | ) -> "McpResultEnvelope": |
| 144 | return cls( |
| 145 | ok=True, |
| 146 | tool=tool, |
| 147 | requestId=request_id, |
| 148 | data=sanitize_response_data(data), |
| 149 | provenance=McpProvenance( |
| 150 | generatedAt=generated_at or datetime.now(timezone.utc), |
| 151 | ruleEngineVersion=rule_engine_version, |
| 152 | ), |
| 153 | warnings=warnings or [], |
| 154 | ) |
| 155 | |
| 156 | @classmethod |
| 157 | def failure( |
| 158 | cls, *, tool: str, request_id: str, code: McpErrorCode |
| 159 | ) -> "McpResultEnvelope": |
| 160 | return cls( |
| 161 | ok=False, |
| 162 | tool=tool, |
| 163 | requestId=request_id, |
| 164 | error=McpError(code=code, message=SAFE_ERROR_MESSAGES[code]), |
| 165 | ) |
| 166 | |
| 167 | def wire(self) -> dict[str, Any]: |
| 168 | return self.model_dump(mode="json", by_alias=True, exclude_none=True) |
| 169 | |
| 170 | |
| 171 | class McpGatewayError(Exception): |
| 172 | def __init__(self, code: McpErrorCode) -> None: |
| 173 | super().__init__(SAFE_ERROR_MESSAGES[code]) |
| 174 | self.code = code |
| 175 | |
| 176 | |
| 177 | @dataclass(frozen=True) |
| 178 | class McpToolExecution: |
| 179 | data: Any |
| 180 | generated_at: datetime | None = None |
| 181 | rule_engine_version: str | None = None |
| 182 | warnings: tuple[str, ...] = () |
| 183 | |
| 184 | |
| 185 | ToolHandler = Callable[[BaseModel, McpAuthContext, str], Awaitable[McpToolExecution]] |
| 186 | |
| 187 | |
| 188 | @dataclass(frozen=True) |
| 189 | class McpToolDefinition: |
| 190 | name: str |
| 191 | description: str |
| 192 | input_model: type[BaseModel] |
| 193 | risk_class: McpRiskClass |
| 194 | kind: McpToolKind |
| 195 | handler: ToolHandler |
| 196 | required_scopes: tuple[str, ...] = ("mcp:read",) |
| 197 | requires_user: bool = False |
| 198 | |
| 199 | |
| 200 | _REQUEST_ID_PATTERN = re.compile(r"^[A-Za-z0-9._:-]{1,128}$") |
| 201 | _SENSITIVE_KEYS = { |
| 202 | "quantity", |
| 203 | "averagecost", |
| 204 | "investedamount", |
| 205 | "pnl", |
| 206 | "pl", |
| 207 | "profitloss", |
| 208 | "profitandloss", |
| 209 | "costbasis", |
| 210 | "allocation", |
| 211 | "brokeraccountid", |
| 212 | "authorization", |
| 213 | "cookie", |
| 214 | "token", |
| 215 | "oauthtoken", |
| 216 | "apikey", |
| 217 | "api_key", |
| 218 | "oauthcode", |
| 219 | "oauthcredential", |
| 220 | } |
| 221 | _SENSITIVE_KEY_PARTS = ( |
| 222 | "password", |
| 223 | "token", |
| 224 | "secret", |
| 225 | "authorization", |
| 226 | "cookie", |
| 227 | "credential", |
| 228 | ) |
| 229 | _BEARER_PATTERN = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/-]+=*") |
| 230 | _JWT_PATTERN = re.compile(r"\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b") |
| 231 | _CREDENTIAL_ASSIGNMENT_PATTERN = re.compile( |
| 232 | r"(?i)\b(password|api[_-]?key|client[_-]?secret|access[_-]?token|refresh[_-]?token)" |
| 233 | r"\s*[:=]\s*[^\s&,;]+" |
| 234 | ) |
| 235 | |
| 236 | |
| 237 | def normalize_request_id(value: str | int | None) -> str: |
| 238 | candidate = "" if value is None else str(value).strip() |
| 239 | return candidate if _REQUEST_ID_PATTERN.fullmatch(candidate) else str(uuid4()) |
| 240 | |
| 241 | |
| 242 | def sanitize_response_data(value: Any) -> Any: |
| 243 | """Recursively remove private financial and credential-shaped response fields.""" |
| 244 | if isinstance(value, BaseModel): |
| 245 | value = value.model_dump(mode="json", by_alias=True, exclude_none=True) |
| 246 | if isinstance(value, dict): |
| 247 | result: dict[str, Any] = {} |
| 248 | for key, item in value.items(): |
| 249 | normalized = re.sub(r"[^a-z0-9_]", "", str(key).lower()) |
| 250 | compact = normalized.replace("_", "") |
| 251 | if normalized in _SENSITIVE_KEYS or compact in _SENSITIVE_KEYS: |
| 252 | continue |
| 253 | if any(part in compact for part in _SENSITIVE_KEY_PARTS): |
| 254 | continue |
| 255 | result[str(key)] = sanitize_response_data(item) |
| 256 | return result |
| 257 | if isinstance(value, (list, tuple)): |
| 258 | return [sanitize_response_data(item) for item in value] |
| 259 | if isinstance(value, str): |
| 260 | if _BEARER_PATTERN.search(value) or _JWT_PATTERN.search(value): |
| 261 | return "<redacted>" |
| 262 | return _CREDENTIAL_ASSIGNMENT_PATTERN.sub(r"\1=<redacted>", value) |
| 263 | return value |