978 lines
39 KiB
Python
978 lines
39 KiB
Python
"""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,
|
|
)
|