diff --git a/docs/phase-foundation-roadmap.md b/docs/phase-foundation-roadmap.md index 3cd46fd..70e8a44 100644 --- a/docs/phase-foundation-roadmap.md +++ b/docs/phase-foundation-roadmap.md @@ -111,7 +111,7 @@ quotas, worker isolation, distributed state, hybrid routing, and E2E proof. - [x] Add isolated shared worker runtime for allowlisted stateless upstreams. - [x] Add distributed deployment state with locking, conflict handling, rollback, recovery, and audit. -- [ ] Add hybrid routing between local edge tools and shared workers. +- [x] Add hybrid routing between local edge tools and shared workers. - [ ] Add shared-runtime E2E proof for tenant isolation, authz denial, quota denial, session affinity, audit, rollback, and degraded mode. diff --git a/docs/protocol.md b/docs/protocol.md index fd6245f..cadec08 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -228,3 +228,12 @@ upstream and original upstream tool name, passes the upstream call timeout from config, and maps upstream timeout, crash, unknown-tool, invalid-argument, disabled-prefix, profile-denial, and mutating-profile-denial failures to broker-owned error codes. + +`mcp_broker.hybrid_router.HybridToolRouter` is the Phase 3 route classifier for +ordinary upstream `tools/call` requests. It preserves the same client-visible +tool name and argument object. By default it calls the local edge upstream path. +Only upstreams tagged `stateless` and `shared-worker`, configured in shared +mode, free of local state/env/session metadata, and supplied with tenant, +team, and quota metadata can route to the in-process fake shared worker proof. +Quota denial on that shared path fails closed and does not retry through the +local edge broker. diff --git a/docs/shared-runtime-guardrails.md b/docs/shared-runtime-guardrails.md index 9d98e9a..19bb516 100644 --- a/docs/shared-runtime-guardrails.md +++ b/docs/shared-runtime-guardrails.md @@ -183,6 +183,31 @@ shared-runtime state under the runtime state directory, appends a rollback journal before recovery replay is needed, and fails closed on lock or revision conflicts. +## Hybrid Routing + +Contract statement: hybrid routing preserves the existing tools/call client shape. +Client-visible tool names and argument objects remain unchanged when a route is +classified for local edge or shared-worker execution. + +Contract statement: local-only tools continue to route to the edge broker. +Local-only includes stateful, browser, file-access, local-secret, OAuth, +unknown, mutating, configured-state, env-backed, request-meta-backed, and +session-env-backed upstreams. + +Contract statement: allowlisted stateless tools can route to shared workers only after quota approval. +The current allowlist is config-derived from upstream tags and safety metadata: +the upstream must be tagged stateless and shared-worker, use shared mode, avoid +local state, and pass the supplied quota snapshot. The default route remains +local edge when shared-runtime metadata is absent. + +Contract statement: shared-worker quota denial does not fall back to local edge execution. +Quota denial is returned as an audited denial for the shared-worker attempt so a +shared route cannot bypass quota by retrying locally. + +The code contract lives in `mcp_broker.hybrid_router`. The daemon delegates +ordinary upstream calls through the hybrid router while preserving broker +catalog tools and the local edge fallback. + ## Mandatory Non-Goals Phase 3 does not add hosted execution. It does not add remote tool calls, remote diff --git a/src/mcp_broker/daemon.py b/src/mcp_broker/daemon.py index 975ffe1..72f494a 100644 --- a/src/mcp_broker/daemon.py +++ b/src/mcp_broker/daemon.py @@ -13,7 +13,7 @@ from mcp_broker import __version__ from mcp_broker.broker import BrokerCore, BrokerToolError -from mcp_broker.catalog import BrokerCatalogFacade, profile_allows_upstream +from mcp_broker.catalog import profile_allows_upstream from mcp_broker.config import BrokerConfig, UpstreamConfig from mcp_broker.daemon_helpers import ( configured_upstream_health as _configured_upstream_health, @@ -30,10 +30,12 @@ from mcp_broker.daemon_provenance import source_provenance as _source_provenance_impl from mcp_broker.daemon_request_context import BrokerDaemonRequestContextMixin from mcp_broker.daemon_status import BrokerDaemonStatusMixin +from mcp_broker.daemon_tool_calls import BrokerDaemonToolCallMixin from mcp_broker.daemon_upstreams import BrokerDaemonUpstreamMixin from mcp_broker.jsonrpc import JsonRpcRequest, JsonRpcResponse from mcp_broker.protocol import McpProtocolHandler from mcp_broker.runtime_reaper import RuntimePaths, write_socket_metadata +from mcp_broker.shared_worker import SharedWorkerRuntime from mcp_broker.upstream_http import HttpUpstreamClient, HttpUpstreamError from mcp_broker.upstream_protocols import ( HttpUpstreamClientProtocol, @@ -75,6 +77,7 @@ class BrokerDaemon( BrokerDaemonStatusMixin, BrokerDaemonUpstreamMixin, BrokerDaemonRequestContextMixin, + BrokerDaemonToolCallMixin, ): runtime_root: Path socket_path: Path @@ -108,6 +111,7 @@ def __post_init__(self) -> None: self._last_request_method: str | None = None self._last_request_status: str | None = None self._auth_repair_stats: dict[str, dict[str, int | str]] = {} + self._shared_worker_runtime = SharedWorkerRuntime(tools=[]) @property def lock_path(self) -> Path: @@ -328,53 +332,6 @@ def _handle_tools_list(self, request: JsonRpcRequest) -> JsonRpcResponse: return JsonRpcResponse.error(request.id, -32000, str(exc)) return JsonRpcResponse.result(request.id, result) - def _handle_tools_call(self, request: JsonRpcRequest) -> JsonRpcResponse: - if self.broker_config is None: - return JsonRpcResponse.error(request.id, -32000, "broker config is not loaded") - params = request.params - if not isinstance(params, dict): - return JsonRpcResponse.error(request.id, -32602, "tools/call params must be an object") - name = params.get("name") - arguments = params.get("arguments", {}) - if not isinstance(name, str) or not isinstance(arguments, dict): - return JsonRpcResponse.error(request.id, -32602, "tools/call name and arguments required") - try: - session_id = self._session_id_from_params(params) - session_context = self._session_context_from_params(params) - profile = self._effective_profile(params, session_context) - except ValueError as exc: - return JsonRpcResponse.error(request.id, -32602, str(exc)) - call_upstream = self._call_upstream_for_session(session_id, session_context) - list_upstream = self._list_upstream_for_session(session_id, session_context) - canonical_name = profile.canonical_broker_tool_name(name) if profile is not None else name - if canonical_name.startswith("broker."): - try: - result = BrokerCatalogFacade( - broker_config=self.broker_config, - profile=profile, - list_upstream=list_upstream, - call_upstream=call_upstream, - call_locks=self._upstream_call_locks, - status_provider=self._upstream_health_for_status, - client_cwd=session_context.get("client_cwd"), - ).call_tool(name, arguments) - except (BrokerToolError, ValueError) as exc: - return JsonRpcResponse.error(request.id, -32000, str(exc)) - return JsonRpcResponse.result(request.id, result) - arguments = self._inject_cwd_project_arg(name, arguments, session_context) - core = BrokerCore( - settings=self.broker_config.broker, - upstreams=self.broker_config.upstreams, - profile=profile, - call_locks=self._upstream_call_locks, - ) - try: - result = core.call_tool(name, arguments, call_upstream) - except (BrokerToolError, ValueError) as exc: - message = exc.message if isinstance(exc, BrokerToolError) else str(exc) - return JsonRpcResponse.error(request.id, -32000, message) - return JsonRpcResponse.result(request.id, result) - def _create_stdio_upstream_process( self, upstream: UpstreamConfig, diff --git a/src/mcp_broker/daemon_tool_calls.py b/src/mcp_broker/daemon_tool_calls.py new file mode 100644 index 0000000..8b1a283 --- /dev/null +++ b/src/mcp_broker/daemon_tool_calls.py @@ -0,0 +1,109 @@ +"""Daemon JSON-RPC tools/call handling.""" + +from __future__ import annotations + +from typing import Any + +from mcp_broker.broker import BrokerToolError +from mcp_broker.catalog import BrokerCatalogFacade +from mcp_broker.hybrid_router import HybridRoutingContext, HybridToolRouter +from mcp_broker.jsonrpc import JsonRpcRequest, JsonRpcResponse + + +class BrokerDaemonToolCallMixin: + def _handle_tools_call(self, request: JsonRpcRequest) -> JsonRpcResponse: + if self.broker_config is None: + return JsonRpcResponse.error(request.id, -32000, "broker config is not loaded") + params = request.params + if not isinstance(params, dict): + return JsonRpcResponse.error(request.id, -32602, "tools/call params must be an object") + name = params.get("name") + arguments = params.get("arguments", {}) + if not isinstance(name, str) or not isinstance(arguments, dict): + return JsonRpcResponse.error(request.id, -32602, "tools/call name and arguments required") + try: + session_id = self._session_id_from_params(params) + session_context = self._session_context_from_params(params) + profile = self._effective_profile(params, session_context) + except ValueError as exc: + return JsonRpcResponse.error(request.id, -32602, str(exc)) + call_upstream = self._call_upstream_for_session(session_id, session_context) + list_upstream = self._list_upstream_for_session(session_id, session_context) + canonical_name = profile.canonical_broker_tool_name(name) if profile is not None else name + if canonical_name.startswith("broker."): + return self._handle_broker_catalog_tool_call( + request_id=request.id, + name=name, + arguments=arguments, + profile=profile, + list_upstream=list_upstream, + call_upstream=call_upstream, + session_context=session_context, + ) + return self._handle_upstream_tool_call( + request_id=request.id, + name=name, + arguments=arguments, + params=params, + profile=profile, + call_upstream=call_upstream, + session_context=session_context, + ) + + def _handle_broker_catalog_tool_call( + self, + *, + request_id: str | int | None, + name: str, + arguments: dict[str, Any], + profile: Any, + list_upstream: Any, + call_upstream: Any, + session_context: dict[str, str], + ) -> JsonRpcResponse: + try: + result = BrokerCatalogFacade( + broker_config=self.broker_config, + profile=profile, + list_upstream=list_upstream, + call_upstream=call_upstream, + call_locks=self._upstream_call_locks, + status_provider=self._upstream_health_for_status, + client_cwd=session_context.get("client_cwd"), + ).call_tool(name, arguments) + except (BrokerToolError, ValueError) as exc: + return JsonRpcResponse.error(request_id, -32000, str(exc)) + return JsonRpcResponse.result(request_id, result) + + def _handle_upstream_tool_call( + self, + *, + request_id: str | int | None, + name: str, + arguments: dict[str, Any], + params: dict[str, Any], + profile: Any, + call_upstream: Any, + session_context: dict[str, str], + ) -> JsonRpcResponse: + arguments = self._inject_cwd_project_arg(name, arguments, session_context) + shared_context = HybridRoutingContext.from_params(params) + try: + result = HybridToolRouter( + upstreams=self.broker_config.upstreams, + settings=self.broker_config.broker, + profile=profile, + call_locks=self._upstream_call_locks, + shared_worker=self._shared_worker_runtime, + ).call_tool( + advertised_name=name, + arguments=arguments, + edge_caller=call_upstream, + tenant_context=shared_context.tenant_context, + team_id=shared_context.team_id, + quota_snapshot=shared_context.quota_snapshot, + ) + except (BrokerToolError, ValueError) as exc: + message = exc.message if isinstance(exc, BrokerToolError) else str(exc) + return JsonRpcResponse.error(request_id, -32000, message) + return JsonRpcResponse.result(request_id, result) diff --git a/src/mcp_broker/hybrid_router.py b/src/mcp_broker/hybrid_router.py new file mode 100644 index 0000000..f2f5c76 --- /dev/null +++ b/src/mcp_broker/hybrid_router.py @@ -0,0 +1,215 @@ +"""Hybrid local-edge and shared-worker routing contracts.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Mapping, Protocol + +from mcp_broker.broker import BrokerCore, BrokerToolError, ToolCaller +from mcp_broker.config import BrokerSettings, UpstreamConfig +from mcp_broker.profiles import ToolExposureProfile +from mcp_broker.quota_policy import QuotaPolicyError, decide_quota +from mcp_broker.shared_runtime_policy import LOCAL_EDGE_BOUNDARY, SHARED_WORKER_BOUNDARY + + +HYBRID_ROUTER_SCHEMA_VERSION = 1 +STATELESS_TAG = "stateless" +SHARED_WORKER_TAGS = frozenset({"shared-worker", "shared_worker"}) + + +class SharedWorkerCaller(Protocol): + def call_tool( + self, + *, + tenant_context: Mapping[str, Any], + upstream_id: str, + upstream_class: str, + allowlisted: bool, + requires_local_state: bool, + tool_name: str, + arguments: Mapping[str, Any], + quota_decision: Mapping[str, Any], + ) -> dict[str, Any]: + ... + + +@dataclass(frozen=True) +class SharedWorkerRoute: + call_tool: Callable[..., dict[str, Any]] + + +@dataclass(frozen=True) +class HybridRoutingContext: + tenant_context: Mapping[str, Any] | None + team_id: str | None + quota_snapshot: Mapping[str, Any] | None + + @classmethod + def from_params(cls, params: Mapping[str, Any]) -> "HybridRoutingContext": + return cls( + tenant_context=_optional_mapping(params.get("tenant_context")), + team_id=_optional_string(params.get("team_id")), + quota_snapshot=_optional_mapping(params.get("quota_snapshot")), + ) + + +class HybridToolRouter: + def __init__( + self, + *, + upstreams: Mapping[str, UpstreamConfig], + shared_worker: SharedWorkerCaller | SharedWorkerRoute, + settings: BrokerSettings | None = None, + profile: ToolExposureProfile | None = None, + call_locks: dict[str, Any] | None = None, + ) -> None: + self._settings = settings or BrokerSettings() + self._profile = profile + self._call_locks = {} if call_locks is None else call_locks + self._upstreams = dict(upstreams) + self._shared_worker = shared_worker + + def call_tool( + self, + *, + advertised_name: str, + arguments: dict[str, Any], + edge_caller: ToolCaller, + tenant_context: Mapping[str, Any] | None, + team_id: str | None, + quota_snapshot: Mapping[str, Any] | None, + ) -> dict[str, Any]: + core = BrokerCore( + settings=self._settings, + upstreams=self._upstreams, + profile=self._profile, + call_locks=self._call_locks, + ) + upstream, upstream_tool_name = core._resolve_call_route(advertised_name) + if not self._should_try_shared_worker( + upstream=upstream, + tenant_context=tenant_context, + team_id=team_id, + quota_snapshot=quota_snapshot, + ): + return core.call_tool(advertised_name, arguments, edge_caller) + quota_decision = _quota_decision( + tenant_context=tenant_context, + team_id=team_id, + upstream_id=upstream.name, + tool_name=advertised_name, + quota_snapshot=quota_snapshot, + ) + worker_response = self._shared_worker.call_tool( + tenant_context=tenant_context, + upstream_id=upstream.name, + upstream_class=STATELESS_TAG, + allowlisted=True, + requires_local_state=False, + tool_name=advertised_name, + arguments=arguments, + quota_decision=quota_decision, + ) + if worker_response.get("allowed") is not True: + return worker_response + result = worker_response.get("result") + if not isinstance(result, dict): + raise BrokerToolError( + code="invalid_shared_worker_response", + message=f"invalid shared worker tools/call response from {upstream.name}", + upstream_name=upstream.name, + tool_name=upstream_tool_name, + ) + BrokerCore._validate_tool_response(upstream.name, result) + return result + + def _should_try_shared_worker( + self, + *, + upstream: UpstreamConfig, + tenant_context: Mapping[str, Any] | None, + team_id: str | None, + quota_snapshot: Mapping[str, Any] | None, + ) -> bool: + if tenant_context is None or team_id is None or quota_snapshot is None: + return False + tags = set(upstream.tags) + return ( + STATELESS_TAG in tags + and bool(tags & SHARED_WORKER_TAGS) + and upstream.mode == "shared" + and not upstream.mutating + and not _requires_local_state(upstream) + ) + + +def build_hybrid_routing_policy() -> dict[str, Any]: + return { + "schema_version": HYBRID_ROUTER_SCHEMA_VERSION, + "default_execution_boundary": LOCAL_EDGE_BOUNDARY, + "client_contract": "unchanged_tools_call_shape", + "shared_worker_boundary": SHARED_WORKER_BOUNDARY, + "shared_worker_requires": [ + "stateless_tag", + "shared_mode", + "allowlisted", + "no_local_state", + "quota_allowed", + ], + "denial_audit_required": True, + } + + +def _requires_local_state(upstream: UpstreamConfig) -> bool: + return bool( + upstream.state_dir + or upstream.env + or upstream.env_files + or upstream.session_env + or upstream.request_meta + ) + + +def _optional_mapping(value: object) -> Mapping[str, Any] | None: + if isinstance(value, Mapping): + return value + return None + + +def _optional_string(value: object) -> str | None: + if isinstance(value, str): + return value + return None + + +def _quota_decision( + *, + tenant_context: Mapping[str, Any] | None, + team_id: str | None, + upstream_id: str, + tool_name: str, + quota_snapshot: Mapping[str, Any] | None, +) -> dict[str, Any]: + if tenant_context is None or team_id is None or quota_snapshot is None: + raise BrokerToolError( + code="shared_routing_context_missing", + message="shared routing requires tenant_context, team_id, and quota_snapshot", + upstream_name=upstream_id, + tool_name=tool_name, + ) + try: + return decide_quota( + tenant_context=tenant_context, + team_id=team_id, + upstream_id=upstream_id, + tool_name=tool_name, + quota_snapshot=quota_snapshot, + ) + except QuotaPolicyError as exc: + raise BrokerToolError( + code="shared_routing_quota_invalid", + message=str(exc), + upstream_name=upstream_id, + tool_name=tool_name, + ) from exc diff --git a/tests/e2e/test_hybrid_routing_e2e.py b/tests/e2e/test_hybrid_routing_e2e.py new file mode 100644 index 0000000..3fcdb80 --- /dev/null +++ b/tests/e2e/test_hybrid_routing_e2e.py @@ -0,0 +1,245 @@ +import pytest + + +pytestmark = pytest.mark.e2e + + +TENANT_A = { + "tenant_id": "tenant-a", + "workspace_id": "workspace-a", + "user_id": "user-a", +} + +TENANT_B = { + "tenant_id": "tenant-b", + "workspace_id": "workspace-b", + "user_id": "user-b", +} + + +def test_hybrid_router_preserves_client_tool_names_across_edge_and_worker() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter + from mcp_broker.shared_worker import SharedWorkerRuntime, SharedWorkerTool + + edge_calls: list[tuple[str, str, dict[str, object], int]] = [] + + def edge_call( + upstream_name: str, + upstream_tool_name: str, + arguments: dict[str, object], + timeout_seconds: int, + ) -> dict[str, object]: + edge_calls.append((upstream_name, upstream_tool_name, arguments, timeout_seconds)) + return { + "content": [{"type": "text", "text": "edge"}], + "structuredContent": { + "execution_boundary": "local_edge", + "upstream_tool_name": upstream_tool_name, + }, + } + + worker = SharedWorkerRuntime( + tools=[ + SharedWorkerTool( + upstream_id="example-stateless", + name="example.echo", + behavior="echo", + ) + ] + ) + router = HybridToolRouter( + upstreams={ + "local-store": UpstreamConfig( + name="local-store", + command="local-store", + tool_prefix="local-store", + mode="per_session", + tags=("stateful",), + state_dir="local-store", + ), + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless", "shared-worker"), + ), + }, + shared_worker=worker, + ) + + edge = router.call_tool( + advertised_name="local-store.search", + arguments={"query": "refund"}, + edge_caller=edge_call, + tenant_context=TENANT_A, + team_id="team-a", + quota_snapshot=_quota_snapshot( + tenant=TENANT_A, + team_id="team-a", + upstream_id="local-store", + tool_name="local-store.search", + ), + ) + shared = router.call_tool( + advertised_name="example.echo", + arguments={"message": "hello"}, + edge_caller=edge_call, + tenant_context=TENANT_A, + team_id="team-a", + quota_snapshot=_quota_snapshot( + tenant=TENANT_A, + team_id="team-a", + upstream_id="example-stateless", + tool_name="example.echo", + ), + ) + + assert edge["structuredContent"] == { + "execution_boundary": "local_edge", + "upstream_tool_name": "search", + } + assert shared == { + "content": [{"type": "text", "text": "hello"}], + "structuredContent": {"message": "hello"}, + } + assert edge_calls == [("local-store", "search", {"query": "refund"}, 60)] + + +def test_hybrid_worker_routing_keeps_tenant_scopes_isolated() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter + from mcp_broker.shared_worker import SharedWorkerRuntime, SharedWorkerTool + + worker = SharedWorkerRuntime( + tools=[ + SharedWorkerTool( + upstream_id="example-stateless", + name="example.echo", + behavior="echo", + ) + ] + ) + router = HybridToolRouter( + upstreams={ + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless", "shared-worker"), + ) + }, + shared_worker=worker, + ) + + first = router.call_tool( + advertised_name="example.echo", + arguments={"message": "tenant-a"}, + edge_caller=_unexpected_edge_call, + tenant_context=TENANT_A, + team_id="team-a", + quota_snapshot=_quota_snapshot( + tenant=TENANT_A, + team_id="team-a", + upstream_id="example-stateless", + tool_name="example.echo", + ), + ) + second = router.call_tool( + advertised_name="example.echo", + arguments={"message": "tenant-b"}, + edge_caller=_unexpected_edge_call, + tenant_context=TENANT_B, + team_id="team-b", + quota_snapshot=_quota_snapshot( + tenant=TENANT_B, + team_id="team-b", + upstream_id="example-stateless", + tool_name="example.echo", + ), + ) + + assert first["structuredContent"] == {"message": "tenant-a"} + assert second["structuredContent"] == {"message": "tenant-b"} + assert worker.state_snapshot() == {} + + +def test_hybrid_shared_worker_quota_denial_does_not_fallback_to_edge() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter + from mcp_broker.shared_worker import SharedWorkerRuntime, SharedWorkerTool + + worker = SharedWorkerRuntime( + tools=[ + SharedWorkerTool( + upstream_id="example-stateless", + name="example.echo", + behavior="echo", + ) + ] + ) + router = HybridToolRouter( + upstreams={ + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless", "shared-worker"), + ) + }, + shared_worker=worker, + ) + + decision = router.call_tool( + advertised_name="example.echo", + arguments={"message": "blocked"}, + edge_caller=_unexpected_edge_call, + tenant_context=TENANT_A, + team_id="team-a", + quota_snapshot=_quota_snapshot( + tenant=TENANT_A, + team_id="team-a", + upstream_id="example-stateless", + tool_name="example.echo", + tool_limit=0, + tool_used=0, + ), + ) + + assert decision["allowed"] is False + assert decision["reason"] == "quota_denied" + assert decision["audit_event"]["denial_reason"] == "tool_quota_exceeded" + + +def _unexpected_edge_call(*_args: object, **_kwargs: object) -> dict[str, object]: + raise AssertionError("shared-worker route must not call the edge broker") + + +def _quota_snapshot( + *, + tenant: dict[str, str], + team_id: str, + upstream_id: str, + tool_name: str, + tool_limit: int = 10, + tool_used: int = 0, +) -> dict[str, object]: + return { + "kill_switches": { + "global": False, + "teams": [], + "users": [], + "upstreams": [], + "tools": [], + }, + "limits": { + "global": {"limit": 100, "used": 0}, + "teams": {team_id: {"limit": 100, "used": 0}}, + "users": {tenant["user_id"]: {"limit": 100, "used": 0}}, + "upstreams": {upstream_id: {"limit": 100, "used": 0}}, + "tools": {tool_name: {"limit": tool_limit, "used": tool_used}}, + }, + } diff --git a/tests/journey/test_shared_runtime_guardrails_contract.py b/tests/journey/test_shared_runtime_guardrails_contract.py index a81e148..241bca1 100644 --- a/tests/journey/test_shared_runtime_guardrails_contract.py +++ b/tests/journey/test_shared_runtime_guardrails_contract.py @@ -26,6 +26,7 @@ def test_shared_runtime_guardrails_document_phase_3_gates() -> None: "## Quota And Cost Controls", "## Shared Worker Runtime", "## Distributed Deployment State", + "## Hybrid Routing", "## Mandatory Non-Goals", ] required_terms = [ @@ -64,6 +65,10 @@ def test_shared_runtime_guardrails_document_phase_3_gates() -> None: "lock acquisition, stale lock recovery, active lock conflict rejection, rollback journal, recovery replay, and audit events", "cloud storage is not implemented", "state mutations require a matching lock token", + "hybrid routing preserves the existing tools/call client shape", + "local-only tools continue to route to the edge broker", + "allowlisted stateless tools can route to shared workers only after quota approval", + "shared-worker quota denial does not fall back to local edge execution", "local edge broker remains the default", "no remote listener", "no shared upstream execution", diff --git a/tests/unit/test_daemon_jsonrpc_contract.py b/tests/unit/test_daemon_jsonrpc_contract.py index 8f6203b..31061b2 100644 --- a/tests/unit/test_daemon_jsonrpc_contract.py +++ b/tests/unit/test_daemon_jsonrpc_contract.py @@ -3018,6 +3018,106 @@ def test_daemon_tools_call_rejects_disabled_upstream_without_starting_stdio( assert daemon._stdio_upstreams == {} +def test_daemon_tools_call_delegates_normal_calls_through_hybrid_router( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import mcp_broker.daemon as daemon_module + import mcp_broker.daemon_tool_calls as daemon_tool_calls_module + from mcp_broker.config import BrokerConfig, BrokerSettings, RuntimeConfig, UpstreamConfig + from mcp_broker.jsonrpc import JsonRpcRequest + + router_calls: list[dict[str, object]] = [] + + class RecordingHybridRouter: + def __init__( # type: ignore[no-untyped-def] + self, + *, + upstreams, + shared_worker, + settings, + profile, + call_locks, + ): + self.upstreams = upstreams + self.shared_worker = shared_worker + self.settings = settings + self.profile = profile + self.call_locks = call_locks + + def call_tool( + self, + *, + advertised_name, + arguments, + edge_caller, + tenant_context, + team_id, + quota_snapshot, + ): # type: ignore[no-untyped-def] + router_calls.append( + { + "advertised_name": advertised_name, + "arguments": arguments, + "tenant_context": tenant_context, + "team_id": team_id, + "quota_snapshot": quota_snapshot, + "upstreams": sorted(self.upstreams), + "shared_worker_type": type(self.shared_worker).__name__, + } + ) + return edge_caller("read-store", "search", arguments, 60) + + monkeypatch.setattr(daemon_tool_calls_module, "HybridToolRouter", RecordingHybridRouter) + config = BrokerConfig( + runtime=RuntimeConfig( + root=tmp_path / "runtime", + socket_path=tmp_path / "broker.sock", + log_dir=tmp_path / "runtime" / "logs", + state_dir=tmp_path / "runtime" / "state", + secrets_dir=tmp_path / "runtime" / "secrets", + ), + broker=BrokerSettings(), + upstreams={ + "read-store": UpstreamConfig( + name="read-store", + command="read-store", + tool_prefix="read-store", + ) + }, + ) + daemon = daemon_module.BrokerDaemon( + runtime_root=config.runtime.root, + socket_path=config.runtime.socket_path, + broker_config=config, + ) + daemon._stdio_upstreams["read-store"] = CatalogClient( + tools=[{"name": "search"}], + response={"content": [{"type": "text", "text": "found"}]}, + ) + request = JsonRpcRequest( + method="tools/call", + id="call", + params={"name": "read-store.search", "arguments": {"query": "refund"}}, + has_id=True, + ) + + response = daemon._handle_tools_call(request) + + assert response.result == {"content": [{"type": "text", "text": "found"}]} + assert router_calls == [ + { + "advertised_name": "read-store.search", + "arguments": {"query": "refund"}, + "tenant_context": None, + "team_id": None, + "quota_snapshot": None, + "upstreams": ["read-store"], + "shared_worker_type": "SharedWorkerRuntime", + } + ] + + def test_daemon_maps_stdio_timeout_and_error(tmp_path: Path) -> None: from mcp_broker.broker import BrokerToolError from mcp_broker.config import BrokerConfig, BrokerSettings, RuntimeConfig, UpstreamConfig diff --git a/tests/unit/test_hybrid_router_contract.py b/tests/unit/test_hybrid_router_contract.py new file mode 100644 index 0000000..1dc6f2c --- /dev/null +++ b/tests/unit/test_hybrid_router_contract.py @@ -0,0 +1,343 @@ +import pytest + + +pytestmark = pytest.mark.unit + + +TENANT_CONTEXT = { + "tenant_id": "tenant-a", + "workspace_id": "workspace-a", + "user_id": "user-a", +} + +ALLOWED_QUOTA = {"allowed": True, "reason": "quota_allowed"} + + +def test_hybrid_router_policy_declares_default_local_boundary() -> None: + from mcp_broker.hybrid_router import build_hybrid_routing_policy + + policy = build_hybrid_routing_policy() + + assert policy["schema_version"] == 1 + assert policy["default_execution_boundary"] == "local_edge" + assert policy["client_contract"] == "unchanged_tools_call_shape" + assert policy["shared_worker_requires"] == [ + "stateless_tag", + "shared_mode", + "allowlisted", + "no_local_state", + "quota_allowed", + ] + assert policy["denial_audit_required"] is True + + +def test_hybrid_routing_context_ignores_invalid_optional_params() -> None: + from mcp_broker.hybrid_router import HybridRoutingContext + + context = HybridRoutingContext.from_params( + { + "tenant_context": TENANT_CONTEXT, + "team_id": ["not-a-string"], + "quota_snapshot": "not-a-map", + } + ) + + assert context.tenant_context == TENANT_CONTEXT + assert context.team_id is None + assert context.quota_snapshot is None + + +def test_local_only_upstream_routes_to_edge_without_calling_shared_worker() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter, SharedWorkerRoute + + edge = RecordingEdgeCaller() + worker = RecordingSharedWorker() + router = HybridToolRouter( + upstreams={ + "local-store": UpstreamConfig( + name="local-store", + command="local-store", + tool_prefix="local-store", + mode="per_session", + tags=("stateful",), + state_dir="local-store", + ) + }, + shared_worker=SharedWorkerRoute(worker.call_tool), + ) + + result = router.call_tool( + advertised_name="local-store.search", + arguments={"query": "refund"}, + edge_caller=edge.call, + tenant_context=TENANT_CONTEXT, + team_id="team-a", + quota_snapshot=_quota_snapshot(), + ) + + assert result == { + "content": [{"type": "text", "text": "edge:local-store.search"}], + "structuredContent": {"execution_boundary": "local_edge"}, + } + assert edge.calls == [("local-store", "search", {"query": "refund"}, 60)] + assert worker.calls == [] + + +def test_allowlisted_stateless_tool_routes_to_shared_worker_with_same_request_shape() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter, SharedWorkerRoute + + edge = RecordingEdgeCaller() + worker = RecordingSharedWorker() + router = HybridToolRouter( + upstreams={ + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless", "shared-worker"), + ) + }, + shared_worker=SharedWorkerRoute(worker.call_tool), + ) + + result = router.call_tool( + advertised_name="example.echo", + arguments={"message": "hello"}, + edge_caller=edge.call, + tenant_context=TENANT_CONTEXT, + team_id="team-a", + quota_snapshot=_quota_snapshot(tool_name="example.echo"), + ) + + assert result == { + "content": [{"type": "text", "text": "worker:example.echo"}], + "structuredContent": { + "execution_boundary": "shared_worker", + "arguments": {"message": "hello"}, + }, + } + assert edge.calls == [] + assert worker.calls == [ + { + "tenant_context": TENANT_CONTEXT, + "upstream_id": "example-stateless", + "upstream_class": "stateless", + "allowlisted": True, + "requires_local_state": False, + "tool_name": "example.echo", + "arguments": {"message": "hello"}, + "quota_decision": { + "allowed": True, + "reason": "quota_allowed", + "checked_scopes": ["global", "team", "user", "upstream", "tool"], + "audit_event": { + "event_type": "quota_decision", + "tenant_id": "tenant-a", + "workspace_id": "workspace-a", + "user_id": "user-a", + "team_id": "team-a", + "upstream_id": "example-stateless", + "tool_name": "example.echo", + "result": "allowed", + }, + }, + } + ] + + +def test_stateless_upstream_without_allowlist_stays_local_edge() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter, SharedWorkerRoute + + edge = RecordingEdgeCaller() + worker = RecordingSharedWorker() + router = HybridToolRouter( + upstreams={ + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless",), + ) + }, + shared_worker=SharedWorkerRoute(worker.call_tool), + ) + + result = router.call_tool( + advertised_name="example.echo", + arguments={"message": "local"}, + edge_caller=edge.call, + tenant_context=TENANT_CONTEXT, + team_id="team-a", + quota_snapshot=_quota_snapshot(tool_name="example.echo"), + ) + + assert result["structuredContent"] == {"execution_boundary": "local_edge"} + assert edge.calls == [("example-stateless", "echo", {"message": "local"}, 60)] + assert worker.calls == [] + + +def test_shared_eligible_quota_denial_fails_closed_without_edge_fallback() -> None: + from mcp_broker.config import UpstreamConfig + from mcp_broker.hybrid_router import HybridToolRouter, SharedWorkerRoute + + edge = RecordingEdgeCaller() + worker = RecordingSharedWorker() + router = HybridToolRouter( + upstreams={ + "example-stateless": UpstreamConfig( + name="example-stateless", + command="example-stateless", + tool_prefix="example", + mode="shared", + tags=("stateless", "shared-worker"), + ) + }, + shared_worker=SharedWorkerRoute(worker.call_tool), + ) + + result = router.call_tool( + advertised_name="example.echo", + arguments={"message": "blocked"}, + edge_caller=edge.call, + tenant_context=TENANT_CONTEXT, + team_id="team-a", + quota_snapshot=_quota_snapshot( + tool_name="example.echo", + tool_limit=0, + tool_used=0, + ), + ) + + assert result["allowed"] is False + assert result["reason"] == "quota_denied" + assert result["audit_event"]["denial_reason"] == "tool_quota_exceeded" + assert edge.calls == [] + assert worker.calls == [ + { + "tenant_context": TENANT_CONTEXT, + "upstream_id": "example-stateless", + "upstream_class": "stateless", + "allowlisted": True, + "requires_local_state": False, + "tool_name": "example.echo", + "arguments": {"message": "blocked"}, + "quota_decision": { + "allowed": False, + "reason": "tool_quota_exceeded", + "blocked_scope": "tool", + "checked_scopes": ["global", "team", "user", "upstream", "tool"], + "audit_event": { + "event_type": "quota_decision", + "tenant_id": "tenant-a", + "workspace_id": "workspace-a", + "user_id": "user-a", + "team_id": "team-a", + "upstream_id": "example-stateless", + "tool_name": "example.echo", + "result": "denied", + "denial_reason": "tool_quota_exceeded", + }, + }, + } + ] + + +def test_unknown_tool_prefix_fails_before_routing() -> None: + from mcp_broker.broker import BrokerToolError + from mcp_broker.hybrid_router import HybridToolRouter, SharedWorkerRoute + + router = HybridToolRouter(upstreams={}, shared_worker=SharedWorkerRoute(lambda **_: {})) + + with pytest.raises(BrokerToolError, match="unknown tool prefix: missing"): + router.call_tool( + advertised_name="missing.search", + arguments={}, + edge_caller=RecordingEdgeCaller().call, + tenant_context=TENANT_CONTEXT, + team_id="team-a", + quota_snapshot=_quota_snapshot(), + ) + + +class RecordingEdgeCaller: + def __init__(self) -> None: + self.calls: list[tuple[str, str, dict[str, object], int]] = [] + + def call( + self, + upstream_name: str, + upstream_tool_name: str, + arguments: dict[str, object], + timeout_seconds: int, + ) -> dict[str, object]: + self.calls.append((upstream_name, upstream_tool_name, arguments, timeout_seconds)) + return { + "content": [{"type": "text", "text": f"edge:{upstream_name}.{upstream_tool_name}"}], + "structuredContent": {"execution_boundary": "local_edge"}, + } + + +class RecordingSharedWorker: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def call_tool(self, **kwargs: object) -> dict[str, object]: + self.calls.append(dict(kwargs)) + quota_decision = kwargs["quota_decision"] + if isinstance(quota_decision, dict) and quota_decision.get("allowed") is False: + return { + "allowed": False, + "reason": "quota_denied", + "result": None, + "audit_event": { + "event_type": "shared_worker_call", + "tenant_id": "tenant-a", + "workspace_id": "workspace-a", + "user_id": "user-a", + "upstream_id": kwargs["upstream_id"], + "tool_name": kwargs["tool_name"], + "result": "denied", + "denial_reason": quota_decision["reason"], + }, + } + tool_name = kwargs["tool_name"] + return { + "allowed": True, + "reason": "worker_call_allowed", + "result": { + "content": [{"type": "text", "text": f"worker:{tool_name}"}], + "structuredContent": { + "execution_boundary": "shared_worker", + "arguments": kwargs["arguments"], + }, + }, + } + + +def _quota_snapshot( + *, + tool_name: str = "local-store.search", + tool_limit: int = 10, + tool_used: int = 0, +) -> dict[str, object]: + return { + "kill_switches": { + "global": False, + "teams": [], + "users": [], + "upstreams": [], + "tools": [], + }, + "limits": { + "global": {"limit": 100, "used": 0}, + "teams": {"team-a": {"limit": 100, "used": 0}}, + "users": {"user-a": {"limit": 100, "used": 0}}, + "upstreams": {"example-stateless": {"limit": 100, "used": 0}}, + "tools": {tool_name: {"limit": tool_limit, "used": tool_used}}, + }, + }