perf: discriminate JSONRPCMessage union by key presence instead of smart-union scoring

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()

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"
+            ),
+        ),
+    ]
+)

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:

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()
+    )