diff --git a/scripts/bench_jsonrpc_codec.py b/scripts/bench_jsonrpc_codec.py new file mode 100644 index 000000000..f461e3e1d --- /dev/null +++ b/scripts/bench_jsonrpc_codec.py @@ -0,0 +1,107 @@ +"""Micro-benchmark: JSONRPCMessage decoding, smart union vs key-presence discriminator. + +Compares the previous bare-union adapter (reconstructed in-process) against the +shipped `jsonrpc_message_adapter` over representative payloads, for both decode +paths used by the SDK: + +- `validate_json(body)` (client SSE/stdio lines) +- `pydantic_core.from_json(body)` + `validate_python(raw)` (server POST path) + +Run with: uv run python scripts/bench_jsonrpc_codec.py [--small N] [--large N] +""" + +import argparse +import json +import time +from collections.abc import Callable +from typing import Any + +import pydantic_core +from mcp_types.jsonrpc import ( + JSONRPCError, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + jsonrpc_message_adapter, +) +from pydantic import TypeAdapter + +# The adapter as it existed before the discriminator change. +bare_union_adapter: TypeAdapter[Any] = TypeAdapter( + JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError +) + +SMALL_REQUEST = json.dumps( + {"jsonrpc": "2.0", "id": 42, "method": "tools/call", "params": {"name": "echo", "arguments": {"text": "hi"}}} +).encode() +SMALL_NOTIFICATION = json.dumps( + {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progressToken": "t", "progress": 0.5}} +).encode() +LARGE_RESULT = json.dumps( + { + "jsonrpc": "2.0", + "id": 42, + "result": { + "content": [{"type": "text", "text": "x" * 256} for _ in range(64)], + "structuredContent": {f"key_{i}": {"value": i, "label": "y" * 32} for i in range(64)}, + }, + } +).encode() + +REPEATS = 5 + + +def bench_pair(old_fn: Callable[[], Any], new_fn: Callable[[], Any], iterations: int) -> tuple[float, float]: + """Return best per-call time in microseconds for each adapter over 2*REPEATS rounds. + + Run order alternates each round so warm-up and CPU-state effects are not + attributed to either adapter. + """ + fns = (old_fn, new_fn) + best = [float("inf"), float("inf")] + for round_index in range(2 * REPEATS): + order = (0, 1) if round_index % 2 == 0 else (1, 0) + for key in order: + start = time.perf_counter() + for _ in range(iterations): + fns[key]() + best[key] = min(best[key], time.perf_counter() - start) + return best[0] / iterations * 1e6, best[1] / iterations * 1e6 + + +def _decode_validate_json(adapter: TypeAdapter[Any], body: bytes) -> Any: + return adapter.validate_json(body, by_name=False) + + +def _decode_two_phase(adapter: TypeAdapter[Any], body: bytes) -> Any: + return adapter.validate_python(pydantic_core.from_json(body), by_name=False) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--small", type=int, default=100_000, help="iterations for small payloads") + parser.add_argument("--large", type=int, default=10_000, help="iterations for the large payload") + args = parser.parse_args() + + payloads = [ + (f"request {len(SMALL_REQUEST)}B", SMALL_REQUEST, args.small), + (f"notification {len(SMALL_NOTIFICATION)}B", SMALL_NOTIFICATION, args.small), + (f"result {len(LARGE_RESULT) / 1024:.1f}KB", LARGE_RESULT, args.large), + ] + decoders: list[tuple[str, Callable[[TypeAdapter[Any], bytes], Any]]] = [ + ("validate_json", _decode_validate_json), + ("two-phase", _decode_two_phase), + ] + print(f"{'payload':<20} {'path':<14} {'smart union':>12} {'discriminator':>14} {'speedup':>8}") + for label, body, iterations in payloads: + for path_label, decode in decoders: + old, new = bench_pair( + lambda: decode(bare_union_adapter, body), + lambda: decode(jsonrpc_message_adapter, body), + iterations, + ) + print(f"{label:<20} {path_label:<14} {old:>10.2f}us {new:>12.2f}us {old / new:>7.2f}x") + + +if __name__ == "__main__": + main() diff --git a/src/mcp-types/mcp_types/jsonrpc.py b/src/mcp-types/mcp_types/jsonrpc.py index fcc3317d8..77765bb54 100644 --- a/src/mcp-types/mcp_types/jsonrpc.py +++ b/src/mcp-types/mcp_types/jsonrpc.py @@ -2,9 +2,9 @@ from __future__ import annotations -from typing import Annotated, Any, Final, Literal +from typing import Annotated, Any, Final, Literal, cast -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Discriminator, Field, Tag, TypeAdapter RequestId = Annotated[int, Field(strict=True)] | str """The ID of a JSON-RPC request.""" @@ -122,4 +122,58 @@ class JSONRPCError(BaseModel): JSONRPCMessage = JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError """Any JSON-RPC envelope that can be decoded off the wire or encoded to be sent.""" -jsonrpc_message_adapter: TypeAdapter[JSONRPCMessage] = TypeAdapter(JSONRPCMessage) + +def _discriminate_jsonrpc_message(value: Any) -> str | None: + """Tag a wire object by key presence per JSON-RPC 2.0. + + Selects exactly one union branch to validate instead of letting smart-union + mode score all four on every message. Classification matches the previous + smart-union outcome for every spec-valid message: a ``method`` member with + a missing, null, or non-int/str ``id`` classifies as a notification + (mirroring ``RequestId``), and ``error`` wins over ``result`` when both are + present. For spec-invalid hybrids that combine ``method`` with ``result``/ + ``error`` members, the ``method`` key deterministically makes the message a + call; smart-union scoring previously preferred whichever branch matched + more fields, which let a malformed frame classify as an error response. + """ + if isinstance(value, dict): + wire = cast("dict[str, Any]", value) + if "method" in wire: + request_id: Any = wire.get("id") + # Mirror `RequestId` (strict int | str): bool/float/None ids do not + # make a request; smart union classified those as notifications. + if isinstance(request_id, str) or (isinstance(request_id, int) and not isinstance(request_id, bool)): + return "request" + return "notification" + if "error" in wire: + return "error" + if "result" in wire: + return "response" + return None + # Revalidation / serialization of already-constructed models. + if isinstance(value, JSONRPCRequest): + return "request" + if isinstance(value, JSONRPCNotification): + return "notification" + if isinstance(value, JSONRPCError): + return "error" + if isinstance(value, JSONRPCResponse): + return "response" + return None + + +jsonrpc_message_adapter: TypeAdapter[JSONRPCMessage] = TypeAdapter( + Annotated[ + Annotated[JSONRPCRequest, Tag("request")] + | Annotated[JSONRPCNotification, Tag("notification")] + | Annotated[JSONRPCResponse, Tag("response")] + | Annotated[JSONRPCError, Tag("error")], + Discriminator( + _discriminate_jsonrpc_message, + custom_error_type="jsonrpc_message_invalid", + custom_error_message=( + "Not a valid JSON-RPC message: expected an object with a 'method', 'result', or 'error' member" + ), + ), + ] +) diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py index d316345c7..45161fba6 100644 --- a/src/mcp/server/streamable_http.py +++ b/src/mcp/server/streamable_http.py @@ -484,6 +484,9 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re # Parse the body - only read it once body = await request.body() + # Two-phase parse is intentional: `validate_json` would re-materialize the + # payload to run the adapter's callable discriminator, which measures ~1.75x + # slower than `from_json` + `validate_python` on large bodies. try: raw_message = pydantic_core.from_json(body) except ValueError as e: diff --git a/tests/types/test_jsonrpc_discriminator.py b/tests/types/test_jsonrpc_discriminator.py new file mode 100644 index 000000000..8f24f591f --- /dev/null +++ b/tests/types/test_jsonrpc_discriminator.py @@ -0,0 +1,148 @@ +"""Behavior pins for the key-presence discriminator on `jsonrpc_message_adapter`. + +The adapter validates exactly one union branch chosen by +`_discriminate_jsonrpc_message`; these tests pin classification parity with the +previous smart-union behavior, including the degenerate shapes. +""" + +import json +from typing import Any + +import pytest +from mcp_types import ( + ErrorData, + JSONRPCError, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + jsonrpc_message_adapter, +) +from pydantic import ValidationError + +REQUEST_WIRE: dict[str, Any] = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"} +NOTIFICATION_WIRE: dict[str, Any] = {"jsonrpc": "2.0", "method": "notifications/progress"} +RESPONSE_WIRE: dict[str, Any] = {"jsonrpc": "2.0", "id": 1, "result": {"tools": []}} +ERROR_WIRE: dict[str, Any] = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": "Parse error"}} + + +@pytest.mark.parametrize( + ("wire", "expected_type"), + [ + (REQUEST_WIRE, JSONRPCRequest), + (NOTIFICATION_WIRE, JSONRPCNotification), + (RESPONSE_WIRE, JSONRPCResponse), + (ERROR_WIRE, JSONRPCError), + ], +) +def test_validate_json_classifies_each_variant(wire: dict[str, object], expected_type: type) -> None: + message = jsonrpc_message_adapter.validate_json(json.dumps(wire), by_name=False) + assert type(message) is expected_type + + +@pytest.mark.parametrize("bad_id", [None, 1.5, True]) +def test_method_with_non_request_id_is_a_notification(bad_id: object) -> None: + """A `method` member with a null/float/bool id downgrades to a notification (smart-union parity).""" + wire = {"jsonrpc": "2.0", "id": bad_id, "method": "m"} + message = jsonrpc_message_adapter.validate_python(wire, by_name=False) + assert type(message) is JSONRPCNotification + + +def test_method_with_string_id_is_a_request() -> None: + message = jsonrpc_message_adapter.validate_python({"jsonrpc": "2.0", "id": "abc", "method": "m"}, by_name=False) + assert type(message) is JSONRPCRequest + + +def test_method_wins_over_result() -> None: + """A degenerate {method, id, result} object classifies as a request (smart-union parity).""" + wire: dict[str, Any] = {"jsonrpc": "2.0", "id": 1, "method": "m", "result": {}} + message = jsonrpc_message_adapter.validate_python(wire, by_name=False) + assert type(message) is JSONRPCRequest + + +def test_error_wins_over_result() -> None: + """A degenerate {id, result, error} object classifies as an error (smart-union parity).""" + wire = {"jsonrpc": "2.0", "id": 1, "result": {}, "error": {"code": -32603, "message": "boom"}} + message = jsonrpc_message_adapter.validate_python(wire, by_name=False) + assert type(message) is JSONRPCError + + +@pytest.mark.parametrize( + ("request_id", "expected_type"), + [ + (1, JSONRPCRequest), + (None, JSONRPCNotification), + ], +) +def test_method_wins_over_error(request_id: object, expected_type: type) -> None: + """A spec-invalid {method, error} hybrid classifies as a call, deliberately diverging from smart union. + + Smart-union scoring preferred `JSONRPCError` here (more matching fields), + which let a malformed frame masquerade as an error response to a pending + request. The `method` key now deterministically makes the message a call. + """ + wire: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, "method": "m", "error": {"code": -1, "message": "x"}} + message = jsonrpc_message_adapter.validate_python(wire, by_name=False) + assert type(message) is expected_type + + +def test_non_string_method_with_result_is_rejected() -> None: + """A `method` key selects the call arm even when its value is invalid, deliberately diverging from smart union. + + Smart union previously fell through to `JSONRPCResponse` for + {method: 123, result: {}}; the call arm now rejects the non-string method. + """ + wire: dict[str, Any] = {"jsonrpc": "2.0", "id": 1, "method": 123, "result": {}} + with pytest.raises(ValidationError) as exc_info: + jsonrpc_message_adapter.validate_python(wire, by_name=False) + assert exc_info.value.errors()[0]["loc"] == ("request", "method") + + +@pytest.mark.parametrize( + "unclassifiable", + [ + b'{"foo": 1}', + b"[]", + ], +) +def test_unclassifiable_json_raises_single_tagged_error(unclassifiable: bytes) -> None: + with pytest.raises(ValidationError) as exc_info: + jsonrpc_message_adapter.validate_json(unclassifiable, by_name=False) + errors = exc_info.value.errors() + assert len(errors) == 1 + assert errors[0]["type"] == "jsonrpc_message_invalid" + + +def test_unclassifiable_python_scalar_raises_tagged_error() -> None: + with pytest.raises(ValidationError) as exc_info: + jsonrpc_message_adapter.validate_python(42, by_name=False) + assert exc_info.value.errors()[0]["type"] == "jsonrpc_message_invalid" + + +def test_chosen_branch_failure_reports_single_branch_location() -> None: + """Once tagged, only the chosen branch validates; its errors carry the tag as the location root.""" + wire = {"jsonrpc": "2.0", "id": 1, "method": "m", "params": "notadict"} + with pytest.raises(ValidationError) as exc_info: + jsonrpc_message_adapter.validate_python(wire, by_name=False) + errors = exc_info.value.errors() + assert len(errors) == 1 + assert errors[0]["loc"] == ("request", "params") + + +@pytest.mark.parametrize( + "instance", + [ + JSONRPCRequest(jsonrpc="2.0", id=1, method="tools/list"), + JSONRPCNotification(jsonrpc="2.0", method="notifications/progress"), + JSONRPCResponse(jsonrpc="2.0", id=1, result={"tools": []}), + JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=-32700, message="Parse error")), + ], +) +def test_model_instances_revalidate_and_dump_identically(instance: JSONRPCMessage) -> None: + """The discriminator also tags already-constructed models (revalidation and dump paths).""" + revalidated = jsonrpc_message_adapter.validate_python(instance, by_name=False) + assert revalidated is instance + assert ( + jsonrpc_message_adapter.dump_json(instance, by_alias=True, exclude_none=True) + == instance.model_dump_json(by_alias=True, exclude_none=True).encode() + )