Skip to content

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

Open
sandole wants to merge 2 commits into
modelcontextprotocol:mainfrom
sandole:perf/jsonrpc-discriminated-union
Open

perf: discriminate JSONRPCMessage union by key presence instead of smart-union scoring#3135
sandole wants to merge 2 commits into
modelcontextprotocol:mainfrom
sandole:perf/jsonrpc-discriminated-union

Conversation

@sandole

@sandole sandole commented Jul 19, 2026

Copy link
Copy Markdown

Closes #3136

Draft - awaiting maintainer feedback on #3136 per the issue-first
workflow; this PR stays in draft until maintainers
have weighed in on the "what/why".

Summary

jsonrpc_message_adapter - the decode choke point for every transport (stdio, SSE,
streamable HTTP; client and server) - validates the 4-member JSONRPCMessage union in
Pydantic smart-union mode, scoring all four branches on every inbound message:

JSONRPCMessage = JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError
jsonrpc_message_adapter: TypeAdapter[JSONRPCMessage] = TypeAdapter(JSONRPCMessage)

The variants are distinguishable by key presence per JSON-RPC 2.0 (method [+ valid
id] → request/notification, error → error, result → response), so this PR adds a
key-presence callable Discriminator and Tags inside the TypeAdapter, selecting
exactly one branch per message.

Measurements

uv run python scripts/bench_jsonrpc_codec.py (darwin arm64, CPython 3.14.2, pydantic 2.12.5; best of 10 rounds with adapter run order alternating each round, per cubic's review note, so warm-up/CPU-state effects are not attributed to either adapter):

payload              path            smart union  discriminator  speedup
request 109B         validate_json        1.68us         1.00us    1.68x
request 109B         two-phase            1.38us         0.87us    1.59x
notification 105B    validate_json        1.85us         0.83us    2.23x
notification 105B    two-phase            1.23us         0.78us    1.59x
result 22.3KB        validate_json       49.23us        30.77us    1.60x
result 22.3KB        two-phase           12.24us        11.78us    1.04x

The benchmark script is included so the numbers are reproducible and to pin the
counter-intuitive finding below; happy to drop it from the PR and keep the numbers
here only, at maintainer discretion.

Compatibility

  • JSONRPCMessage stays a plain union (isinstance keeps working); Tag/Discriminator live only inside the TypeAdapter. Zero call-site changes.
  • Classification parity with smart-union outcomes for all spec-valid messages, pinned by new tests: null/1.5/bool/absent-id method objects still downgrade to notifications (interaction with Requests with "id": null silently misclassified as notifications #2057/Reject JSON-RPC requests with null id instead of misclassifying as notifications #2075 noted - the id-guard is the natural locus for that change if it lands); degenerate {method,id,result} → Request and {id,result,error} → Error unchanged.
  • Deliberate divergence on spec-invalid hybrids (surfaced by cubic's review, pinned by tests): {method, error} hybrids now classify as a call (the method key wins deterministically) instead of JSONRPCError, and {method: <non-str>, result} is rejected instead of falling through to JSONRPCResponse. Smart union picked those branches via field-count scoring - a heuristic that let a malformed frame masquerade as the error response to a pending request. No fixed key-priority order can reproduce the scoring behavior on these shapes, and the key-presence rule is the spec-faithful choice for input that violates JSON-RPC 2.0's member exclusivity anyway.
  • Wire output byte-identical (dump_json with by_alias/exclude_none round-trip tested).
  • Only visible change: unclassifiable input raises a single jsonrpc_message_invalid error with an actionable message instead of a 4-branch error dump. Error codes and HTTP statuses unchanged.
  • Adds a comment pinning the server POST path's two-phase from_jsonvalidate_python parse: with a callable discriminator, single-pass validate_json must materialize the payload to call the discriminator and measures ~1.75x slower on large bodies, so the two-phase form is intentional (guards against a future "cleanup" regression).

Tests

  • New: tests/types/test_jsonrpc_discriminator.py - 21 tests, one per discriminator branch, including instance-revalidation identity, dump-parity pins, and pins for the documented spec-invalid-hybrid divergences.
  • Full suite: 5464 passed / 8 skipped / 1 xfailed, branch coverage 100%. (One pre-existing, environment-sensitive failure on the dev machine - test_modern_post_with_deeply_nested_body_is_parse_error_not_a_crash - fails identically on clean main; unrelated.)
  • ruff format --check, ruff check, repo-wide pyright (strict), and the mcp-types standalone import check: all clean.

AI disclosure

AI assistance (Claude Code) was used while preparing this PR. I reviewed the final
change, ran the measurements and full verification matrix above locally, and take
responsibility for it.

…art-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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 4 files

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread src/mcp-types/mcp_types/jsonrpc.py
Comment thread scripts/bench_jsonrpc_codec.py Outdated
…k 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: <non-str>, 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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant