refactor(supervisor): separate MCP dispatch from transport
test / image-input-builds (pull_request) Failing after 13m22s
test / unit (pull_request) Failing after 13m27s
test / coverage (pull_request) Has been skipped
test / integration-docker (pull_request) Has been cancelled
tracker-policy-pr / check-pr (pull_request) Successful in 14s
test / image-input-builds (pull_request) Failing after 13m22s
test / unit (pull_request) Failing after 13m27s
test / coverage (pull_request) Has been skipped
test / integration-docker (pull_request) Has been cancelled
tracker-policy-pr / check-pr (pull_request) Successful in 14s
This commit is contained in:
@@ -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",
|
||||||
|
]
|
||||||
@@ -58,10 +58,15 @@ import typing
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
from bot_bottle.constants import IDENTITY_HEADER
|
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
|
||||||
from bot_bottle.gateway.egress.schema import load_config, route_to_yaml_dict
|
|
||||||
from bot_bottle.gateway.egress.types import LOG_OFF
|
from bot_bottle.gateway.egress.types import LOG_OFF
|
||||||
from bot_bottle.gateway.policy_resolver import PolicyResolveError, PolicyResolver
|
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
|
from bot_bottle.supervisor import types as _sv
|
||||||
|
|
||||||
|
|
||||||
@@ -611,6 +616,11 @@ class MCPHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
result = self._dispatch(req, config)
|
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:
|
except _RpcClientError as e:
|
||||||
self._write_jsonrpc(jsonrpc_error(req.id, e.code, e.message))
|
self._write_jsonrpc(jsonrpc_error(req.id, e.code, e.message))
|
||||||
return
|
return
|
||||||
@@ -633,41 +643,37 @@ class MCPHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._write_jsonrpc(jsonrpc_result(req.id, result))
|
self._write_jsonrpc(jsonrpc_result(req.id, result))
|
||||||
|
|
||||||
def _dispatch(self, req: JsonRpcRequest, config: ServerConfig) -> object:
|
def _dispatch(self, req: JsonRpcRequest, config: ServerConfig) -> object:
|
||||||
method = req.method
|
def check(params: dict[str, object]) -> object:
|
||||||
if method == "initialize":
|
return handle_check_proposal(
|
||||||
return handle_initialize(req.params)
|
params,
|
||||||
if method == "notifications/initialized":
|
resolver=self._resolver_or_fail(),
|
||||||
return None # ack-only
|
source_ip=self.client_address[0],
|
||||||
if method == "tools/list":
|
identity_token=self._identity_token(),
|
||||||
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,
|
|
||||||
)
|
)
|
||||||
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:
|
def _identity_token(self) -> str:
|
||||||
"""The agent's per-bottle identity token from the request header (the
|
"""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")
|
raise _RpcInternalError("supervise server has no policy resolver")
|
||||||
return 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:
|
def _write_jsonrpc(self, body: bytes) -> None:
|
||||||
self.send_response(200)
|
self.send_response(200)
|
||||||
self.send_header("Content-Type", "application/json")
|
self.send_header("Content-Type", "application/json")
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ import tempfile
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
import types
|
import types
|
||||||
|
import typing
|
||||||
import unittest
|
import unittest
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -47,6 +48,7 @@ from bot_bottle.gateway.supervisor.server import (
|
|||||||
jsonrpc_error,
|
jsonrpc_error,
|
||||||
jsonrpc_result,
|
jsonrpc_result,
|
||||||
parse_jsonrpc,
|
parse_jsonrpc,
|
||||||
|
resolved_routes_payload,
|
||||||
validate_proposed_file,
|
validate_proposed_file,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -701,9 +703,14 @@ class TestResolvedRoutesPayload(unittest.TestCase):
|
|||||||
" - host: api.anthropic.com\n"
|
" - host: api.anthropic.com\n"
|
||||||
" - host: www.google.com\n"
|
" - host: www.google.com\n"
|
||||||
)
|
)
|
||||||
payload = _handler(
|
payload = resolved_routes_payload(
|
||||||
_FakeSuperviseResolver(bottle_id="b1", policy=policy)
|
typing.cast(
|
||||||
)._resolved_routes_payload()
|
supervise_server.PolicyResolver,
|
||||||
|
_FakeSuperviseResolver(bottle_id="b1", policy=policy),
|
||||||
|
),
|
||||||
|
_SRC,
|
||||||
|
_TOK,
|
||||||
|
)
|
||||||
assert payload is not None
|
assert payload is not None
|
||||||
self.assertFalse(payload["isError"]) # type: ignore[index]
|
self.assertFalse(payload["isError"]) # type: ignore[index]
|
||||||
data = json.loads(payload["content"][0]["text"]) # 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:
|
def test_orchestrator_error_fails_closed_to_empty(self) -> None:
|
||||||
# resolve_client_context swallows resolver errors → deny-all (empty),
|
# resolve_client_context swallows resolver errors → deny-all (empty),
|
||||||
# never another bottle's routes.
|
# never another bottle's routes.
|
||||||
payload = _handler(
|
payload = resolved_routes_payload(
|
||||||
_FakeSuperviseResolver(raises=True)
|
typing.cast(
|
||||||
)._resolved_routes_payload()
|
supervise_server.PolicyResolver,
|
||||||
|
_FakeSuperviseResolver(raises=True),
|
||||||
|
),
|
||||||
|
_SRC,
|
||||||
|
_TOK,
|
||||||
|
)
|
||||||
assert payload is not None
|
assert payload is not None
|
||||||
data = json.loads(payload["content"][0]["text"]) # type: ignore[index]
|
data = json.loads(payload["content"][0]["text"]) # type: ignore[index]
|
||||||
self.assertEqual([], data["routes"])
|
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
|
# A server without a resolver is a misconfig, not a mode: raise rather
|
||||||
# than list anything.
|
# than list anything.
|
||||||
with self.assertRaises(_RpcInternalError):
|
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):
|
class TestNonBlockingSupervise(unittest.TestCase):
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user