From 72f8a1d04025543e5612d9ef9ab98e386d9f6967 Mon Sep 17 00:00:00 2001 From: Alex Soffronow Pagonidis <237136924+alex-clickhouse@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:09:27 +0000 Subject: [PATCH] Port nerve/mcp_server/ to the mcp 2.x lowlevel API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A fresh install has been failing at startup with ImportError: cannot import name 'request_ctx' from 'mcp.server.lowlevel.server' mcp was never declared in pyproject.toml — it arrived transitively via claude-agent-sdk, whose constraint is `mcp<3.0.0,>=1.23.0`. Nothing else in the tree capped it, so once mcp 2.0.0 landed on PyPI every fresh resolve picked it up. mcp 2.0 is an API redesign of the lowlevel server, so this ports the endpoint and declares `mcp>=2,<3` explicitly. Declaring it is correct regardless: three modules under nerve/mcp_server/ import mcp directly. What changed upstream, and what it forced here: - Handler registration moved from `@server.list_tools()` / `@server.call_tool()` decorators to `Server(..., on_list_tools=, on_call_tool=)` constructor callbacks. `Server.call_tool` no longer exists. - The request context is now an argument, not a contextvar. `request_ctx` is gone with no replacement, so `_resolve_client_info` and `_bound_identity_from_request` take a `ServerRequestContext` explicitly and `build_ctx_resolver`'s closure threads it down from the call handler. - `on_list_tools` returns `ListToolsResult` rather than `list[Tool]`, and `on_call_tool` receives `CallToolRequestParams` rather than positional `(name, arguments)`. - Model attributes are snake_case in Python now (the camelCase wire aliases are unchanged). Constructing with `isError=` still works via `populate_by_name`, but reads had to move: `client_params.clientInfo` -> `.client_info`. Two things worth a reviewer's attention: mcp 1.x validated tool arguments against `inputSchema` inside its `call_tool` decorator (`validate_input=True` by default). The 2.x callback does not, so porting mechanically would have silently dropped argument validation on an endpoint external clients can reach. `build_mcp_server` now validates explicitly with jsonschema — declared as a direct dependency, since we import it — and returns the same isError shape as the other failure paths. Tests pin the behaviour, including that a rejected call never reaches the handler. The `clientInfo` -> `client_info` rename was invisible to the entire existing suite: every unit test builds a context with `session=None`, so that line never executed, and all 3303 tests passed with it broken. Over real HTTP it raised AttributeError inside the resolver, which `_call_tool` turns into "Context error" — every external MCP tool call would have failed. The two new `test_mcp_http_integration` cases drive `initialize` -> `tools/call` through the real transport and assert satellite attribution, which is what caught it. The module docstring had claimed a tools/call flow that was never actually there. Verified: 3310 tests pass under mcp 2.0.0 on a fresh `--refresh` resolve; `import nerve.gateway.server` clean; the SDK's in-process MCP bridge (which serves every agent tool, not just /mcp/v1) round-trips a real tool call on its mcp-2.x branch; and forcing mcp 1.29.0 fails at import, confirming the `>=2` floor is load-bearing. Fixes #316 Co-Authored-By: Claude Opus 5 --- nerve/mcp_server/http.py | 49 ++++--- nerve/mcp_server/server.py | 118 +++++++++++----- pyproject.toml | 20 ++- tests/test_mcp_http_integration.py | 112 +++++++++++++++ tests/test_mcp_server.py | 210 ++++++++++++++++++++++------- tests/test_mcp_session_binding.py | 30 ++--- 6 files changed, 416 insertions(+), 123 deletions(-) diff --git a/nerve/mcp_server/http.py b/nerve/mcp_server/http.py index bcf5706f9..072b1aa68 100644 --- a/nerve/mcp_server/http.py +++ b/nerve/mcp_server/http.py @@ -29,7 +29,7 @@ import logging from typing import TYPE_CHECKING, Callable -from mcp.server.lowlevel.server import request_ctx +from mcp.server.context import ServerRequestContext from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from starlette.types import Receive, Scope, Send @@ -71,22 +71,28 @@ async def _send_status(send: Send, status: int, message: str) -> None: await send({"type": "http.response.body", "body": body, "more_body": False}) -def _resolve_client_info() -> tuple[str | None, str | None, str | None]: - """Read client metadata from the active MCP request context. +def _resolve_client_info( + rctx: ServerRequestContext | None, +) -> tuple[str | None, str | None, str | None]: + """Read client metadata from the supplied MCP request context. Returns ``(client_name, mcp_session_id, request_path)`` — any field may be ``None`` if the corresponding data isn't available (e.g. during the very first ``initialize`` call the session may not yet have ``client_params``). + + mcp 2.x hands the request context to handlers as an argument rather + than exposing it through a contextvar, so callers thread it down from + the handler. ``None`` is accepted for callers with no live request. """ - try: - rctx = request_ctx.get() - except LookupError: + if rctx is None: return None, None, None client_name: str | None = None if rctx.session and rctx.session.client_params: - info = rctx.session.client_params.clientInfo + # mcp 2.x exposes the model fields as snake_case in Python; the wire + # form is still camelCase (`clientInfo`) via the alias generator. + info = rctx.session.client_params.client_info if info is not None: client_name = info.name @@ -105,6 +111,7 @@ def _resolve_client_info() -> tuple[str | None, str | None, str | None]: def _bound_identity_from_request( config: "NerveConfig", + rctx: ServerRequestContext | None, ) -> tuple[str | None, dict[str, str]]: """Session id bound into the request's bearer token, if any. @@ -120,9 +127,7 @@ def _bound_identity_from_request( """ if not config.auth.jwt_secret: return None, {} - try: - rctx = request_ctx.get() - except LookupError: + if rctx is None: return None, {} request = getattr(rctx, "request", None) if request is None: @@ -150,28 +155,34 @@ def _bound_identity_from_request( return bound_session_id(payload), runtime -def _bound_session_from_request(config: "NerveConfig") -> str | None: +def _bound_session_from_request( + config: "NerveConfig", + rctx: ServerRequestContext | None, +) -> str | None: """Backward-compatible session-only view used by tests/callers.""" - return _bound_identity_from_request(config)[0] + return _bound_identity_from_request(config, rctx)[0] def build_ctx_resolver(engine: "AgentEngine", resolver: SatelliteSessionResolver): """Build the per-call_tool ``ToolContext`` resolver closure. - The Server's ``call_tool`` handler invokes this for every tool call - to attribute the call to the correct session: a session-bound token - (backend-managed agents) binds directly to its engine session; - everything else goes through satellite attribution. Per-call + The Server's ``call_tool`` handler invokes this for every tool call, + passing the request's ``ServerRequestContext``, to attribute the call + to the correct session: a session-bound token (backend-managed + agents) binds directly to its engine session; everything else goes + through satellite attribution. Per-call resolution is cheap (the satellite session id is deterministic and the underlying ``get_session`` / ``create_session`` check is O(1) on the indexed primary key). """ - async def _resolve() -> ToolContext: - session_id, runtime_metadata = _bound_identity_from_request(engine.config) + async def _resolve(rctx: ServerRequestContext | None = None) -> ToolContext: + session_id, runtime_metadata = _bound_identity_from_request( + engine.config, rctx + ) if session_id is None: - client_name, mcp_session_id, _ = _resolve_client_info() + client_name, mcp_session_id, _ = _resolve_client_info(rctx) if mcp_session_id is None: # Stateless requests / pre-initialize calls can land here. diff --git a/nerve/mcp_server/server.py b/nerve/mcp_server/server.py index 149a10ec8..cfcae5ff1 100644 --- a/nerve/mcp_server/server.py +++ b/nerve/mcp_server/server.py @@ -7,8 +7,17 @@ The ``ctx_resolver`` callable is invoked per ``call_tool`` request to build a fresh :class:`ToolContext` for the satellite session that owns -this MCP connection. It returns an awaitable so the resolver can fetch -state from the DB on first use and cache it for subsequent calls. +this MCP connection. It receives the request's +:class:`~mcp.server.context.ServerRequestContext` — mcp 2.x hands that to +handlers as an argument rather than exposing it through a contextvar — and +returns an awaitable so the resolver can fetch state from the DB on first +use and cache it for subsequent calls. + +Argument validation against each tool's ``inputSchema`` happens here, +explicitly. mcp 1.x's ``@server.call_tool()`` decorator did it for us by +default (``validate_input=True``); the 2.x ``on_call_tool`` callback does +not, so relying on the library would have silently dropped validation on an +endpoint external clients can reach. The ``audit_writer`` callable is invoked after every successful tool call to persist an ``external_tool_call`` event into ``session_events``. @@ -21,15 +30,24 @@ import logging from typing import Any, Awaitable, Callable +import jsonschema +from mcp.server.context import ServerRequestContext from mcp.server.lowlevel import Server -from mcp.types import CallToolResult, TextContent, Tool +from mcp.types import ( + CallToolRequestParams, + CallToolResult, + ListToolsResult, + PaginatedRequestParams, + TextContent, + Tool, +) from nerve.agent.tools import ToolContext, ToolRegistry, ToolResult logger = logging.getLogger(__name__) -CtxResolver = Callable[[], Awaitable[ToolContext]] +CtxResolver = Callable[[ServerRequestContext], Awaitable[ToolContext]] AuditWriter = Callable[[ToolContext, str, dict, ToolResult, float, bool], Awaitable[None]] @@ -67,46 +85,68 @@ def build_mcp_server( """ import time - server: Server = Server(name=name, version=version) - - @server.list_tools() - async def _list_tools() -> list[Tool]: - return [ - Tool( - name=spec.name, - description=spec.description, - inputSchema=spec.input_schema, - ) - for spec in registry.list(include_hoa=include_hoa) - ] - - @server.call_tool() - async def _call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: + def _error(message: str) -> CallToolResult: + """A failed call, shaped exactly like every other failure here.""" + return CallToolResult( + content=[TextContent(type="text", text=message)], + isError=True, + ) + + async def _list_tools( + rctx: ServerRequestContext, + params: PaginatedRequestParams | None = None, + ) -> ListToolsResult: + return ListToolsResult( + tools=[ + Tool( + name=spec.name, + description=spec.description, + inputSchema=spec.input_schema, + ) + for spec in registry.list(include_hoa=include_hoa) + ] + ) + + async def _call_tool( + rctx: ServerRequestContext, + params: CallToolRequestParams, + ) -> CallToolResult: + name = params.name + arguments: dict[str, Any] = dict(params.arguments or {}) + spec = registry.get(name) if spec is None: - return CallToolResult( - content=[TextContent(type="text", text=f"Unknown tool: {name!r}")], - isError=True, - ) + return _error(f"Unknown tool: {name!r}") # HoA gating: registry.list() filters by include_hoa, but a # malicious caller could still invoke a HoA tool by name. Enforce # the same allowlist here. if not include_hoa and name.startswith("hoa_"): - return CallToolResult( - content=[TextContent(type="text", text=f"Tool not available: {name!r}")], - isError=True, - ) + return _error(f"Tool not available: {name!r}") + + # Validate arguments against the tool's declared inputSchema before + # any handler sees them. mcp 1.x's call_tool decorator did this by + # default; the 2.x callback does not, and this endpoint is reachable + # by external clients, so the check lives here explicitly rather + # than depending on a library default. + if spec.input_schema: + try: + jsonschema.validate(instance=arguments, schema=spec.input_schema) + except jsonschema.ValidationError as e: + logger.info("Invalid arguments for %s: %s", name, e.message) + return _error(f"Invalid arguments for {name!r}: {e.message}") + except jsonschema.SchemaError: + # A malformed schema is our bug, not the caller's. Refuse the + # call rather than dispatching unvalidated arguments. + logger.exception("Tool %s has an invalid inputSchema", name) + return _error(f"Tool {name!r} has an invalid input schema") start = time.monotonic() try: - ctx = await ctx_resolver() + ctx = await ctx_resolver(rctx) except Exception as e: logger.exception("Failed to resolve ToolContext for %s", name) - return CallToolResult( - content=[TextContent(type="text", text=f"Context error: {e}")], - isError=True, - ) + return _error(f"Context error: {e}") try: result = await spec.handler(ctx, arguments) @@ -126,10 +166,7 @@ async def _call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: ) except Exception: logger.exception("Audit writer failed for %s", name) - return CallToolResult( - content=[TextContent(type="text", text=f"Tool error: {e}")], - isError=True, - ) + return _error(f"Tool error: {e}") duration_ms = (time.monotonic() - start) * 1000.0 if audit_writer is not None: @@ -161,4 +198,11 @@ async def _call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: return CallToolResult(content=content, isError=result.is_error) - return server + # mcp 2.x registers handlers as constructor callbacks; the 1.x + # ``@server.list_tools()`` / ``@server.call_tool()`` decorators are gone. + return Server( + name=name, + version=version, + on_list_tools=_list_tools, + on_call_tool=_call_tool, + ) diff --git a/pyproject.toml b/pyproject.toml index aae066895..c1254ea39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,25 @@ dependencies = [ "aiosqlite>=0.21.0", "pyyaml>=6.0", "python-telegram-bot>=21.0", - "claude-agent-sdk>=0.2.82", + # Floor raised from 0.2.82 for mcp 2.x: 0.2.140 is the first release whose + # own constraint is `mcp<3.0.0` (0.2.100 through 0.2.139 all declare + # `mcp<2.0.0`). It matters because the SDK builds the in-process MCP server + # that serves EVERY agent tool, so an SDK without 2.x support breaks the + # whole Claude backend, not just /mcp/v1. 0.2.82 is the dangerous case: it + # declares `mcp>=1.23.0` with no upper bound, so `0.2.82 + mcp 2.x` is a + # resolvable combination that cannot work. + "claude-agent-sdk>=0.2.140", + # Declared explicitly because nerve/mcp_server/ imports `mcp` directly + # rather than only through claude-agent-sdk. Both bounds are load-bearing: + # nerve/mcp_server/ is written against the 2.x lowlevel API (constructor + # callbacks, request context as a handler argument), which does not exist + # in 1.x, and 3.x has not been vetted. + "mcp>=2,<3", + # Used directly by nerve/mcp_server/server.py to validate tool arguments + # against each tool's inputSchema. mcp 1.x's call_tool decorator did this + # for us; the 2.x callback does not, so the check is ours now and the + # dependency is declared rather than borrowed from claude-agent-sdk. + "jsonschema>=4.20", "apscheduler>=3.11.0", "pyjwt>=2.10.0", "bcrypt>=4.2.0", diff --git a/tests/test_mcp_http_integration.py b/tests/test_mcp_http_integration.py index 4fcd206b8..25ceaf520 100644 --- a/tests/test_mcp_http_integration.py +++ b/tests/test_mcp_http_integration.py @@ -183,3 +183,115 @@ def test_mcp_list_tools(app_with_mcp): assert "notify" in names # HoA tools excluded by default. assert not any(n.startswith("hoa_") for n in names) + + +def test_mcp_call_tool_attributes_to_satellite_session(app_with_mcp): + """``tools/call`` over real HTTP, asserting satellite attribution. + + This is the only flow that exercises the ``ServerRequestContext`` mcp 2.x + hands to handlers: ``build_ctx_resolver`` reads the bearer token and the + client's ``clientInfo`` off it to decide which session a call belongs to. + Every unit test passes a stand-in context object, so a real request through + the transport is the only thing that proves that wiring end to end — and + the resulting session id is the observable proof the context arrived + populated rather than empty. + """ + from nerve.agent.tools import ToolResult, ToolSpec + from nerve.gateway import server as gw + + with TestClient(app_with_mcp) as client: + async def _probe(ctx, args): + return ToolResult.text(f"probe:{ctx.session_id}") + + # Registered on the live registry the manager closes over; lookup + # happens per call, so a late registration is visible. + gw._engine.registry.register(ToolSpec( + name="probe_attribution", + description="report the resolved session id", + input_schema={"type": "object", "properties": {}, "required": []}, + handler=_probe, + )) + + init_resp = _post_jsonrpc(client, { + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "probe-client", "version": "0.1"}, + }, + }) + sid = init_resp.headers["mcp-session-id"] + + notif_resp = _post_jsonrpc(client, { + "jsonrpc": "2.0", "method": "notifications/initialized", + }, session_id=sid) + assert notif_resp.status_code in (200, 202), notif_resp.text + + call_resp = _post_jsonrpc(client, { + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "probe_attribution", "arguments": {}}, + }, session_id=sid) + assert call_resp.status_code == 200, call_resp.text + body = _parse_response(call_resp) + assert "error" not in body, body + result = body["result"] + assert not result.get("isError"), result + + text = result["content"][0]["text"] + # external:: — the client name comes + # from the initialize handshake above, so seeing it here means the + # request context reached the resolver intact. + assert text.startswith("probe:external:"), text + assert "probe-client" in text, text + + +def test_mcp_call_tool_rejects_invalid_arguments(app_with_mcp): + """Schema validation is enforced on the real transport, not just in unit tests. + + mcp 1.x validated arguments inside its ``call_tool`` decorator; the 2.x + callback does not, so ``build_mcp_server`` does it. Worth asserting over + HTTP because this endpoint is the one external clients can reach. + """ + from nerve.agent.tools import ToolResult, ToolSpec + from nerve.gateway import server as gw + + with TestClient(app_with_mcp) as client: + calls: list[dict] = [] + + async def _typed(ctx, args): + calls.append(args) + return ToolResult.text("should not run") + + gw._engine.registry.register(ToolSpec( + name="probe_typed", + description="requires an integer count", + input_schema={ + "type": "object", + "properties": {"count": {"type": "integer"}}, + "required": ["count"], + }, + handler=_typed, + )) + + init_resp = _post_jsonrpc(client, { + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "probe-client", "version": "0.1"}, + }, + }) + sid = init_resp.headers["mcp-session-id"] + _post_jsonrpc(client, { + "jsonrpc": "2.0", "method": "notifications/initialized", + }, session_id=sid) + + call_resp = _post_jsonrpc(client, { + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "probe_typed", "arguments": {"count": "nope"}}, + }, session_id=sid) + assert call_resp.status_code == 200, call_resp.text + result = _parse_response(call_resp)["result"] + assert result["isError"] is True, result + assert "Invalid arguments" in result["content"][0]["text"] + assert calls == [], "handler ran despite invalid arguments" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 89d4ec8d0..e21fc3034 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -12,18 +12,16 @@ from __future__ import annotations +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock import pytest from mcp.types import ( - CallToolRequest, CallToolRequestParams, - ClientRequest, - ListToolsRequest, - ListToolsResult, CallToolResult, + ListToolsResult, ) from nerve.agent.tools import ToolContext, ToolRegistry, ToolResult, ToolSpec @@ -36,10 +34,18 @@ def _make_spec( response_text: str = "ok", is_error: bool = False, raises: Exception | None = None, + input_schema: dict | None = None, + seen: list[dict] | None = None, ) -> ToolSpec: - """Build a deterministic ToolSpec for protocol tests.""" + """Build a deterministic ToolSpec for protocol tests. + + ``seen`` records the arguments each invocation receives, so a test can + assert a rejected call never reached the handler at all. + """ async def handler(ctx: ToolContext, args: dict) -> ToolResult: + if seen is not None: + seen.append(args) if raises is not None: raise raises return ToolResult.text(response_text, is_error=is_error) @@ -47,7 +53,11 @@ async def handler(ctx: ToolContext, args: dict) -> ToolResult: return ToolSpec( name=name, description=f"test tool {name}", - input_schema={"type": "object", "properties": {}, "required": []}, + input_schema=( + input_schema + if input_schema is not None + else {"type": "object", "properties": {}, "required": []} + ), handler=handler, ) @@ -57,7 +67,9 @@ async def _resolve_static_ctx(session_id: str = "external:test:s1") -> ToolConte def _ctx_resolver(session_id: str = "external:test:s1"): - async def _r() -> ToolContext: + # Takes the request context mcp 2.x passes to handlers; these tests don't + # exercise per-request attribution, so it's accepted and ignored. + async def _r(rctx: Any = None) -> ToolContext: return ToolContext(session_id=session_id) return _r @@ -72,12 +84,9 @@ async def test_list_tools_returns_registry_entries(self): registry.register(_make_spec("beta")) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - # Invoke the registered list_tools handler directly via the - # SDK's request_handlers dispatch table. - handler = server.request_handlers[ListToolsRequest] - result = await handler(ListToolsRequest(method="tools/list")) - # ServerResult is the wrapping discriminated union. - tools_result: ListToolsResult = result.root + # Invoke the registered list_tools handler directly, bypassing the + # transport. mcp 2.x hands back the result model unwrapped. + tools_result: ListToolsResult = await _invoke_list_tools(server) assert isinstance(tools_result, ListToolsResult) assert {t.name for t in tools_result.tools} == {"alpha", "beta"} @@ -87,9 +96,8 @@ async def test_list_tools_hides_hoa_by_default(self): registry.register(_make_spec("hoa_execute")) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - handler = server.request_handlers[ListToolsRequest] - result = await handler(ListToolsRequest(method="tools/list")) - names = {t.name for t in result.root.tools} + result = await _invoke_list_tools(server) + names = {t.name for t in result.tools} assert names == {"regular"} async def test_list_tools_includes_hoa_when_opted_in(self): @@ -100,16 +108,29 @@ async def test_list_tools_includes_hoa_when_opted_in(self): registry, ctx_resolver=_ctx_resolver(), include_hoa=True, ) - handler = server.request_handlers[ListToolsRequest] - result = await handler(ListToolsRequest(method="tools/list")) - names = {t.name for t in result.root.tools} + result = await _invoke_list_tools(server) + names = {t.name for t in result.tools} assert names == {"regular", "hoa_execute"} -def _build_call_request(name: str, args: dict | None = None) -> CallToolRequest: - return CallToolRequest( - method="tools/call", - params=CallToolRequestParams(name=name, arguments=args or {}), +# mcp 2.x dispatches through `get_request_handler(method)`, keyed by method +# string, and hands the handler the request context plus a params model. The +# old `request_handlers[RequestClass]` table and the `ServerResult` root +# wrapper are both gone, so these two helpers are the whole test seam. +_FAKE_RCTX = SimpleNamespace(request=None, session=None) + + +async def _invoke_list_tools(server: Any) -> ListToolsResult: + entry = server.get_request_handler("tools/list") + return await entry.handler(_FAKE_RCTX, None) + + +async def _invoke_call_tool( + server: Any, name: str, args: dict | None = None, +) -> CallToolResult: + entry = server.get_request_handler("tools/call") + return await entry.handler( + _FAKE_RCTX, CallToolRequestParams(name=name, arguments=args or {}), ) @@ -122,10 +143,9 @@ async def test_dispatches_to_registered_handler(self): registry.register(_make_spec("alpha", response_text="hello-alpha")) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - handler = server.request_handlers[CallToolRequest] - result = await handler(_build_call_request("alpha")) - call_result: CallToolResult = result.root - assert call_result.isError is False + result = await _invoke_call_tool(server, "alpha") + call_result: CallToolResult = result + assert call_result.is_error is False assert call_result.content[0].text == "hello-alpha" async def test_returns_error_for_unknown_tool(self): @@ -133,10 +153,9 @@ async def test_returns_error_for_unknown_tool(self): registry.register(_make_spec("alpha")) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - handler = server.request_handlers[CallToolRequest] - result = await handler(_build_call_request("missing")) - call_result: CallToolResult = result.root - assert call_result.isError is True + result = await _invoke_call_tool(server, "missing") + call_result: CallToolResult = result + assert call_result.is_error is True assert "Unknown tool" in call_result.content[0].text async def test_propagates_handler_is_error(self): @@ -144,10 +163,9 @@ async def test_propagates_handler_is_error(self): registry.register(_make_spec("alpha", response_text="boom", is_error=True)) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - handler = server.request_handlers[CallToolRequest] - result = await handler(_build_call_request("alpha")) - call_result: CallToolResult = result.root - assert call_result.isError is True + result = await _invoke_call_tool(server, "alpha") + call_result: CallToolResult = result + assert call_result.is_error is True assert call_result.content[0].text == "boom" async def test_handler_exception_returns_tool_error(self): @@ -157,10 +175,9 @@ async def test_handler_exception_returns_tool_error(self): ) server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) - handler = server.request_handlers[CallToolRequest] - result = await handler(_build_call_request("crash")) - call_result: CallToolResult = result.root - assert call_result.isError is True + result = await _invoke_call_tool(server, "crash") + call_result: CallToolResult = result + assert call_result.is_error is True assert "explosion" in call_result.content[0].text async def test_hoa_tool_rejected_when_include_hoa_false(self): @@ -170,10 +187,9 @@ async def test_hoa_tool_rejected_when_include_hoa_false(self): registry, ctx_resolver=_ctx_resolver(), include_hoa=False, ) - handler = server.request_handlers[CallToolRequest] - result = await handler(_build_call_request("hoa_execute")) - call_result: CallToolResult = result.root - assert call_result.isError is True + result = await _invoke_call_tool(server, "hoa_execute") + call_result: CallToolResult = result + assert call_result.is_error is True assert "not available" in call_result.content[0].text async def test_audit_writer_called_on_success(self): @@ -185,8 +201,7 @@ async def test_audit_writer_called_on_success(self): audit_writer=audit, ) - handler = server.request_handlers[CallToolRequest] - await handler(_build_call_request("alpha", {"foo": "bar"})) + await _invoke_call_tool(server, "alpha", {"foo": "bar"}) audit.assert_awaited_once() args, _kwargs = audit.call_args @@ -207,11 +222,112 @@ async def test_audit_writer_called_on_exception(self): audit_writer=audit, ) - handler = server.request_handlers[CallToolRequest] - await handler(_build_call_request("crash")) + await _invoke_call_tool(server, "crash") audit.assert_awaited_once() args, _kwargs = audit.call_args _sid, _name, _args, result_obj, _ms, is_error = args assert is_error is True assert "nope" in result_obj.content[0]["text"] + + +_TYPED_SCHEMA = { + "type": "object", + "properties": {"count": {"type": "integer"}, "label": {"type": "string"}}, + "required": ["count"], + "additionalProperties": False, +} + + +@pytest.mark.asyncio +class TestArgumentValidation: + """Arguments are validated against the tool's inputSchema before dispatch. + + mcp 1.x's ``@server.call_tool()`` decorator validated by default + (``validate_input=True``); the 2.x ``on_call_tool`` callback does not. The + check therefore lives in ``build_mcp_server`` itself, and these tests pin + it there — an externally reachable endpoint must not hand unvalidated + arguments to a tool handler, and nothing in the library will complain if + that protection quietly disappears again. + """ + + async def test_wrong_type_is_rejected_before_the_handler_runs(self): + seen: list[dict] = [] + registry = ToolRegistry() + registry.register( + _make_spec("typed", input_schema=_TYPED_SCHEMA, seen=seen), + ) + server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) + + result = await _invoke_call_tool(server, "typed", {"count": "not-an-int"}) + + assert result.is_error is True + assert "Invalid arguments" in result.content[0].text + assert seen == [], "handler ran despite invalid arguments" + + async def test_missing_required_argument_is_rejected(self): + seen: list[dict] = [] + registry = ToolRegistry() + registry.register( + _make_spec("typed", input_schema=_TYPED_SCHEMA, seen=seen), + ) + server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) + + result = await _invoke_call_tool(server, "typed", {"label": "x"}) + + assert result.is_error is True + assert "Invalid arguments" in result.content[0].text + assert seen == [] + + async def test_unknown_property_is_rejected(self): + seen: list[dict] = [] + registry = ToolRegistry() + registry.register( + _make_spec("typed", input_schema=_TYPED_SCHEMA, seen=seen), + ) + server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) + + result = await _invoke_call_tool( + server, "typed", {"count": 1, "surprise": "!"}, + ) + + assert result.is_error is True + assert seen == [] + + async def test_valid_arguments_reach_the_handler(self): + seen: list[dict] = [] + registry = ToolRegistry() + registry.register( + _make_spec( + "typed", response_text="done", + input_schema=_TYPED_SCHEMA, seen=seen, + ), + ) + server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) + + result = await _invoke_call_tool( + server, "typed", {"count": 3, "label": "ok"}, + ) + + assert result.is_error is False + assert result.content[0].text == "done" + assert seen == [{"count": 3, "label": "ok"}] + + async def test_malformed_schema_refuses_the_call(self): + """A broken schema is our bug — refuse rather than skip validation.""" + seen: list[dict] = [] + registry = ToolRegistry() + registry.register( + _make_spec( + "broken", + input_schema={"type": "object", "properties": {"x": {"type": 5}}}, + seen=seen, + ), + ) + server = build_mcp_server(registry, ctx_resolver=_ctx_resolver()) + + result = await _invoke_call_tool(server, "broken", {"x": 1}) + + assert result.is_error is True + assert "invalid input schema" in result.content[0].text + assert seen == [] diff --git a/tests/test_mcp_session_binding.py b/tests/test_mcp_session_binding.py index fece3dd10..4872fecc6 100644 --- a/tests/test_mcp_session_binding.py +++ b/tests/test_mcp_session_binding.py @@ -125,20 +125,16 @@ def get(self, key, default=None): query_params={}, ) fake_rctx = SimpleNamespace(request=fake_request, session=None) - cv_token = mcp_http.request_ctx.set(fake_rctx) - try: - assert mcp_http._bound_session_from_request(cfg) == "engine-sess-7" + assert mcp_http._bound_session_from_request(cfg, fake_rctx) == "engine-sess-7" - # Plain token → no binding (satellite path). - fake_request.headers = _Headers( - {"authorization": f"Bearer {create_token(SECRET)}"}, - ) - assert mcp_http._bound_session_from_request(cfg) is None - finally: - mcp_http.request_ctx.reset(cv_token) + # Plain token → no binding (satellite path). + fake_request.headers = _Headers( + {"authorization": f"Bearer {create_token(SECRET)}"}, + ) + assert mcp_http._bound_session_from_request(cfg, fake_rctx) is None - # No request context set → no binding. - assert mcp_http._bound_session_from_request(cfg) is None + # No request context → no binding. + assert mcp_http._bound_session_from_request(cfg, None) is None @pytest.mark.asyncio async def test_worker_token_adds_runtime_attribution(self, tmp_path): @@ -162,13 +158,9 @@ def get(self, key, default=None): headers=_Headers({"authorization": f"Bearer {token}"}), query_params={}, ) - cv_token = mcp_http.request_ctx.set( - SimpleNamespace(request=fake_request, session=None), + session_id, runtime = mcp_http._bound_identity_from_request( + cfg, SimpleNamespace(request=fake_request, session=None), ) - try: - session_id, runtime = mcp_http._bound_identity_from_request(cfg) - finally: - mcp_http.request_ctx.reset(cv_token) assert session_id == "engine-sess-8" assert runtime == {"worker_id": worker_id, "runtime": "ultracode"} payload = decode_mcp_token(token, SECRET) @@ -181,4 +173,4 @@ async def test_dev_mode_never_binds(self, tmp_path): cfg = NerveConfig.from_dict({"workspace": str(tmp_path)}) cfg.auth.jwt_secret = "" - assert mcp_http._bound_session_from_request(cfg) is None + assert mcp_http._bound_session_from_request(cfg, None) is None