From bb1776a858dfbbf7a52cc8c09a438557bce6d72f Mon Sep 17 00:00:00 2001 From: codex Date: Mon, 27 Jul 2026 02:38:49 +0000 Subject: [PATCH] refactor(supervisor): separate MCP dispatch from transport --- bot_bottle/gateway/supervisor/mcp_dispatch.py | 79 ++++++++++++++++ bot_bottle/gateway/supervisor/server.py | 92 +++++++++---------- tests/unit/test_supervise_server.py | 32 +++++-- tests/unit/test_supervisor_mcp_dispatch.py | 81 ++++++++++++++++ 4 files changed, 227 insertions(+), 57 deletions(-) create mode 100644 bot_bottle/gateway/supervisor/mcp_dispatch.py create mode 100644 tests/unit/test_supervisor_mcp_dispatch.py diff --git a/bot_bottle/gateway/supervisor/mcp_dispatch.py b/bot_bottle/gateway/supervisor/mcp_dispatch.py new file mode 100644 index 00000000..79f7a18a --- /dev/null +++ b/bot_bottle/gateway/supervisor/mcp_dispatch.py @@ -0,0 +1,79 @@ +"""Framework-neutral MCP method and tool dispatch.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Callable, Protocol + +from bot_bottle.gateway.egress.context import resolve_client_context +from bot_bottle.gateway.egress.schema import route_to_yaml_dict +from bot_bottle.gateway.policy_resolver import PolicyResolver +from bot_bottle.supervisor import types as _sv + + +class Request(Protocol): + @property + def method(self) -> str: ... + + @property + def params(self) -> dict[str, object]: ... + + +class MethodNotFoundError(Exception): + """Raised when a JSON-RPC method has no MCP handler.""" + + +Handler = Callable[[dict[str, object]], object] + + +@dataclass(frozen=True) +class Handlers: + initialize: Handler + tools_list: Handler + list_routes: Handler + check_proposal: Handler + propose: Handler + + +def dispatch(request: Request, handlers: Handlers) -> object: + """Route one parsed request without depending on the HTTP server.""" + if request.method == "initialize": + return handlers.initialize(request.params) + if request.method == "notifications/initialized": + return None + if request.method == "tools/list": + return handlers.tools_list(request.params) + if request.method != "tools/call": + raise MethodNotFoundError(request.method) + + tool = request.params.get("name") + if tool == _sv.TOOL_LIST_EGRESS_ROUTES: + return handlers.list_routes(request.params) + if tool == _sv.TOOL_CHECK_PROPOSAL: + return handlers.check_proposal(request.params) + return handlers.propose(request.params) + + +def resolved_routes_payload( + resolver: PolicyResolver, + source_ip: str, + identity_token: str, +) -> dict[str, object]: + """Render the calling bottle's routes, failing closed to an empty list.""" + config, _slug, _tokens = resolve_client_context( + resolver, source_ip, identity_token, + ) + body = json.dumps( + {"routes": [route_to_yaml_dict(route) for route in config.routes]}, + indent=2, + ) + return {"content": [{"type": "text", "text": body}], "isError": False} + + +__all__ = [ + "Handlers", + "MethodNotFoundError", + "dispatch", + "resolved_routes_payload", +] diff --git a/bot_bottle/gateway/supervisor/server.py b/bot_bottle/gateway/supervisor/server.py index 8c4cce64..8bfd2fd8 100644 --- a/bot_bottle/gateway/supervisor/server.py +++ b/bot_bottle/gateway/supervisor/server.py @@ -58,10 +58,15 @@ import typing from dataclasses import dataclass from bot_bottle.constants import IDENTITY_HEADER -from bot_bottle.gateway.egress.context import resolve_client_context -from bot_bottle.gateway.egress.schema import load_config, route_to_yaml_dict +from bot_bottle.gateway.egress.schema import load_config from bot_bottle.gateway.egress.types import LOG_OFF from bot_bottle.gateway.policy_resolver import PolicyResolveError, PolicyResolver +from bot_bottle.gateway.supervisor.mcp_dispatch import ( + Handlers as DispatchHandlers, + MethodNotFoundError, + dispatch, + resolved_routes_payload, +) from bot_bottle.supervisor import types as _sv @@ -611,6 +616,11 @@ class MCPHandler(http.server.BaseHTTPRequestHandler): try: result = self._dispatch(req, config) + except MethodNotFoundError as e: + self._write_jsonrpc( + jsonrpc_error(req.id, ERR_METHOD_NOT_FOUND, f"method not found: {e}"), + ) + return except _RpcClientError as e: self._write_jsonrpc(jsonrpc_error(req.id, e.code, e.message)) return @@ -633,41 +643,37 @@ class MCPHandler(http.server.BaseHTTPRequestHandler): self._write_jsonrpc(jsonrpc_result(req.id, result)) def _dispatch(self, req: JsonRpcRequest, config: ServerConfig) -> object: - method = req.method - if method == "initialize": - return handle_initialize(req.params) - if method == "notifications/initialized": - return None # ack-only - if method == "tools/list": - return handle_tools_list(req.params) - if method == "tools/call": - # `list-egress-routes` is read-only introspection. The shared gateway - # has no static route table (routes are resolved per request by - # source IP), so answer it from the calling bottle's resolved policy. - # Otherwise the agent sees an empty allowlist and composes an egress - # proposal that *replaces* the live routes instead of extending them - # — silently dropping base routes like api.anthropic.com on approval. - if req.params.get("name") == _sv.TOOL_LIST_EGRESS_ROUTES: - return self._resolved_routes_payload() - resolver = self._resolver_or_fail() - source_ip = self.client_address[0] - token = self._identity_token() - # `check-proposal` is a non-blocking read of the calling bottle's - # own queue — attributed by (source_ip, identity_token) like a - # proposal, but it never queues or blocks. - if req.params.get("name") == _sv.TOOL_CHECK_PROPOSAL: - return handle_check_proposal( - req.params, resolver=resolver, - source_ip=source_ip, identity_token=token, - ) - # The control plane attributes the proposal to the source-IP + token - # resolved bottle, so the one shared queue holds each bottle's - # proposal under its own id — no slug is asserted by this daemon. - return handle_tools_call( - req.params, config, resolver=resolver, - source_ip=source_ip, identity_token=token, + def check(params: dict[str, object]) -> object: + return handle_check_proposal( + params, + resolver=self._resolver_or_fail(), + source_ip=self.client_address[0], + identity_token=self._identity_token(), ) - raise _RpcClientError(ERR_METHOD_NOT_FOUND, f"method not found: {method}") + + def propose(params: dict[str, object]) -> object: + return handle_tools_call( + params, + config, + resolver=self._resolver_or_fail(), + source_ip=self.client_address[0], + identity_token=self._identity_token(), + ) + + return dispatch( + req, + DispatchHandlers( + initialize=handle_initialize, + tools_list=handle_tools_list, + list_routes=lambda _params: resolved_routes_payload( + self._resolver_or_fail(), + self.client_address[0], + self._identity_token(), + ), + check_proposal=check, + propose=propose, + ), + ) def _identity_token(self) -> str: """The agent's per-bottle identity token from the request header (the @@ -686,20 +692,6 @@ class MCPHandler(http.server.BaseHTTPRequestHandler): raise _RpcInternalError("supervise server has no policy resolver") return resolver - def _resolved_routes_payload(self) -> dict[str, object]: - """The calling bottle's live egress routes as the `list-egress-routes` - JSON payload, resolved by (source_ip, identity token). Fail-closed: an - unattributed source or an unreachable orchestrator yields an empty route - list (never another bottle's), courtesy of `resolve_client_context`.""" - resolver = self._resolver_or_fail() - conf, _slug, _tokens = resolve_client_context( - resolver, self.client_address[0], self._identity_token(), - ) - body = json.dumps( - {"routes": [route_to_yaml_dict(r) for r in conf.routes]}, indent=2, - ) - return {"content": [{"type": "text", "text": body}], "isError": False} - def _write_jsonrpc(self, body: bytes) -> None: self.send_response(200) self.send_header("Content-Type", "application/json") diff --git a/tests/unit/test_supervise_server.py b/tests/unit/test_supervise_server.py index a4c5117d..61f30e39 100644 --- a/tests/unit/test_supervise_server.py +++ b/tests/unit/test_supervise_server.py @@ -13,6 +13,7 @@ import tempfile import threading import time import types +import typing import unittest from pathlib import Path @@ -47,6 +48,7 @@ from bot_bottle.gateway.supervisor.server import ( jsonrpc_error, jsonrpc_result, parse_jsonrpc, + resolved_routes_payload, validate_proposed_file, ) @@ -701,9 +703,14 @@ class TestResolvedRoutesPayload(unittest.TestCase): " - host: api.anthropic.com\n" " - host: www.google.com\n" ) - payload = _handler( - _FakeSuperviseResolver(bottle_id="b1", policy=policy) - )._resolved_routes_payload() + payload = resolved_routes_payload( + typing.cast( + supervise_server.PolicyResolver, + _FakeSuperviseResolver(bottle_id="b1", policy=policy), + ), + _SRC, + _TOK, + ) assert payload is not None self.assertFalse(payload["isError"]) # type: ignore[index] data = json.loads(payload["content"][0]["text"]) # type: ignore[index] @@ -713,9 +720,14 @@ class TestResolvedRoutesPayload(unittest.TestCase): def test_orchestrator_error_fails_closed_to_empty(self) -> None: # resolve_client_context swallows resolver errors → deny-all (empty), # never another bottle's routes. - payload = _handler( - _FakeSuperviseResolver(raises=True) - )._resolved_routes_payload() + payload = resolved_routes_payload( + typing.cast( + supervise_server.PolicyResolver, + _FakeSuperviseResolver(raises=True), + ), + _SRC, + _TOK, + ) assert payload is not None data = json.loads(payload["content"][0]["text"]) # type: ignore[index] self.assertEqual([], data["routes"]) @@ -724,7 +736,13 @@ class TestResolvedRoutesPayload(unittest.TestCase): # A server without a resolver is a misconfig, not a mode: raise rather # than list anything. with self.assertRaises(_RpcInternalError): - _handler(None)._resolved_routes_payload() + _handler(None)._dispatch( + parse_jsonrpc( + b'{"jsonrpc":"2.0","id":1,"method":"tools/call",' + b'"params":{"name":"list-egress-routes"}}', + ), + ServerConfig(), + ) class TestNonBlockingSupervise(unittest.TestCase): diff --git a/tests/unit/test_supervisor_mcp_dispatch.py b/tests/unit/test_supervisor_mcp_dispatch.py new file mode 100644 index 00000000..ab3b2d91 --- /dev/null +++ b/tests/unit/test_supervisor_mcp_dispatch.py @@ -0,0 +1,81 @@ +"""Unit tests for framework-neutral supervisor MCP dispatch.""" + +from __future__ import annotations + +import unittest +from dataclasses import dataclass + +from bot_bottle.gateway.supervisor.mcp_dispatch import ( + Handlers, + MethodNotFoundError, + dispatch, +) +from bot_bottle.supervisor import types as _sv + + +@dataclass(frozen=True) +class _Request: + method: str + params: dict[str, object] + + +class TestDispatch(unittest.TestCase): + def setUp(self) -> None: + self.calls: list[str] = [] + + def handler(name: str): + def call(_params: dict[str, object]) -> str: + self.calls.append(name) + return name + + return call + + self.handlers = Handlers( + initialize=handler("initialize"), + tools_list=handler("tools_list"), + list_routes=handler("list_routes"), + check_proposal=handler("check_proposal"), + propose=handler("propose"), + ) + + def request(self, method: str, **params: object) -> _Request: + return _Request(method=method, params=params) + + def test_routes_protocol_methods(self) -> None: + self.assertEqual( + "initialize", dispatch(self.request("initialize"), self.handlers), + ) + self.assertEqual( + "tools_list", dispatch(self.request("tools/list"), self.handlers), + ) + self.assertIsNone( + dispatch(self.request("notifications/initialized"), self.handlers), + ) + + def test_routes_each_tool_class(self) -> None: + self.assertEqual( + "list_routes", + dispatch( + self.request("tools/call", name=_sv.TOOL_LIST_EGRESS_ROUTES), + self.handlers, + ), + ) + self.assertEqual( + "check_proposal", + dispatch( + self.request("tools/call", name=_sv.TOOL_CHECK_PROPOSAL), + self.handlers, + ), + ) + self.assertEqual( + "propose", + dispatch(self.request("tools/call", name=_sv.TOOL_EGRESS_ALLOW), self.handlers), + ) + + def test_unknown_method_is_typed(self) -> None: + with self.assertRaisesRegex(MethodNotFoundError, "unknown"): + dispatch(self.request("unknown"), self.handlers) + + +if __name__ == "__main__": + unittest.main()