Skip to content

perf: JSONRPCMessage smart union scores all four branches on every inbound message #3136

Description

@sandole

Description

jsonrpc_message_adapter (src/mcp-types/mcp_types/jsonrpc.py) validates the
JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError union in
Pydantic's smart-union mode, which attempts/scores every branch for every message.
This adapter is the single decode choke point for every transport:

  • server/streamable_http.py (POST body), server/sse.py, server/stdio.py
  • client/streamable_http.py (SSE data + JSON responses), client/sse.py, client/stdio.py

The four variants are trivially distinguishable by key presence per JSON-RPC 2.0
(method [+ valid id] → request/notification, error → error, result → response),
so a callable Discriminator + Tags (available at the repo's pydantic floor,
>=2.12.0) selects exactly one branch instead.

Measured impact

scripts/bench_jsonrpc_codec.py (proposed alongside the fix), darwin arm64 / CPython 3.14.2 / pydantic 2.12.5, best-of-5:

payload path smart union discriminator speedup
request 109B validate_json 1.76µs 1.05µs 1.69x
request 109B from_json + validate_python 1.45µs 0.90µs 1.60x
notification 105B validate_json 1.88µs 0.86µs 2.18x
notification 105B from_json + validate_python 1.29µs 0.80µs 1.60x
result 22.3KB validate_json 51.7µs 32.1µs 1.61x
result 22.3KB from_json + validate_python 12.7µs 12.2µs 1.04x

For high-throughput agent communication (many small requests/notifications per
second per session) this halves envelope-decode CPU on the hot path of every
transport.

Behavior compatibility

The fix keeps classification parity with today's smart-union outcomes for all
spec-valid messages, verified empirically and pinned by tests:

  • {"id": null | 1.5 | true, "method": ...} and absent-id → JSONRPCNotification (current downgrade behavior preserved; note the interaction with Requests with "id": null silently misclassified as notifications #2057 / PR Reject JSON-RPC requests with null id instead of misclassifying as notifications #2075 - if that lands, the discriminator's id-guard becomes the single, clean place to implement the rejection)
  • degenerate {method, id, result}JSONRPCRequest; {id, result, error}JSONRPCError
  • JSONRPCMessage stays a plain union (so isinstance(x, JSONRPCMessage) keeps working); the Annotated/Tag wrapping lives only inside the TypeAdapter
  • wire output (dump_json, by_alias=True, exclude_none=True) is byte-identical
  • deliberate divergence on spec-invalid hybrids (surfaced by cubic's review on the PR):
    {method, error} hybrids classify as a call instead of JSONRPCError, and
    {method: <non-str>, result} is rejected instead of falling through to
    JSONRPCResponse - smart union picked those via field-count scoring, which let a
    malformed frame masquerade as an error response to a pending request; pinned by tests
  • only visible change: ValidationError for unclassifiable input becomes a single
    jsonrpc_message_invalid error with a clear message instead of a 4-branch error dump
    (no in-repo test asserts the old text; error codes and HTTP statuses unchanged)

Additional finding worth recording: with a callable discriminator,
validate_json must materialize the payload to call the discriminator, so the
server's existing two-phase parse (pydantic_core.from_jsonvalidate_python)
is the faster path on large bodies (12.2µs vs 32.1µs on 22KB) - the fix adds a
comment pinning that so it doesn't get "cleaned up" into a regression later.

Proposed fix

PR ready: key-presence callable discriminator on the adapter, zero call-site
changes, new parity/branch tests (100% coverage maintained), standalone
micro-benchmark script.


Disclosure: this issue and the accompanying PR were developed with AI assistance
(Claude Code); all measurements and parity checks were run and verified locally.

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions