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

Closed
SA
11d ago
+315-3

Closes #3136

Draft - awaiting maintainer feedback on #3136 per the issue-first
workflow

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 #2057/#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_json → validate_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.