Files

154 lines
5.6 KiB
Python

"""Typed, public-safe contracts for the dataplane router."""
from __future__ import annotations
import ipaddress
import re
from datetime import datetime
from typing import Literal, Protocol
from urllib.parse import urlparse
from uuid import UUID
from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, field_validator, model_validator
EnvironmentName = Literal["dev", "test", "staging", "prod"]
SUPPORTED_PROTOCOLS = frozenset({"2026-07-28", "2025-11-25"})
class ClientBinding(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
client_id: UUID
issuer: str = Field(min_length=1, max_length=2048)
subject: str = Field(min_length=1, max_length=512)
enabled: bool = False
class AliasRecord(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
project_slug: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,62}$")
workspace_id: UUID | None = None
project_id: UUID | None = None
environment: EnvironmentName
alias: str
upstream_url: AnyHttpUrl
enabled: bool = False
protocol_revisions: tuple[str, ...] = ("2026-07-28", "2025-11-25")
bindings: tuple[ClientBinding, ...] = ()
@field_validator("alias")
@classmethod
def validate_alias(cls, value: str) -> str:
if len(value) > 128 or not re.fullmatch(
r"/mcp/[a-z0-9][a-z0-9-]{0,62}/(?:dev|test|staging|prod)/?", value
):
raise ValueError("alias must be a stable /mcp path")
return value.rstrip("/")
@field_validator("upstream_url")
@classmethod
def reject_upstream_userinfo(cls, value: AnyHttpUrl) -> AnyHttpUrl:
if value.username is not None or value.password is not None:
raise ValueError("upstream URL userinfo is not permitted")
internal_contextforge = (
value.scheme == "http"
and value.host in {"contextforge", "contextforge-control"}
and value.port == 4444
and (
(value.path or "").rstrip("/") == "/mcp"
or bool(
re.fullmatch(
r"/servers/[0-9a-f]{32}/mcp",
(value.path or "").rstrip("/"),
)
)
)
)
if value.scheme != "https" and not internal_contextforge:
raise ValueError("dataplane upstreams must use HTTPS or the fixed internal gateway")
hostname = (urlparse(str(value)).hostname or "").casefold()
if hostname in {"localhost", "metadata", "metadata.google.internal"}:
raise ValueError("dataplane upstream host is not allowed")
try:
address = ipaddress.ip_address(hostname)
except ValueError:
address = None
if address is not None and (
address.is_private
or address.is_loopback
or address.is_link_local
or address.is_reserved
or address.is_unspecified
or address.is_multicast
):
raise ValueError("dataplane upstream host is not allowed")
return value
@field_validator("protocol_revisions")
@classmethod
def validate_protocols(cls, value: tuple[str, ...]) -> tuple[str, ...]:
if not value or not set(value).issubset(SUPPORTED_PROTOCOLS):
raise ValueError("unsupported MCP protocol revision")
return tuple(dict.fromkeys(value))
@model_validator(mode="after")
def validate_identity(self) -> AliasRecord:
if not self.alias.endswith(f"/{self.environment}"):
raise ValueError("alias environment does not match the record environment")
if len({binding.client_id for binding in self.bindings}) != len(self.bindings):
raise ValueError("duplicate client binding")
return self
class InvocationMetadata(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
trace_id: str = Field(pattern=r"^[A-Za-z0-9._:-]{1,128}$")
project_slug: str
environment: EnvironmentName
client_id: UUID
subject_digest: str = Field(pattern=r"^[0-9a-f]{64}$")
method: str = Field(min_length=1, max_length=16)
mcp_method: str | None = Field(default=None, max_length=128)
mcp_name: str | None = Field(default=None, max_length=256)
protocol_revision: str
request_bytes: int = Field(ge=0, le=4 * 1024 * 1024)
response_status: int = Field(ge=100, le=599)
response_bytes: int = Field(ge=0, le=4 * 1024 * 1024)
response_content_type: str | None = Field(default=None, max_length=128)
jsonrpc_response: bool = False
response_id_matches: bool | None = None
response_shape: str = Field(default="unknown", pattern=r"^(result|error|unknown)$")
duration_ms: int = Field(ge=0, le=120_000)
class RouterSettings(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
max_body_bytes: int = Field(default=4 * 1024 * 1024, ge=1024, le=4 * 1024 * 1024)
max_response_bytes: int = Field(default=4 * 1024 * 1024, ge=1024, le=4 * 1024 * 1024)
timeout_seconds: float = Field(default=30.0, gt=0, le=30)
alias_cache_ttl_seconds: int = Field(default=30, ge=1, le=60)
max_upstream_concurrency: int = Field(default=24, ge=1, le=256)
upstream_queue_timeout_seconds: float = Field(default=5.0, ge=0.1, le=30)
class PrincipalLike(Protocol):
"""Minimal validated principal contract consumed by the dataplane."""
@property
def subject(self) -> str: ...
@property
def issuer(self) -> str: ...
@property
def kind(self) -> object: ...
@property
def client_id(self) -> UUID | None: ...
@property
def expires_at(self) -> datetime: ...