Initial curated ITWorx_MCP_Hub source from 368fba691ea8d23403b4d56b53b5ec670d7c5020

services/dataplane_router/proxy.py

@@ -0,0 +1,977 @@
+"""Bounded streaming proxy implementation; request credentials never enter logs."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import re
+import time
+from collections.abc import AsyncGenerator, Callable, Mapping
+from typing import Any, Protocol, cast
+from uuid import uuid4
+
+import httpx
+from fastapi import Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from packages.hub_policy import Governance, PolicyDenied, RateLimitExceeded
+
+from .authz import AliasAuthorizer
+from .errors import (
+    AUTHENTICATION_REQUIRED,
+    INTERNAL_ERROR,
+    OVERLOADED,
+    RESOURCE_NOT_REGISTERED,
+    RESULT_TOO_LARGE,
+    TARGET_UNAVAILABLE,
+    TIMEOUT,
+    TOKEN_INVALID,
+    VALIDATION_FAILED,
+    RouterError,
+)
+from .models import AliasRecord, EnvironmentName, InvocationMetadata, PrincipalLike, RouterSettings
+from .registry import AliasRegistry
+
+HOP_BY_HOP = frozenset(
+    {
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailer",
+        "transfer-encoding",
+        "upgrade",
+        "host",
+        "content-length",
+    }
+)
+AUTO_RESPONSE_HEADERS = frozenset({"date", "server"})
+TRACE_PATTERN = re.compile(r"^[A-Za-z0-9._:-]{1,128}$")
+ENVIRONMENTS = frozenset({"dev", "test", "staging", "prod"})
+IDENTITY_HEADERS = frozenset(
+    {
+        "authorization",
+        "x-itworx-project",
+        "x-itworx-environment",
+        "x-itworx-client-id",
+        "x-itworx-identity-assertion",
+    }
+)
+# Transport-session state must never cross the public Hub boundary. MCP
+# 2025-11-25 permits a Streamable HTTP server to omit Mcp-Session-Id entirely,
+# so the Hub can remain stateless even for initialize-handshake clients.
+MCP_TRANSPORT_STATE_HEADERS = frozenset({"mcp-session-id"})
+# These arrive from an untrusted client-side proxy hop.  The router must not
+# relay them into ContextForge, which applies its own source-network policy.
+# Forwarding them both lets callers spoof routing metadata and makes a valid
+# external MCP request appear to originate outside the trusted connector subnet.
+UNTRUSTED_FORWARDING_HEADERS = frozenset(
+    {
+        "forwarded",
+        "x-forwarded-for",
+        "x-forwarded-host",
+        "x-forwarded-port",
+        "x-forwarded-proto",
+        "x-forwarded-scheme",
+        "x-real-ip",
+    }
+)
+
+
+class InvocationRecorder:
+    def record(self, metadata: InvocationMetadata) -> None:
+        raise NotImplementedError
+
+
+class McpExtensionRouter(Protocol):
+    def handle_public_method(
+        self, project_slug: str, method: str, params: dict[str, Any]
+    ) -> dict[str, Any] | None: ...
+
+    def augment_initialize(self, project_slug: str, content: bytes, content_type: str) -> bytes: ...
+
+
+class InMemoryInvocationRecorder(InvocationRecorder):
+    def __init__(self) -> None:
+        self.records: list[InvocationMetadata] = []
+
+    def record(self, metadata: InvocationMetadata) -> None:
+        self.records.append(metadata)
+
+
+def _trace_id(value: str | None) -> str:
+    if value and TRACE_PATTERN.fullmatch(value):
+        return value
+    return uuid4().hex
+
+
+def _safe_headers(headers: Mapping[str, str], *, trace_id: str) -> dict[str, str]:
+    filtered = {
+        key: value
+        for key, value in headers.items()
+        if key.lower()
+        not in (
+            HOP_BY_HOP
+            | IDENTITY_HEADERS
+            | MCP_TRANSPORT_STATE_HEADERS
+            | UNTRUSTED_FORWARDING_HEADERS
+            | AUTO_RESPONSE_HEADERS
+        )
+    }
+    filtered["x-itworx-trace-id"] = trace_id
+    return filtered
+
+
+def _subject_digest(principal: PrincipalLike) -> str:
+    return hashlib.sha256(f"{principal.issuer}|{principal.subject}".encode()).hexdigest()
+
+
+def _mcp_metadata(body: bytes) -> tuple[str | None, str | None]:
+    if not body:
+        return None, None
+    try:
+        value = json.loads(body)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return None, None
+    if not isinstance(value, dict):
+        return None, None
+    method = value.get("method")
+    if not isinstance(method, str) or not re.fullmatch(r"[a-z][a-z0-9_/.-]{0,127}", method):
+        method = None
+    params = value.get("params")
+    name = params.get("name") if isinstance(params, dict) else None
+    if not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:/-]{0,255}", name):
+        name = None
+    return method, name
+
+
+def _jsonrpc_response_contract(body: bytes, request_id: Any) -> tuple[bool, bool | None, str]:
+    try:
+        value = json.loads(body)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return False, None, "unknown"
+    if not isinstance(value, dict) or value.get("jsonrpc") != "2.0":
+        return False, None, "unknown"
+    response_id_matches = value.get("id") == request_id
+    if "result" in value and "error" not in value:
+        return True, response_id_matches, "result"
+    if "error" in value and "result" not in value:
+        return True, response_id_matches, "error"
+    return True, response_id_matches, "unknown"
+
+
+def _rewrite_tool_list_names(content: bytes, aliases: Mapping[str, str]) -> bytes:
+    """Translate ContextForge's internal unique names to Hub canonical names."""
+
+    if not aliases:
+        return content
+    try:
+        document = json.loads(content)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return content
+    if not isinstance(document, dict):
+        return content
+    result = document.get("result")
+    if not isinstance(result, dict):
+        return content
+    tools = result.get("tools")
+    if not isinstance(tools, list):
+        return content
+    changed = False
+    rewritten: list[Any] = []
+    for tool in tools:
+        if not isinstance(tool, dict):
+            rewritten.append(tool)
+            continue
+        name = tool.get("name")
+        canonical = aliases.get(name) if isinstance(name, str) else None
+        if canonical is None:
+            rewritten.append(tool)
+            continue
+        changed = True
+        rewritten.append({**tool, "name": canonical})
+    if not changed:
+        return content
+    return json.dumps(
+        {**document, "result": {**result, "tools": rewritten}},
+        separators=(",", ":"),
+    ).encode()
+
+
+def _resolve_local_schema_reference(
+    root_schema: Mapping[str, Any], schema: object
+) -> dict[str, Any] | None:
+    """Resolve a chain of local JSON Schema references without remote lookups."""
+
+    if not isinstance(schema, dict):
+        return None
+    current = schema
+    visited: set[str] = set()
+    while "$ref" in current:
+        # Ref siblings may add constraints. Decline to infer their meaning here;
+        # leaving metadata absent is safer than attesting to a partial schema.
+        if set(current) != {"$ref"}:
+            return None
+        reference = current.get("$ref")
+        if not isinstance(reference, str) or not reference.startswith("#/") or reference in visited:
+            return None
+        visited.add(reference)
+        target: Any = root_schema
+        for raw_token in reference[2:].split("/"):
+            token_parts: list[str] = []
+            index = 0
+            while index < len(raw_token):
+                character = raw_token[index]
+                if character == "~":
+                    if index + 1 >= len(raw_token) or raw_token[index + 1] not in "01":
+                        return None
+                    token_parts.append("/" if raw_token[index + 1] == "1" else "~")
+                    index += 2
+                else:
+                    token_parts.append(character)
+                    index += 1
+            token = "".join(token_parts)
+            if isinstance(target, Mapping):
+                if token not in target:
+                    return None
+                target = target[token]
+            elif isinstance(target, list):
+                if re.fullmatch(r"(?:0|[1-9][0-9]*)", token) is None:
+                    return None
+                try:
+                    array_index = int(token)
+                except ValueError:
+                    return None
+                if array_index >= len(target):
+                    return None
+                target = target[array_index]
+            else:
+                return None
+        if not isinstance(target, dict):
+            return None
+        current = target
+    return current
+
+
+def _add_knowledge_hub_file_params(content: bytes) -> bytes:
+    """Restore ChatGPT's file binding for the two vetted attachment tools."""
+
+    try:
+        document = json.loads(content)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return content
+    if not isinstance(document, dict):
+        return content
+    result = document.get("result")
+    if not isinstance(result, dict):
+        return content
+    tools = result.get("tools")
+    if not isinstance(tools, list):
+        return content
+
+    changed = False
+    rewritten: list[Any] = []
+    for tool in tools:
+        if not isinstance(tool, dict) or tool.get("name") not in {
+            "analyze_knowledge_files",
+            "extract_knowledge_files",
+        }:
+            rewritten.append(tool)
+            continue
+        raw_input_schema = tool.get("inputSchema")
+        input_schema = _resolve_local_schema_reference(
+            raw_input_schema if isinstance(raw_input_schema, dict) else {}, raw_input_schema
+        )
+        if input_schema is None:
+            rewritten.append(tool)
+            continue
+        properties = input_schema.get("properties")
+        raw_files_schema = properties.get("files") if isinstance(properties, dict) else None
+        files_schema = _resolve_local_schema_reference(input_schema, raw_files_schema)
+        if not isinstance(files_schema, dict) or files_schema.get("type") != "array":
+            rewritten.append(tool)
+            continue
+        file_item = _resolve_local_schema_reference(input_schema, files_schema.get("items"))
+        file_properties = file_item.get("properties") if isinstance(file_item, dict) else None
+        required = file_item.get("required") if isinstance(file_item, dict) else None
+        expected_properties = {"download_url", "file_id", "mime_type", "file_name"}
+        if (
+            not isinstance(file_item, dict)
+            or file_item.get("type") != "object"
+            or not isinstance(file_properties, dict)
+            or set(file_properties) != expected_properties
+            or any(
+                not isinstance(value, dict) or value.get("type") != "string"
+                for value in file_properties.values()
+            )
+            or not isinstance(required, list)
+            or len(required) != 2
+            or not all(isinstance(field, str) for field in required)
+            or set(required) != {"download_url", "file_id"}
+        ):
+            rewritten.append(tool)
+            continue
+        metadata = tool.get("_meta")
+        metadata = dict(metadata) if isinstance(metadata, dict) else {}
+        metadata["openai/fileParams"] = ["files"]
+        rewritten.append({**tool, "_meta": metadata})
+        changed = True
+    if not changed:
+        return content
+    return json.dumps(
+        {**document, "result": {**result, "tools": rewritten}},
+        separators=(",", ":"),
+    ).encode()
+
+
+def _contextforge_tool_name(connector: str, exposed_name: str) -> str:
+    slug = re.sub(r"[^a-z0-9]+", "-", exposed_name.casefold()).strip("-")
+    return f"itworx-{connector}-{slug}"
+
+
+_RETRYABLE_MCP_METHODS = frozenset(
+    {
+        "initialize",
+        "ping",
+        "tools/list",
+        "resources/list",
+        "resources/read",
+        "prompts/list",
+        "prompts/get",
+    }
+)
+_RETRYABLE_UPSTREAM_STATUSES = frozenset({502, 503, 504})
+
+
+class DataplaneProxy:
+    def __init__(
+        self,
+        registry: AliasRegistry,
+        authenticator: Callable[[str | None], PrincipalLike],
+        *,
+        authorizer: AliasAuthorizer | None = None,
+        recorder: InvocationRecorder | None = None,
+        settings: RouterSettings | None = None,
+        transport: httpx.AsyncBaseTransport | None = None,
+        governance: Governance | None = None,
+        tool_metadata: Mapping[str, tuple[str, str] | tuple[str, str, str]] | None = None,
+        upstream_authorization: Callable[[], str] | None = None,
+        identity_assertion: Callable[[str, str, str, str], str] | None = None,
+        mcp_extension: McpExtensionRouter | None = None,
+    ) -> None:
+        self.registry = registry
+        self.authenticator = authenticator
+        self.authorizer = authorizer or AliasAuthorizer()
+        self.recorder = recorder or InMemoryInvocationRecorder()
+        self.settings = settings or registry.settings
+        self.governance = governance
+        self.tool_metadata = {
+            exposed_name: (
+                item[0],
+                item[1],
+                item[2] if len(item) == 3 else exposed_name,
+            )
+            for exposed_name, item in (tool_metadata or {}).items()
+        }
+        self.contextforge_tool_aliases = {
+            _contextforge_tool_name(connector, exposed_name): exposed_name
+            for exposed_name, (connector, _, _) in self.tool_metadata.items()
+        }
+        self.upstream_authorization = upstream_authorization
+        self.identity_assertion = identity_assertion
+        self.mcp_extension = mcp_extension
+        # ContextForge sees the router as one source IP, so its generic
+        # per-IP limiter cannot fairly protect multi-client MCP traffic.  The
+        # public router owns admission control and keeps a bounded queue here.
+        self._upstream_slots = asyncio.BoundedSemaphore(self.settings.max_upstream_concurrency)
+        self.client = httpx.AsyncClient(
+            timeout=httpx.Timeout(self.settings.timeout_seconds),
+            transport=transport,
+        )
+
+    async def close(self) -> None:
+        await self.client.aclose()
+
+    async def handle(
+        self, request: Request, project_slug: str, environment: str
+    ) -> StreamingResponse | JSONResponse:
+        trace_id = _trace_id(request.headers.get("x-trace-id"))
+        started = time.monotonic()
+        governance_context = None
+        governance_trace = None
+        governance_started = None
+        capacity_acquired = False
+        mcp_method: str | None = None
+        mcp_name: str | None = None
+        retry_safe = False
+        request_document: dict[str, Any] | None = None
+        request_id: Any = None
+        try:
+            if (
+                not re.fullmatch(r"[a-z0-9][a-z0-9-]{0,62}", project_slug)
+                or environment not in ENVIRONMENTS
+            ):
+                raise RESOURCE_NOT_REGISTERED
+            record = self.registry.resolve(project_slug, cast(EnvironmentName, environment))
+            if record is None:
+                raise RESOURCE_NOT_REGISTERED
+            auth_header = request.headers.get("authorization")
+            if not auth_header:
+                raise AUTHENTICATION_REQUIRED
+            try:
+                principal = self.authenticator(auth_header)
+            except Exception as exc:
+                code = getattr(exc, "code", None)
+                if code in {
+                    "invalid_token",
+                    "unknown_machine_principal",
+                    "revoked_machine_principal",
+                }:
+                    raise TOKEN_INVALID from exc
+                raise AUTHENTICATION_REQUIRED from exc
+            protocol_revision = self._protocol_revision(request)
+            content_length = request.headers.get("content-length")
+            if content_length and (
+                not content_length.isdigit() or int(content_length) > self.settings.max_body_bytes
+            ):
+                raise VALIDATION_FAILED
+            body = await self._read_body(request)
+            if body:
+                try:
+                    candidate = json.loads(body)
+                except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+                    raise VALIDATION_FAILED from exc
+                if not isinstance(candidate, dict):
+                    # Batch requests cannot be assigned one risk class safely. A
+                    # mixed R1/R3 batch would otherwise inherit the generic R1
+                    # request classification and bypass per-tool governance.
+                    raise VALIDATION_FAILED
+                request_document = candidate
+                request_id = candidate.get("id")
+            mcp_method, mcp_name = _mcp_metadata(body)
+            if self.governance is not None:
+                connector = "dataplane-router"
+                risk_class = "R1"
+                policy_tool_name = mcp_name
+                if mcp_method == "tools/call":
+                    metadata = self.tool_metadata.get(mcp_name or "")
+                    if metadata is None:
+                        connector, risk_class = "unknown", "RX"
+                    else:
+                        connector, risk_class, policy_tool_name = metadata
+                binding_present = any(
+                    binding.enabled
+                    and binding.client_id == principal.client_id
+                    and binding.issuer == principal.issuer
+                    and binding.subject == principal.subject
+                    for binding in record.bindings
+                )
+                governance_context = self.governance.context(
+                    workspace=str(record.workspace_id or "unknown"),
+                    project=record.project_slug,
+                    environment=record.environment,
+                    principal=principal.subject,
+                    client_id=principal.client_id,
+                    binding_present=binding_present,
+                    client_enabled=binding_present and record.enabled,
+                    tool=(policy_tool_name or f"mcp.{mcp_method}" if mcp_method else "mcp.request"),
+                    connector=connector,
+                    risk_class=risk_class,
+                    arguments={
+                        "method": request.method,
+                        "mcpMethod": mcp_method,
+                        "mcpName": mcp_name,
+                    },
+                    resource_type="tool" if mcp_method == "tools/call" else "mcp-request",
+                    environment_classification=(
+                        "production" if record.environment == "prod" else "development"
+                    ),
+                    trace_id=trace_id,
+                )
+                governance_trace, governance_started = self.governance.begin(governance_context)
+            self.authorizer.authorize(principal, record)
+            if self.mcp_extension is not None and mcp_method and request_document is not None:
+                params = request_document.get("params", {})
+                if not isinstance(params, dict):
+                    params = {}
+                try:
+                    local_result = self.mcp_extension.handle_public_method(
+                        record.project_slug, mcp_method, params
+                    )
+                    local_payload = (
+                        {
+                            "jsonrpc": "2.0",
+                            "id": request_document.get("id"),
+                            "result": local_result,
+                        }
+                        if local_result is not None
+                        else None
+                    )
+                except ValueError:
+                    local_payload = {
+                        "jsonrpc": "2.0",
+                        "id": request_document.get("id"),
+                        "error": {"code": -32602, "message": "Invalid extension parameters"},
+                    }
+                if local_payload is not None:
+                    response = httpx.Response(
+                        200,
+                        headers={"content-type": "application/json"},
+                        content=json.dumps(local_payload, separators=(",", ":")).encode(),
+                    )
+                    return StreamingResponse(
+                        self._response_stream(
+                            response,
+                            trace_id=trace_id,
+                            record=record,
+                            principal=principal,
+                            protocol_revision=protocol_revision,
+                            method=request.method,
+                            request_bytes=len(body),
+                            request_id=request_id,
+                            mcp_method=mcp_method,
+                            mcp_name=mcp_name,
+                            started=started,
+                            governance=self.governance,
+                            governance_context=governance_context,
+                            governance_trace=governance_trace,
+                            governance_started=governance_started,
+                        ),
+                        status_code=200,
+                        headers={
+                            "content-type": "application/json",
+                            "content-length": str(len(response.content)),
+                            "x-itworx-trace-id": trace_id,
+                        },
+                        media_type=None,
+                    )
+            forwarded_headers = _safe_headers(request.headers, trace_id=trace_id)
+            forwarded_headers["x-itworx-project"] = record.project_slug
+            forwarded_headers["x-itworx-environment"] = record.environment
+            forwarded_headers["x-itworx-client-id"] = str(principal.client_id)
+            if self.identity_assertion is None:
+                raise RuntimeError("connector identity assertion is unavailable")
+            forwarded_headers["x-itworx-identity-assertion"] = self.identity_assertion(
+                record.project_slug,
+                record.environment,
+                str(principal.client_id),
+                trace_id,
+            )
+            if self.upstream_authorization is not None:
+                forwarded_headers["authorization"] = self.upstream_authorization()
+            try:
+                await asyncio.wait_for(
+                    self._upstream_slots.acquire(),
+                    timeout=self.settings.upstream_queue_timeout_seconds,
+                )
+            except TimeoutError as exc:
+                raise OVERLOADED from exc
+            capacity_acquired = True
+            retry_safe = self._retry_safe(mcp_method, mcp_name)
+            response, retry_count = await self._send_buffered(
+                request.method,
+                str(record.upstream_url),
+                headers=forwarded_headers,
+                body=body,
+                retry_safe=retry_safe,
+            )
+            if self.mcp_extension is not None and mcp_method == "initialize":
+                transformed = self.mcp_extension.augment_initialize(
+                    record.project_slug,
+                    response.content,
+                    response.headers.get("content-type", ""),
+                )
+                response = httpx.Response(
+                    response.status_code,
+                    headers=response.headers,
+                    content=transformed,
+                    request=response.request,
+                )
+            elif mcp_method == "tools/list":
+                transformed = response.content
+                if self.contextforge_tool_aliases:
+                    transformed = _rewrite_tool_list_names(
+                        transformed, self.contextforge_tool_aliases
+                    )
+                if record.project_slug == "knowledge-hub" and record.environment == "prod":
+                    transformed = _add_knowledge_hub_file_params(transformed)
+                response = httpx.Response(
+                    response.status_code,
+                    headers=response.headers,
+                    content=transformed,
+                    request=response.request,
+                )
+            if (
+                response.headers.get("content-length", "").isdigit()
+                and int(response.headers["content-length"]) > self.settings.max_response_bytes
+            ):
+                await response.aclose()
+                raise RESULT_TOO_LARGE from None
+            # The upstream body is fully buffered, so slow downstream clients
+            # no longer need to reserve scarce upstream request capacity.
+            self._upstream_slots.release()
+            capacity_acquired = False
+            streaming_response = StreamingResponse(
+                self._response_stream(
+                    response,
+                    trace_id=trace_id,
+                    record=record,
+                    principal=principal,
+                    protocol_revision=protocol_revision,
+                    method=request.method,
+                    request_bytes=len(body),
+                    request_id=request_id,
+                    mcp_method=mcp_method,
+                    mcp_name=mcp_name,
+                    started=started,
+                    governance=self.governance,
+                    governance_context=governance_context,
+                    governance_trace=governance_trace,
+                    governance_started=governance_started,
+                    release_upstream_capacity=False,
+                ),
+                status_code=response.status_code,
+                headers={
+                    **_safe_headers(response.headers, trace_id=trace_id),
+                    # _send_buffered has already consumed and bounded the complete
+                    # decoded entity. Give downstream HTTP clients an unambiguous
+                    # message boundary instead of turning this fixed body back into
+                    # an open-ended chunked stream.
+                    "content-length": str(len(response.content)),
+                    "x-itworx-upstream-retries": str(retry_count),
+                },
+                media_type=None,
+            )
+            return streaming_response
+        except RouterError as exc:
+            self._finish_governance_error(
+                governance_context, governance_trace, governance_started, exc.code
+            )
+            return self._error(
+                exc,
+                trace_id,
+                operation=mcp_name or mcp_method,
+                safe_to_retry=exc is OVERLOADED or (retry_safe and exc.retryable),
+            )
+        except PolicyDenied as exc:
+            return JSONResponse(
+                status_code=403,
+                content={
+                    "error": {
+                        "code": "POLICY_DENIED",
+                        "message": "The policy denied this invocation.",
+                        "reason": exc.trace.reason_code,
+                        "traceId": trace_id,
+                    }
+                },
+                headers={"x-trace-id": trace_id},
+            )
+        except RateLimitExceeded:
+            return JSONResponse(
+                status_code=429,
+                content={
+                    "error": {
+                        "code": "RATE_LIMITED",
+                        "message": "The invocation rate limit was exceeded.",
+                        "retryable": True,
+                        "traceId": trace_id,
+                    }
+                },
+                headers={"x-trace-id": trace_id},
+            )
+        except httpx.TimeoutException:
+            self._finish_governance_error(
+                governance_context, governance_trace, governance_started, TIMEOUT.code
+            )
+            return self._error(
+                TIMEOUT,
+                trace_id,
+                operation=mcp_name or mcp_method,
+                safe_to_retry=retry_safe,
+            )
+        except httpx.HTTPError:
+            self._finish_governance_error(
+                governance_context, governance_trace, governance_started, TARGET_UNAVAILABLE.code
+            )
+            return self._error(
+                TARGET_UNAVAILABLE,
+                trace_id,
+                operation=mcp_name or mcp_method,
+                safe_to_retry=retry_safe,
+            )
+        except Exception:
+            self._finish_governance_error(
+                governance_context, governance_trace, governance_started, INTERNAL_ERROR.code
+            )
+            return self._error(
+                INTERNAL_ERROR,
+                trace_id,
+                operation=mcp_name or mcp_method,
+                safe_to_retry=False,
+            )
+        finally:
+            if capacity_acquired:
+                self._upstream_slots.release()
+
+    def _finish_governance_error(
+        self, context: Any, trace: Any, started: Any, error_code: str
+    ) -> None:
+        if self.governance is not None and context is not None and trace is not None:
+            self.governance.finish(
+                context,
+                trace,
+                started,
+                outcome="error",
+                error_code=error_code,
+            )
+
+    def _protocol_revision(self, request: Request) -> str:
+        value = request.headers.get("mcp-protocol-version")
+        if value is None:
+            return "2025-11-25"
+        if value not in {"2026-07-28", "2025-11-25"}:
+            raise VALIDATION_FAILED
+        return value
+
+    async def _read_body(self, request: Request) -> bytes:
+        chunks: list[bytes] = []
+        total = 0
+        async for chunk in request.stream():
+            total += len(chunk)
+            if total > self.settings.max_body_bytes:
+                raise VALIDATION_FAILED
+            chunks.append(chunk)
+        return b"".join(chunks)
+
+    def _retry_safe(self, mcp_method: str | None, mcp_name: str | None) -> bool:
+        if mcp_method in _RETRYABLE_MCP_METHODS:
+            return True
+        if mcp_method != "tools/call" or mcp_name is None:
+            return False
+        metadata = self.tool_metadata.get(mcp_name)
+        return metadata is not None and metadata[1] == "R1"
+
+    async def _send_buffered(
+        self,
+        method: str,
+        url: str,
+        *,
+        headers: Mapping[str, str],
+        body: bytes,
+        retry_safe: bool,
+    ) -> tuple[httpx.Response, int]:
+        """Buffer a bounded response before exposing it to the MCP client.
+
+        This prevents a late upstream stream failure from tearing down a client
+        response after its status and headers have already been committed.
+        """
+
+        # The public client's compression preferences belong to the external
+        # hop.  Reusing them for the internal ContextForge hop can negotiate a
+        # codec that the router's httpx runtime does not decode (notably br or
+        # zstd).  The representation header is stripped after buffering, so an
+        # undecoded entity would otherwise be exposed as invalid JSON.  Keep
+        # the internal hop deterministic; the public proxy may compress the
+        # final, correctly framed response independently.
+        upstream_headers = dict(headers)
+        upstream_headers["accept-encoding"] = "identity"
+
+        attempts = 2 if retry_safe else 1
+        last_error: httpx.HTTPError | None = None
+        for attempt in range(attempts):
+            upstream_request = self.client.build_request(
+                method,
+                url,
+                headers=upstream_headers,
+                content=body,
+            )
+            response: httpx.Response | None = None
+            try:
+                response = await self.client.send(upstream_request, stream=True)
+                content = await self._read_upstream_response(response)
+                response_headers = response.headers.copy()
+                # ``aiter_bytes`` returns the decoded entity. Forwarding the
+                # upstream representation headers would make httpx (and later
+                # the client) decode the already-decoded body a second time.
+                # Length and integrity metadata are also stale after MCP
+                # response rewriting.
+                for header in (
+                    "content-encoding",
+                    "content-length",
+                    "content-md5",
+                    "digest",
+                    "etag",
+                    "transfer-encoding",
+                ):
+                    response_headers.pop(header, None)
+                buffered = httpx.Response(
+                    response.status_code,
+                    headers=response_headers,
+                    content=content,
+                    request=upstream_request,
+                )
+                if attempt + 1 < attempts and buffered.status_code in _RETRYABLE_UPSTREAM_STATUSES:
+                    await buffered.aclose()
+                    continue
+                return buffered, attempt
+            except (httpx.TimeoutException, httpx.NetworkError) as exc:
+                last_error = exc
+                if attempt + 1 >= attempts:
+                    raise
+            finally:
+                if response is not None:
+                    await response.aclose()
+        if last_error is not None:  # pragma: no cover - loop always raises or returns
+            raise last_error
+        raise RuntimeError("upstream retry loop ended without a response")
+
+    async def _read_upstream_response(self, response: httpx.Response) -> bytes:
+        chunks: list[bytes] = []
+        total = 0
+        try:
+            # Bound the decoded entity exposed to the MCP client. Bounding raw
+            # compressed bytes would permit a small compression bomb to exceed
+            # the public output budget after decompression.
+            async for chunk in response.aiter_bytes():
+                total += len(chunk)
+                if total > self.settings.max_response_bytes:
+                    raise RESULT_TOO_LARGE from None
+                chunks.append(chunk)
+        except httpx.StreamConsumed:
+            content = response.content
+            if len(content) > self.settings.max_response_bytes:
+                raise RESULT_TOO_LARGE from None
+            return content
+        return b"".join(chunks)
+
+    async def _response_stream(
+        self,
+        response: httpx.Response,
+        *,
+        trace_id: str,
+        record: AliasRecord,
+        principal: PrincipalLike,
+        protocol_revision: str,
+        method: str,
+        request_bytes: int,
+        request_id: Any,
+        mcp_method: str | None,
+        mcp_name: str | None,
+        started: float,
+        governance: Governance | None = None,
+        governance_context: Any = None,
+        governance_trace: Any = None,
+        governance_started: Any = None,
+        release_upstream_capacity: bool = False,
+    ) -> AsyncGenerator[bytes]:
+        response_bytes = 0
+        stream_error: str | None = None
+        try:
+            try:
+                async for chunk in response.aiter_raw():
+                    response_bytes += len(chunk)
+                    if response_bytes > self.settings.max_response_bytes:
+                        raise RESULT_TOO_LARGE from None
+                    yield chunk
+            except httpx.StreamConsumed:
+                chunk = response.content
+                response_bytes = len(chunk)
+                if response_bytes > self.settings.max_response_bytes:
+                    raise RESULT_TOO_LARGE from None
+                yield chunk
+        except Exception as exc:
+            stream_error = getattr(exc, "code", type(exc).__name__)
+            raise
+        finally:
+            try:
+                await response.aclose()
+            finally:
+                # Never leak a capacity slot, even when an upstream transport
+                # raises while it is being closed after cancellation.
+                if release_upstream_capacity:
+                    self._upstream_slots.release()
+            if (
+                governance is not None
+                and governance_context is not None
+                and governance_trace is not None
+                and governance_started is not None
+            ):
+                governance.finish(
+                    governance_context,
+                    governance_trace,
+                    governance_started,
+                    outcome=(
+                        "error"
+                        if stream_error is not None or response.status_code >= 400
+                        else "success"
+                    ),
+                    error_code=stream_error
+                    or (
+                        f"UPSTREAM_{response.status_code}" if response.status_code >= 400 else None
+                    ),
+                )
+            duration_ms = min(120_000, int((time.monotonic() - started) * 1000))
+            if principal.client_id is not None:
+                jsonrpc_response, response_id_matches, response_shape = _jsonrpc_response_contract(
+                    response.content, request_id
+                )
+                self.recorder.record(
+                    InvocationMetadata(
+                        trace_id=trace_id,
+                        project_slug=record.project_slug,
+                        environment=record.environment,
+                        client_id=principal.client_id,
+                        subject_digest=_subject_digest(principal),
+                        method=method,
+                        mcp_method=mcp_method,
+                        mcp_name=mcp_name,
+                        protocol_revision=protocol_revision,
+                        request_bytes=request_bytes,
+                        response_status=response.status_code,
+                        response_bytes=min(response_bytes, self.settings.max_response_bytes),
+                        response_content_type=response.headers.get("content-type"),
+                        jsonrpc_response=jsonrpc_response,
+                        response_id_matches=response_id_matches,
+                        response_shape=response_shape,
+                        duration_ms=duration_ms,
+                    )
+                )
+
+    @staticmethod
+    def _error(
+        error: RouterError,
+        trace_id: str,
+        *,
+        operation: str | None = None,
+        safe_to_retry: bool = False,
+    ) -> JSONResponse:
+        headers = {"x-trace-id": trace_id}
+        if error is OVERLOADED:
+            headers["retry-after"] = "1"
+        return JSONResponse(
+            status_code=error.status_code,
+            content={
+                "error": {
+                    "code": error.code,
+                    "category": "transient" if error.retryable else "permanent",
+                    "message": error.message,
+                    "operation": operation,
+                    "subsystem": "dataplane-router",
+                    "retryable": error.retryable,
+                    "safeToRetry": safe_to_retry,
+                    "traceId": trace_id,
+                    "sessionId": None,
+                    "upstreamStatus": None,
+                    "expectedState": None,
+                    "actualState": None,
+                    "recoveryHint": (
+                        "Retry the same operation; the stateless MCP session remains usable."
+                        if safe_to_retry
+                        else "Correct the request or authorization before retrying."
+                    ),
+                }
+            },
+            headers=headers,
+        )