"""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, )