From 2e506e70be56f6d472c56a52061d6c5534e70b5d Mon Sep 17 00:00:00 2001 From: johnseong Date: Sun, 19 Jul 2026 18:25:36 -0400 Subject: [PATCH 1/2] perf: discriminate JSONRPCMessage union by key presence instead of smart-union scoring Every transport (stdio, SSE, streamable HTTP; client and server) decodes inbound messages through jsonrpc_message_adapter, which validated the 4-member JSONRPCMessage union in Pydantic smart-union mode - scoring all four branches per message. The variants are distinguishable by key presence per JSON-RPC 2.0, so a callable Discriminator selects exactly one branch: 1.6-2.2x faster envelope decoding (see scripts/bench_jsonrpc_codec.py), byte-identical wire output, and exact classification parity with the previous behavior (including the null/float/bool-id notification downgrade and error-over-result precedence for degenerate shapes), pinned by tests. JSONRPCMessage itself stays a plain union so isinstance() keeps working; the Tag/Discriminator wrapping lives only inside the TypeAdapter, so no call sites change. Unclassifiable input now raises a single jsonrpc_message_invalid error with an actionable message instead of a four-branch error dump; error codes and HTTP statuses are unchanged. Also documents why the server POST path keeps its two-phase from_json + validate_python parse: with a callable discriminator, validate_json must materialize the payload to call the discriminator and measures ~1.75x slower on large bodies. --- scripts/bench_jsonrpc_codec.py | 97 ++++++++++++++++++ src/mcp-types/mcp_types/jsonrpc.py | 56 ++++++++++- src/mcp/server/streamable_http.py | 3 + tests/types/test_jsonrpc_discriminator.py | 117 ++++++++++++++++++++++ 4 files changed, 270 insertions(+), 3 deletions(-) create mode 100644 scripts/bench_jsonrpc_codec.py create mode 100644 tests/types/test_jsonrpc_discriminator.py diff --git a/scripts/bench_jsonrpc_codec.py b/scripts/bench_jsonrpc_codec.py new file mode 100644 index 000000000..10df85747 --- /dev/null +++ b/scripts/bench_jsonrpc_codec.py @@ -0,0 +1,97 @@ +"""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(fn: Callable[[], Any], iterations: int) -> float: + """Return best-of-REPEATS per-call time in microseconds.""" + best = float("inf") + for _ in range(REPEATS): + start = time.perf_counter() + for _ in range(iterations): + fn() + best = min(best, time.perf_counter() - start) + return best / 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 = bench(lambda: decode(bare_union_adapter, body), iterations) + new = bench(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..d2fe12c36 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,54 @@ 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. Mirrors the smart-union outcomes + exactly, including the rule that a ``method`` member with a missing, null, + or non-int/str ``id`` classifies as a notification, and that ``error`` wins + over ``result`` when both are present. + """ + 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..267b82517 --- /dev/null +++ b/tests/types/test_jsonrpc_discriminator.py @@ -0,0 +1,117 @@ +"""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( + "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() + ) From 4ad8fc066e56f756b80878e1bd8f607a4fcc7a94 Mon Sep 17 00:00:00 2001 From: johnseong Date: Sun, 19 Jul 2026 19:04:28 -0400 Subject: [PATCH 2/2] Address review: pin spec-invalid hybrid divergences, de-bias benchmark ordering The discriminator does not reproduce smart-union scoring on spec-invalid hybrids: {method, error} now classifies as a call instead of JSONRPCError, and {method: , result} is rejected instead of falling through to JSONRPCResponse. No fixed key-priority order can mirror field-count scoring on these shapes, and the key-presence rule is the safer choice (a malformed frame can no longer masquerade as the error response to a pending request), so pin the divergence with tests and correct the docstring's parity claim to spec-valid messages. Also alternate adapter run order each benchmark round so warm-up and CPU-state effects are not attributed to either adapter; speedups are unchanged (1.6-2.2x). --- scripts/bench_jsonrpc_codec.py | 32 +++++++++++++++-------- src/mcp-types/mcp_types/jsonrpc.py | 12 ++++++--- tests/types/test_jsonrpc_discriminator.py | 31 ++++++++++++++++++++++ 3 files changed, 60 insertions(+), 15 deletions(-) diff --git a/scripts/bench_jsonrpc_codec.py b/scripts/bench_jsonrpc_codec.py index 10df85747..f461e3e1d 100644 --- a/scripts/bench_jsonrpc_codec.py +++ b/scripts/bench_jsonrpc_codec.py @@ -51,15 +51,22 @@ REPEATS = 5 -def bench(fn: Callable[[], Any], iterations: int) -> float: - """Return best-of-REPEATS per-call time in microseconds.""" - best = float("inf") - for _ in range(REPEATS): - start = time.perf_counter() - for _ in range(iterations): - fn() - best = min(best, time.perf_counter() - start) - return best / iterations * 1e6 +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: @@ -88,8 +95,11 @@ def main() -> None: 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 = bench(lambda: decode(bare_union_adapter, body), iterations) - new = bench(lambda: decode(jsonrpc_message_adapter, body), iterations) + 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") diff --git a/src/mcp-types/mcp_types/jsonrpc.py b/src/mcp-types/mcp_types/jsonrpc.py index d2fe12c36..77765bb54 100644 --- a/src/mcp-types/mcp_types/jsonrpc.py +++ b/src/mcp-types/mcp_types/jsonrpc.py @@ -127,10 +127,14 @@ 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. Mirrors the smart-union outcomes - exactly, including the rule that a ``method`` member with a missing, null, - or non-int/str ``id`` classifies as a notification, and that ``error`` wins - over ``result`` when both are present. + 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) diff --git a/tests/types/test_jsonrpc_discriminator.py b/tests/types/test_jsonrpc_discriminator.py index 267b82517..8f24f591f 100644 --- a/tests/types/test_jsonrpc_discriminator.py +++ b/tests/types/test_jsonrpc_discriminator.py @@ -67,6 +67,37 @@ def test_error_wins_over_result() -> None: 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", [