from __future__ import annotations import asyncio import gzip import json import tempfile import time import unittest from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Literal from uuid import UUID, uuid4 import httpx from fastapi import FastAPI from fastapi.testclient import TestClient from pydantic import AnyHttpUrl, TypeAdapter from packages.hub_policy import Decision, Governance, InMemoryAuditLog, PolicyEngine, PolicyRule from services.control_api.auth.models import AuthenticatedPrincipal, PrincipalKind from services.dataplane_router import ( AliasRecord, AliasRegistry, ClientBinding, DataplaneProxy, InMemoryInvocationRecorder, RouterSettings, StaticAliasSource, build_router, ) from services.dataplane_router.authz import AliasAuthorizer class StaticAuthenticator: def __init__(self, principal: AuthenticatedPrincipal | None) -> None: self.principal = principal self.seen_headers: list[str | None] = [] def __call__(self, header: str | None) -> AuthenticatedPrincipal: self.seen_headers.append(header) if self.principal is None: from services.control_api.auth.errors import InvalidToken raise InvalidToken() return self.principal class ConnectorTransport: def __init__(self) -> None: self.requests: list[httpx.Request] = [] self.mode = "ok" async def handle(self, request: httpx.Request) -> httpx.Response: self.requests.append(request) if self.mode == "timeout": raise httpx.ReadTimeout("simulated timeout") payload = await request.aread() return httpx.Response( 200, headers={"content-type": "application/json", "connection": "close"}, content=json.dumps( { "jsonrpc": "2.0", "id": 1, "result": { "protocol": request.headers.get("mcp-protocol-version", "compat"), "bodyBytes": len(payload), }, } ).encode(), request=request, ) def principal(client_id: UUID, subject: str = "svc-ludarium") -> AuthenticatedPrincipal: now = datetime.now(UTC) return AuthenticatedPrincipal( subject=subject, issuer="https://fake.oidc.test", audience=("itworx-mcp-hub",), kind=PrincipalKind.MACHINE_CLIENT, client_id=client_id, machine_subject=subject, expires_at=now + timedelta(minutes=5), issued_at=now, ) def record( client_id: UUID, *, project: str = "ludarium", environment: Literal["dev", "test", "staging", "prod"] = "prod", ) -> AliasRecord: return AliasRecord( project_slug=project, environment=environment, alias=f"/mcp/{project}/{environment}", upstream_url=TypeAdapter(AnyHttpUrl).validate_python("https://gateway.invalid/registered"), enabled=True, bindings=( ClientBinding( client_id=client_id, issuer="https://fake.oidc.test", subject="svc-ludarium", enabled=True, ), ), ) class DataplaneRouterTests(unittest.TestCase): def setUp(self) -> None: self.client_id = uuid4() self.connector = ConnectorTransport() self.authenticator = StaticAuthenticator(principal(self.client_id)) self.recorder = InMemoryInvocationRecorder() self.registry = AliasRegistry( StaticAliasSource((record(self.client_id),)), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), ) self.proxy = DataplaneProxy( self.registry, self.authenticator, recorder=self.recorder, transport=httpx.MockTransport(self.connector.handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), identity_assertion=lambda project, environment, client_id, trace_id: ( f"signed:{project}:{environment}:{client_id}:{trace_id}" ), ) self.app = FastAPI() self.app.include_router(build_router(self.proxy)) self.http = TestClient(self.app) def tearDown(self) -> None: asyncio.run(self.proxy.close()) def test_native_and_compatibility_protocols_reach_registered_connector(self) -> None: native = self.http.post( "/mcp/ludarium/prod", headers={ "Authorization": "Bearer opaque-token", "Mcp-Protocol-Version": "2026-07-28", "Mcp-Method": "tools/list", "Mcp-Name": "ludarium", "X-ITWorx-Identity-Assertion": "attacker-value", "X-Forwarded-For": "198.51.100.10", "X-Forwarded-Proto": "https", "X-Real-IP": "198.51.100.10", }, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) compat = self.http.post( "/mcp/ludarium/prod", headers={ "Authorization": "Bearer opaque-token", "Mcp-Protocol-Version": "2025-11-25", }, content=b"{}", ) self.assertEqual(native.status_code, 200) self.assertEqual(compat.status_code, 200) self.assertEqual( [request.headers["mcp-protocol-version"] for request in self.connector.requests], ["2026-07-28", "2025-11-25"], ) self.assertEqual( self.authenticator.seen_headers, ["Bearer opaque-token", "Bearer opaque-token"] ) self.assertEqual(len(self.recorder.records), 2) self.assertEqual(self.recorder.records[0].mcp_method, "tools/list") self.assertTrue(self.recorder.records[0].jsonrpc_response) self.assertTrue(self.recorder.records[0].response_id_matches) self.assertEqual(self.recorder.records[0].response_shape, "result") forwarded = self.connector.requests[0].headers self.assertEqual(forwarded["x-itworx-project"], "ludarium") self.assertEqual(forwarded["x-itworx-environment"], "prod") self.assertEqual(forwarded["x-itworx-client-id"], str(self.client_id)) self.assertTrue(forwarded["x-itworx-identity-assertion"].startswith("signed:ludarium:")) self.assertNotIn("authorization", forwarded) self.assertNotIn("x-forwarded-for", forwarded) self.assertNotIn("x-forwarded-proto", forwarded) self.assertNotIn("x-real-ip", forwarded) def test_knowledge_hub_prod_tools_list_advertises_file_params(self) -> None: item_schema = { "type": "object", "properties": { "download_url": {"type": "string"}, "file_id": {"type": "string"}, "mime_type": {"type": "string"}, "file_name": {"type": "string"}, }, "required": ["download_url", "file_id"], } file_schema = { "type": "object", "$defs": {"OpenAIFile": item_schema}, "properties": { "files": { "type": "array", "items": {"$ref": "#/$defs/OpenAIFile"}, } }, } preserve_vendor_metadata = False async def handle(request: httpx.Request) -> httpx.Response: nonlocal preserve_vendor_metadata metadata = ( {"vendor": "preserved", "openai/fileParams": ["wrong"]} if preserve_vendor_metadata else None ) document = { "jsonrpc": "2.0", "id": 7, "result": { "tools": [ { "name": "itworx-ingestion-analyze-knowledge-files", "description": "Analyze attached knowledge files.", "inputSchema": file_schema, "_meta": metadata, "extensionMetadata": None, }, { "name": "search_knowledge", "inputSchema": {"type": "object", "properties": {}}, "_meta": {"vendor": "unchanged"}, }, ], "nextCursor": None, }, } return httpx.Response(200, json=document, request=request) client_registry = AliasRegistry( StaticAliasSource( ( record(self.client_id, project="knowledge-hub"), record(self.client_id, project="knowledge-hub", environment="dev"), record(self.client_id), ) ), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), ) proxy = DataplaneProxy( client_registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), tool_metadata={"analyze_knowledge_files": ("ingestion", "R1")}, identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/knowledge-hub/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 7, "method": "tools/list"}, ) self.assertEqual(response.status_code, 200) public_document = response.json() tools = public_document["result"]["tools"] analyze, search = tools self.assertEqual(analyze["name"], "analyze_knowledge_files") self.assertEqual(analyze["_meta"], {"openai/fileParams": ["files"]}) self.assertIsNone(analyze["extensionMetadata"]) self.assertEqual(search["_meta"], {"vendor": "unchanged"}) self.assertEqual(public_document["result"]["nextCursor"], None) preserve_vendor_metadata = True response = client.post( "/mcp/knowledge-hub/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 7, "method": "tools/list"}, ) self.assertEqual( response.json()["result"]["tools"][0]["_meta"], {"vendor": "preserved", "openai/fileParams": ["files"]}, ) for endpoint in ("/mcp/knowledge-hub/dev", "/mcp/ludarium/prod"): response = client.post( endpoint, headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 7, "method": "tools/list"}, ) self.assertEqual(response.status_code, 200) self.assertEqual( response.json()["result"]["tools"][0]["_meta"], {"vendor": "preserved", "openai/fileParams": ["wrong"]}, ) finally: asyncio.run(proxy.close()) def test_file_params_requires_the_connector_file_object_schema(self) -> None: from services.dataplane_router.proxy import _add_knowledge_hub_file_params without_files_array = { "jsonrpc": "2.0", "id": 1, "result": { "tools": [ { "name": "analyze_knowledge_files", "inputSchema": { "type": "object", "properties": { "files": { "type": "array", "items": { "type": "object", "properties": {"download_url": {"type": "string"}}, "required": ["download_url"], }, } }, }, "_meta": {"vendor": "preserved"}, } ] }, } encoded = json.dumps(without_files_array).encode() self.assertEqual(_add_knowledge_hub_file_params(encoded), encoded) def test_file_params_metadata_is_restored_only_for_the_two_attachment_tools(self) -> None: from services.dataplane_router.proxy import _add_knowledge_hub_file_params item_schema = { "type": "object", "additionalProperties": False, "properties": { "download_url": {"type": "string"}, "file_id": {"type": "string"}, "mime_type": {"type": "string"}, "file_name": {"type": "string"}, }, "required": ["download_url", "file_id"], } document = { "jsonrpc": "2.0", "id": 1, "result": { "tools": [ { "name": "extract_knowledge_files", "inputSchema": { "type": "object", "properties": {"files": {"type": "array", "items": item_schema}}, }, "_meta": {"vendor": "keep"}, }, { "name": "unregistered_tool", "inputSchema": { "type": "object", "properties": {"files": {"type": "array", "items": item_schema}}, }, }, ] }, } rewritten = json.loads(_add_knowledge_hub_file_params(json.dumps(document).encode())) extract, unrelated = rewritten["result"]["tools"] self.assertEqual( extract["_meta"], {"vendor": "keep", "openai/fileParams": ["files"]}, ) self.assertNotIn("_meta", unrelated) def test_file_params_metadata_resolves_local_item_references_and_fails_closed(self) -> None: from services.dataplane_router.proxy import _add_knowledge_hub_file_params item_schema = { "type": "object", "additionalProperties": False, "properties": { "download_url": {"type": "string"}, "file_id": {"type": "string"}, "mime_type": {"type": "string"}, "file_name": {"type": "string"}, }, "required": ["download_url", "file_id"], } document = { "jsonrpc": "2.0", "id": 1, "result": { "tools": [ { "name": "extract_knowledge_files", "inputSchema": { "type": "object", "$defs": {"Attachment": item_schema}, "properties": { "files": { "type": "array", "items": {"$ref": "#/$defs/Attachment"}, } }, }, "_meta": {"vendor": "preserved"}, }, { "name": "analyze_knowledge_files", "inputSchema": { "type": "object", "properties": { "files": { "type": "array", "items": {"$ref": "https://example.test/Attachment"}, } }, }, }, ] }, } rewritten = json.loads(_add_knowledge_hub_file_params(json.dumps(document).encode())) extract, analyze = rewritten["result"]["tools"] self.assertEqual( extract["_meta"], {"vendor": "preserved", "openai/fileParams": ["files"]}, ) self.assertNotIn("_meta", analyze) def test_file_params_metadata_resolves_local_array_index_references(self) -> None: from services.dataplane_router.proxy import _add_knowledge_hub_file_params item_schema = { "type": "object", "properties": { "download_url": {"type": "string"}, "file_id": {"type": "string"}, "mime_type": {"type": "string"}, "file_name": {"type": "string"}, }, "required": ["download_url", "file_id"], } document = { "jsonrpc": "2.0", "id": 1, "result": { "tools": [ { "name": "extract_knowledge_files", "inputSchema": { "type": "object", "allOf": [{"$defs": {"Attachment": item_schema}}], "properties": { "files": { "type": "array", "items": {"$ref": "#/allOf/0/$defs/Attachment"}, } }, }, } ] }, } rewritten = json.loads(_add_knowledge_hub_file_params(json.dumps(document).encode())) extract = rewritten["result"]["tools"][0] self.assertEqual(extract["_meta"], {"openai/fileParams": ["files"]}) def test_tool_list_exposes_canonical_names_while_preserving_unknown_tools(self) -> None: class PrefixedToolTransport: async def handle(self, request: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "application/json"}, json={ "jsonrpc": "2.0", "id": 1, "result": { "tools": [ {"name": "itworx-gitea-gitea-repository-create"}, {"name": "future-tool"}, ] }, }, request=request, ) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(PrefixedToolTransport().handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), tool_metadata={"gitea.repository.create": ("gitea", "R2")}, identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(response.status_code, 200) self.assertEqual( [tool["name"] for tool in response.json()["result"]["tools"]], ["gitea.repository.create", "future-tool"], ) finally: asyncio.run(proxy.close()) def test_capacity_queue_fails_fast_with_retry_guidance(self) -> None: constrained = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(self.connector.handle), settings=RouterSettings( max_body_bytes=1024, max_response_bytes=4096, max_upstream_concurrency=1, upstream_queue_timeout_seconds=0.1, ), identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) async def exercise() -> None: await constrained._upstream_slots.acquire() app = FastAPI() app.include_router(build_router(constrained)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(response.status_code, 503) self.assertEqual(response.json()["error"]["code"], "OVERLOADED") self.assertEqual(response.headers["retry-after"], "1") finally: constrained._upstream_slots.release() await constrained.close() asyncio.run(exercise()) def test_request_burst_is_bounded_without_rejections(self) -> None: class LoadTransport: def __init__(self) -> None: self.active = 0 self.maximum = 0 async def handle(self, request: httpx.Request) -> httpx.Response: self.active += 1 self.maximum = max(self.maximum, self.active) try: await asyncio.sleep(0.02) return httpx.Response( 200, headers={"content-type": "application/json"}, content=b'{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}', request=request, ) finally: self.active -= 1 load = LoadTransport() constrained = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(load.handle), settings=RouterSettings( max_body_bytes=1024, max_response_bytes=4096, max_upstream_concurrency=4, upstream_queue_timeout_seconds=2, ), identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) async def exercise() -> None: app = FastAPI() app.include_router(build_router(constrained)) async with httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://test" ) as client: responses = await asyncio.gather( *( client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": index, "method": "tools/list"}, ) for index in range(40) ) ) self.assertTrue(all(response.status_code == 200 for response in responses)) self.assertEqual(load.maximum, 4) await constrained.close() asyncio.run(exercise()) def test_gitea_skills_and_app_templates_are_served_by_authenticated_dataplane(self) -> None: from services.runtime.mcp_skills import StaticSkills skills = StaticSkills( Path(__file__).resolve().parents[1] / "plugins/itworx-gitea/skills", serve_public_widgets=True, ) registry = AliasRegistry(StaticAliasSource((record(self.client_id, project="gitea-all"),))) proxy = DataplaneProxy( registry, self.authenticator, recorder=self.recorder, transport=httpx.MockTransport(self.connector.handle), identity_assertion=lambda project, environment, client_id, trace_id: "signed-fixture", mcp_extension=skills, ) app = FastAPI() app.include_router(build_router(proxy)) client = TestClient(app) headers = { "Authorization": "Bearer opaque-token", "Mcp-Protocol-Version": "2026-07-28", } initialized = client.post( "/mcp/gitea-all/prod", headers=headers, json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, ) listed = client.post( "/mcp/gitea-all/prod", headers=headers, json={"jsonrpc": "2.0", "id": 2, "method": "skills/list", "params": {}}, ) uri = listed.json()["result"]["skills"][0]["resources"][0]["uri"] resource = client.post( "/mcp/gitea-all/prod", headers=headers, json={ "jsonrpc": "2.0", "id": 3, "method": "resources/read", "params": {"uri": uri}, }, ) widgets = client.post( "/mcp/gitea-all/prod", headers=headers, json={ "jsonrpc": "2.0", "id": 4, "method": "resources/list", "params": {"_meta": {"openai/locale": "nl-NL"}}, }, ) widget_uris = {item["uri"] for item in widgets.json()["result"]["resources"]} widget = client.post( "/mcp/gitea-all/prod", headers=headers, json={ "jsonrpc": "2.0", "id": 5, "method": "resources/read", "params": { "uri": "ui://itworx-gitea/operations-v4.html", "_meta": {"progressToken": 5}, }, }, ) invalid = client.post( "/mcp/gitea-all/prod", headers=headers, json={ "jsonrpc": "2.0", "id": 6, "method": "skills/get", "params": {"uri": "skill://itworx-gitea/missing/SKILL.md"}, }, ) self.assertIn( "io.modelcontextprotocol/skills", initialized.json()["result"]["capabilities"]["extensions"], ) self.assertEqual( initialized.json()["result"]["capabilities"]["resources"], {"listChanged": False}, ) self.assertEqual(len(listed.json()["result"]["skills"]), 5) self.assertEqual(resource.json()["result"]["contents"][0]["uri"], uri) self.assertEqual( widget_uris, { "ui://itworx-gitea/change-review-v2.html", "ui://itworx-gitea/change-review-v3.html", "ui://itworx-gitea/change-review-v4.html", "ui://itworx-gitea/operations-v2.html", "ui://itworx-gitea/operations-v3.html", "ui://itworx-gitea/operations-v4.html", }, ) legacy_widget = client.post( "/mcp/gitea-all/prod", headers=headers, json={ "jsonrpc": "2.0", "id": 7, "method": "resources/read", "params": { "uri": "ui://itworx-gitea/operations-v2.html", "_meta": {"openai/locale": "nl-NL"}, }, }, ) widget_content = widget.json()["result"]["contents"][0] self.assertEqual(widget_content["mimeType"], "text/html;profile=mcp-app") self.assertIn("", widget_content["text"]) self.assertIn("hydrationAttempts<50", widget_content["text"]) self.assertEqual(widget_content["_meta"]["ui"]["csp"]["connectDomains"], []) self.assertEqual( legacy_widget.json()["result"]["contents"][0]["mimeType"], "text/html;profile=mcp-app", ) self.assertEqual(invalid.json()["error"]["code"], -32602) self.assertEqual(len(self.connector.requests), 1) asyncio.run(proxy.close()) def test_tool_catalog_metadata_drives_r2_policy_and_unknown_tools_fail_closed(self) -> None: audit = InMemoryAuditLog() governance = Governance( PolicyEngine( ( PolicyRule( slug="deny-unknown", priority=100, description="Unknown tool calls are privileged and denied.", match={"riskClasses": ["RX"]}, decision=Decision.DENY, terminal=True, ), PolicyRule( slug="allow-gitea-write", project="gitea-all", priority=50, description="Allow the bounded Gitea write set.", match={ "project": "gitea-all", "riskClasses": ["R2"], "toolPatterns": ["gitea.*.create", "gitea.file.update"], }, decision=Decision.ALLOW, ), ) ), audit, ) registry = AliasRegistry(StaticAliasSource((record(self.client_id, project="gitea-all"),))) proxy = DataplaneProxy( registry, self.authenticator, governance=governance, tool_metadata={"gitea.branch.create": ("gitea", "R2")}, transport=httpx.MockTransport(self.connector.handle), identity_assertion=lambda project, environment, client_id, trace_id: "signed-fixture", ) app = FastAPI() app.include_router(build_router(proxy)) client = TestClient(app) allowed = client.post( "/mcp/gitea-all/prod", headers={"Authorization": "Bearer token"}, json={ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "gitea.branch.create", "arguments": {}}, }, ) denied = client.post( "/mcp/gitea-all/prod", headers={"Authorization": "Bearer token"}, json={ "jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": "gitea.repository.admin", "arguments": {}}, }, ) self.assertEqual(allowed.status_code, 200) self.assertEqual(denied.status_code, 403) self.assertEqual(denied.json()["error"]["code"], "POLICY_DENIED") self.assertEqual(len(self.connector.requests), 1) self.assertEqual(audit.events[0].tool, "gitea.branch.create") self.assertEqual(audit.events[0].policy_decision, "allow") self.assertEqual(audit.events[1].tool, "gitea.repository.admin") self.assertEqual(audit.events[1].outcome, "denied") asyncio.run(proxy.close()) def test_contextforge_published_tool_name_uses_canonical_policy_identity(self) -> None: audit = InMemoryAuditLog() governance = Governance( PolicyEngine( ( PolicyRule( slug="deny-unknown", priority=100, description="Unknown tool calls fail closed.", match={"riskClasses": ["RX"]}, decision=Decision.DENY, terminal=True, ), PolicyRule( slug="allow-read", priority=50, description="Known reads are allowed.", match={ "riskClasses": ["R1"], "toolPatterns": ["gitea.repository.list"], }, decision=Decision.ALLOW, ), ) ), audit, ) registry = AliasRegistry(StaticAliasSource((record(self.client_id, project="gitea-all"),))) proxy = DataplaneProxy( registry, self.authenticator, governance=governance, tool_metadata={ "itworx-gitea-gitea-repository-list": ( "gitea", "R1", "gitea.repository.list", ) }, transport=httpx.MockTransport(self.connector.handle), identity_assertion=lambda project, environment, client_id, trace_id: "signed-fixture", ) app = FastAPI() app.include_router(build_router(proxy)) response = TestClient(app).post( "/mcp/gitea-all/prod", headers={"Authorization": "Bearer token"}, json={ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { "name": "itworx-gitea-gitea-repository-list", "arguments": {}, }, }, ) self.assertEqual(response.status_code, 200) self.assertEqual(audit.events[0].tool, "gitea.repository.list") self.assertEqual(audit.events[0].policy_decision, "allow") asyncio.run(proxy.close()) def test_registered_short_alias_resolves_to_authoritative_project(self) -> None: short = record(self.client_id, project="general-infrastructure").model_copy( update={"alias": "/mcp/general/prod"} ) registry = AliasRegistry(StaticAliasSource((short,))) proxy = DataplaneProxy( registry, self.authenticator, transport=httpx.MockTransport(self.connector.handle), upstream_authorization=lambda: "Basic internal-only", identity_assertion=lambda project, environment, client_id, trace_id: "signed-fixture", ) app = FastAPI() app.include_router(build_router(proxy)) response = TestClient(app).post( "/mcp/general/prod", headers={ "Authorization": "Bearer client-secret", "X-ITWorx-Project": "attacker-project", }, content=b"{}", ) self.assertEqual(response.status_code, 200) forwarded = self.connector.requests[-1].headers self.assertEqual(forwarded["x-itworx-project"], "general-infrastructure") self.assertEqual(forwarded["authorization"], "Basic internal-only") asyncio.run(proxy.close()) def test_unknown_cross_project_and_dev_to_prod_aliases_are_denied(self) -> None: unknown = self.http.post("/mcp/unknown/prod", headers={"Authorization": "Bearer t"}) cross_project = self.http.post("/mcp/ludarium/dev", headers={"Authorization": "Bearer t"}) self.assertEqual(unknown.status_code, 404) self.assertEqual(cross_project.status_code, 404) self.assertEqual(len(self.connector.requests), 0) def test_binding_denies_wrong_client_and_attacker_cannot_select_upstream(self) -> None: wrong = StaticAuthenticator(principal(uuid4())) proxy = DataplaneProxy( self.registry, wrong, transport=httpx.MockTransport(self.connector.handle), ) app = FastAPI() app.include_router(build_router(proxy)) response = TestClient(app).post( "/mcp/ludarium/prod?upstream_url=https://attacker.invalid", headers={"Authorization": "Bearer token"}, content=b"{}", ) self.assertEqual(response.status_code, 403) self.assertEqual(len(self.connector.requests), 0) asyncio.run(proxy.close()) def test_oversized_body_timeout_and_invalid_protocol_are_predictable(self) -> None: oversized = self.http.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer token", "Content-Length": "1025"}, content=b"x" * 1025, ) self.assertEqual(oversized.status_code, 400) self.assertEqual(oversized.json()["error"]["code"], "VALIDATION_FAILED") invalid = self.http.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer token", "Mcp-Protocol-Version": "1999"}, content=b"{}", ) self.assertEqual(invalid.status_code, 400) self.connector.mode = "timeout" timeout = self.http.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer token"}, content=b"{}", ) self.assertEqual(timeout.status_code, 504) self.assertEqual(timeout.json()["error"]["code"], "TIMEOUT") def test_missing_authentication_is_rejected_before_connector_access(self) -> None: response = self.http.post("/mcp/ludarium/prod", content=b"{}") self.assertEqual(response.status_code, 401) self.assertEqual(response.json()["error"]["code"], "AUTHENTICATION_REQUIRED") self.assertEqual(len(self.connector.requests), 0) def test_auth_header_never_enters_audit_and_revocation_invalidation_drops_cache(self) -> None: first = self.http.get( "/mcp/ludarium/prod", headers={"Authorization": "Bearer secret-token"} ) self.assertEqual(first.status_code, 200) self.assertNotIn( "secret-token", json.dumps([item.model_dump(mode="json") for item in self.recorder.records]), ) self.registry.invalidate(client_id=self.client_id) self.assertEqual(self.registry.health()["cached_aliases"], 0) resolved = self.registry.resolve("ludarium", "prod") self.assertIsNotNone(resolved) assert resolved is not None self.assertEqual(resolved.alias, "/mcp/ludarium/prod") def test_stream_cancellation_closes_upstream_response(self) -> None: response = httpx.Response( 200, content=b"chunk", request=httpx.Request("GET", "https://gateway.invalid/registered"), ) async def cancel_after_first_chunk() -> None: stream = self.proxy._response_stream( response, trace_id="trace-cancel", record=record(self.client_id), principal=principal(self.client_id), protocol_revision="2026-07-28", method="GET", request_bytes=0, request_id=None, mcp_method=None, mcp_name=None, started=time.monotonic(), ) await stream.__anext__() await stream.aclose() asyncio.run(cancel_after_first_chunk()) self.assertTrue(response.is_closed) def test_sqlalchemy_source_projects_only_enabled_bound_hub_state(self) -> None: from services.control_api.persistence.dataplane_source import SqlAlchemyAliasSource from services.control_api.persistence.db import ( Base, create_engine_from_url, session_factory, ) from services.control_api.persistence.models import ( ClientBindingRow, ClientRow, EnvironmentRow, ProjectRow, VirtualServerRow, WorkspaceRow, ) engine = create_engine_from_url("sqlite:///:memory:") Base.metadata.create_all(engine) sessions = session_factory(engine) workspace_id = uuid4() project_id = uuid4() environment_id = uuid4() server_id = uuid4() client_id = uuid4() with sessions() as session: with session.begin(): workspace = WorkspaceRow(id=workspace_id, slug="itworx", name="ITWorx") session.add(workspace) session.flush() project = ProjectRow( id=project_id, workspace_id=workspace_id, slug="ludarium", name="Ludarium", ) session.add(project) session.flush() environment_row = EnvironmentRow( id=environment_id, workspace_id=workspace_id, project_id=project_id, name="prod", classification="production", ) session.add(environment_row) session.flush() client = ClientRow( id=client_id, workspace_id=workspace_id, subject="svc-ludarium", issuer="https://fake.oidc.test", status="active", ) session.add(client) session.flush() server = VirtualServerRow( id=server_id, workspace_id=workspace_id, project_id=project_id, environment_id=environment_id, slug="ludarium-prod-readonly", desired_hash="a" * 64, enabled=True, health_state="healthy", record_metadata={"endpointAlias": "/mcp/ludarium/prod"}, ) session.add(server) session.flush() session.add( ClientBindingRow( id=uuid4(), workspace_id=workspace_id, project_id=project_id, client_id=client_id, virtual_server_id=server_id, enabled=True, ) ) with sessions() as session: aliases = SqlAlchemyAliasSource( session, workspace_id=workspace_id, gateway_base_url="https://mcp.example.com", ).aliases() self.assertEqual(len(aliases), 1) self.assertEqual( str(aliases[0].upstream_url), "https://mcp.example.com/mcp/ludarium/prod", ) self.assertEqual(aliases[0].bindings[0].client_id, client_id) with sessions() as session: internal = SqlAlchemyAliasSource( session, workspace_id=workspace_id, gateway_base_url="http://contextforge-control:4444/mcp", ).aliases() self.assertEqual(internal, ()) with sessions() as session, session.begin(): stored = session.get(VirtualServerRow, server_id) assert stored is not None stored.gateway_external_id = str(server_id) with sessions() as session: internal = SqlAlchemyAliasSource( session, workspace_id=workspace_id, gateway_base_url="http://contextforge-control:4444/mcp", ).aliases() self.assertEqual(internal, ()) with sessions() as session, session.begin(): stored = session.get(VirtualServerRow, server_id) assert stored is not None stored.gateway_external_id = server_id.hex with sessions() as session: internal = SqlAlchemyAliasSource( session, workspace_id=workspace_id, gateway_base_url="http://contextforge-control:4444/mcp", ).aliases() self.assertEqual( str(internal[0].upstream_url), f"http://contextforge-control:4444/servers/{server_id.hex}/mcp", ) engine.dispose() def test_health_readiness_are_safe(self) -> None: health = self.http.get("/health") ready = self.http.get("/health/ready") self.assertEqual(health.status_code, 200) self.assertEqual(ready.status_code, 200) self.assertNotIn("gateway.invalid", ready.text) def test_late_stream_failure_retries_r1_and_preserves_next_request(self) -> None: class FailingStream(httpx.AsyncByteStream): async def __aiter__(self): yield b'{"jsonrpc":"2.0","id":1,' raise httpx.ReadError("late failure") class FlakyTransport: def __init__(self) -> None: self.calls = 0 async def handle(self, request: httpx.Request) -> httpx.Response: self.calls += 1 if self.calls == 1: return httpx.Response(200, stream=FailingStream(), request=request) return httpx.Response( 200, json={"jsonrpc": "2.0", "id": 1, "result": {"ok": True}}, request=request, ) transport = FlakyTransport() proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(transport.handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), tool_metadata={"gitea.repository.get": ("gitea", "R1")}, identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: first = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "gitea.repository.get", "arguments": {}}, }, ) second = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"}, ) self.assertEqual(first.status_code, 200) self.assertEqual(first.headers["x-itworx-upstream-retries"], "1") self.assertEqual(second.status_code, 200) self.assertEqual(transport.calls, 3) finally: asyncio.run(proxy.close()) def test_compressed_upstream_response_is_decoded_once_and_rebounded(self) -> None: document = {"jsonrpc": "2.0", "id": 1, "result": {"tools": []}} encoded = gzip.compress(json.dumps(document).encode()) async def handle(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={ "content-type": "application/json", "content-encoding": "gzip", "content-length": str(len(encoded)), "etag": '"compressed-representation"', "date": "Sat, 29 Aug 2026 15:48:15 GMT", "server": "upstream-test-server", }, content=encoded, request=request, ) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(response.status_code, 200) self.assertEqual(response.json(), document) self.assertEqual(response.headers["content-length"], str(len(response.content))) self.assertNotIn("content-encoding", response.headers) self.assertNotIn("etag", response.headers) self.assertNotEqual(response.headers.get("date"), "Sat, 29 Aug 2026 15:48:15 GMT") self.assertNotEqual(response.headers.get("server"), "upstream-test-server") finally: asyncio.run(proxy.close()) def test_client_compression_preferences_are_not_forwarded_upstream(self) -> None: document = {"jsonrpc": "2.0", "id": 1, "result": {"tools": ["x" * 12000]}} observed_accept_encoding: str | None = None async def handle(request: httpx.Request) -> httpx.Response: nonlocal observed_accept_encoding observed_accept_encoding = request.headers.get("accept-encoding") return httpx.Response(200, json=document, request=request) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=16384), identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={ "Authorization": "Bearer opaque-token", "Accept-Encoding": "gzip, deflate, br, zstd", }, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(response.status_code, 200) self.assertEqual(response.json(), document) self.assertEqual(observed_accept_encoding, "identity") finally: asyncio.run(proxy.close()) def test_compressed_upstream_response_is_bounded_after_decompression(self) -> None: encoded = gzip.compress(b"A" * 8192) async def handle(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-encoding": "gzip"}, content=encoded, request=request, ) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(response.status_code, 502) self.assertEqual(response.json()["error"]["code"], "RESULT_TOO_LARGE") finally: asyncio.run(proxy.close()) def test_late_stream_failure_never_retries_r3_mutation(self) -> None: class FailingStream(httpx.AsyncByteStream): async def __aiter__(self): yield b"partial" raise httpx.ReadError("late failure") calls = 0 async def handle(request: httpx.Request) -> httpx.Response: nonlocal calls calls += 1 return httpx.Response(200, stream=FailingStream(), request=request) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=1024, max_response_bytes=4096), tool_metadata={"gitea.pullrequest.sync": ("gitea", "R3")}, identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json={ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "gitea.pullrequest.sync", "arguments": {}}, }, ) self.assertEqual(response.status_code, 503) self.assertEqual(response.json()["error"]["operation"], "gitea.pullrequest.sync") self.assertFalse(response.json()["error"]["safeToRetry"]) self.assertEqual(calls, 1) finally: asyncio.run(proxy.close()) def test_json_rpc_batch_is_rejected_before_risk_classification(self) -> None: calls = 0 async def handle(request: httpx.Request) -> httpx.Response: nonlocal calls calls += 1 return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {}}) proxy = DataplaneProxy( self.registry, self.authenticator, transport=httpx.MockTransport(handle), settings=RouterSettings(max_body_bytes=4096, max_response_bytes=4096), tool_metadata={"gitea.pullrequest.sync": ("gitea", "R3")}, identity_assertion=lambda project, environment, client_id, trace_id: "signed", ) app = FastAPI() app.include_router(build_router(proxy)) try: with TestClient(app) as client: response = client.post( "/mcp/ludarium/prod", headers={"Authorization": "Bearer opaque-token"}, json=[ {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, { "jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": "gitea.pullrequest.sync", "arguments": {}}, }, ], ) self.assertEqual(response.status_code, 400) self.assertEqual(response.json()["error"]["code"], "VALIDATION_FAILED") self.assertFalse(response.json()["error"]["safeToRetry"]) self.assertEqual(calls, 0) finally: asyncio.run(proxy.close()) class DemoTokenEndToEndScopeTests(unittest.TestCase): """One REAL minted demo token, through the REAL DemoTokenAuthenticator, CompositeAuthenticator, and AliasAuthorizer: allowed into mobilityops, denied everywhere else.""" def setUp(self) -> None: from services.control_api.persistence.demo_client import ( DEMO_CLIENT_ID, DEMO_CLIENT_ISSUER, DEMO_CLIENT_SUBJECT, ) from services.runtime.dataplane import CompositeAuthenticator from services.runtime.demo_token import DemoTokenAuthenticator, mint_demo_token self.key = b"x" * 32 self.temp_dir = tempfile.TemporaryDirectory() # Deliberately absent: nothing is revoked in this test. self.revocations_path = Path(self.temp_dir.name) / "revocations.json" # Real wall-clock mint -- DemoTokenAuthenticator has no test-clock hook # and AliasAuthorizer independently checks expiry against real "now". self.token, self.jti = mint_demo_token(self.key, ttl_hours=1) self.oidc = StaticAuthenticator(None) self.authenticator = CompositeAuthenticator( demo=DemoTokenAuthenticator(self.key, self.revocations_path), oidc=self.oidc, ) demo_record = AliasRecord( project_slug="mobilityops", environment="prod", alias="/mcp/mobilityops/prod", upstream_url=TypeAdapter(AnyHttpUrl).validate_python( "https://gateway.invalid/mobilityops" ), enabled=True, bindings=( ClientBinding( client_id=DEMO_CLIENT_ID, issuer=DEMO_CLIENT_ISSUER, subject=DEMO_CLIENT_SUBJECT, enabled=True, ), ), ) # A second, real project the demo identity has no binding for. other_record = record(uuid4(), project="ludarium") self.connector = ConnectorTransport() self.registry = AliasRegistry(StaticAliasSource((demo_record, other_record))) self.proxy = DataplaneProxy( self.registry, self.authenticator, authorizer=AliasAuthorizer(), transport=httpx.MockTransport(self.connector.handle), identity_assertion=lambda project, environment, client_id, trace_id: ( f"signed:{project}:{environment}:{client_id}:{trace_id}" ), ) self.app = FastAPI() self.app.include_router(build_router(self.proxy)) self.http = TestClient(self.app) def tearDown(self) -> None: asyncio.run(self.proxy.close()) self.temp_dir.cleanup() def test_real_demo_token_is_allowed_into_mobilityops_and_denied_elsewhere(self) -> None: from services.control_api.persistence.demo_client import DEMO_CLIENT_ID allowed = self.http.post( "/mcp/mobilityops/prod", headers={ "Authorization": f"Bearer {self.token}", "Mcp-Protocol-Version": "2026-07-28", }, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(allowed.status_code, 200) self.assertEqual(allowed.json()["result"]["protocol"], "2026-07-28") self.assertEqual(len(self.connector.requests), 1) forwarded = self.connector.requests[0].headers self.assertEqual(forwarded["x-itworx-project"], "mobilityops") self.assertEqual(forwarded["x-itworx-environment"], "prod") self.assertEqual(forwarded["x-itworx-client-id"], str(DEMO_CLIENT_ID)) self.assertNotIn("authorization", forwarded) # The EXACT same token against a project it holds no binding for. denied = self.http.post( "/mcp/ludarium/prod", headers={"Authorization": f"Bearer {self.token}"}, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, ) self.assertEqual(denied.status_code, 403) self.assertEqual(denied.json()["error"]["code"], "ACCESS_DENIED") # Still exactly one upstream request: the deny never reached a connector. self.assertEqual(len(self.connector.requests), 1) # Every request carried a demo-prefixed bearer, so CompositeAuthenticator # must have routed both to the demo path and never consulted OIDC. self.assertEqual(self.oidc.seen_headers, []) def test_non_demo_bearer_still_routes_to_the_oidc_path(self) -> None: """Companion to the assertion above: proves seen_headers == [] means 'never routed here', not 'this authenticator is inert'.""" response = self.http.post( "/mcp/mobilityops/prod", headers={"Authorization": "Bearer an-ordinary-oidc-token"}, content=b"{}", ) self.assertEqual(response.status_code, 401) self.assertEqual(self.oidc.seen_headers, ["Bearer an-ordinary-oidc-token"]) self.assertEqual(len(self.connector.requests), 0) if __name__ == "__main__": unittest.main()