Skip to content

Commit 31b76cb

Browse files
maxisbeymaxisbey
authored
Drop later-revision cache-hint fields on pre-2026 sessions (modelcontextprotocol#3223)
1 parent 959569b commit 31b76cb

2 files changed

Lines changed: 117 additions & 3 deletions

File tree

‎src/mcp/client/session.py‎

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44
import logging
55
from collections.abc import Callable, Mapping, Sequence
66
from dataclasses import dataclass
7-
from functools import reduce
7+
from functools import cache, reduce
88
from operator import or_
9-
from types import TracebackType
10-
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, overload
9+
from types import TracebackType, UnionType
10+
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, get_args, overload
1111

1212
import anyio
1313
import anyio.abc
@@ -30,6 +30,7 @@
3030
from mcp_types import methods as _methods
3131
from mcp_types.version import (
3232
HANDSHAKE_PROTOCOL_VERSIONS,
33+
KNOWN_PROTOCOL_VERSIONS,
3334
LATEST_HANDSHAKE_VERSION,
3435
LATEST_MODERN_VERSION,
3536
MODERN_PROTOCOL_VERSIONS,
@@ -77,6 +78,45 @@ def _clamp_inbound_ttl(raw: dict[str, Any]) -> None:
7778
raw["ttlMs"] = 0
7879

7980

81+
@cache
82+
def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]:
83+
"""Top-level wire keys `target` declares (its members', for a union).
84+
85+
A `RootModel` row (e.g. an empty result carried as `RootModel[Result]`)
86+
reports its wrapped type's keys, not the pydantic-internal `root`.
87+
"""
88+
members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,)
89+
models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)]
90+
fields: set[str] = set()
91+
for model in models:
92+
if getattr(model, "__pydantic_root_model__", False): # a RootModel wrapper row
93+
fields |= _wire_fields(model.model_fields["root"].annotation)
94+
else:
95+
fields.update(field.alias or name for name, field in model.model_fields.items())
96+
return frozenset(fields)
97+
98+
99+
@cache
100+
def _later_revision_fields(method: str, version: str) -> frozenset[str]:
101+
"""Result keys a revision newer than `version` declares for `method` but `version` doesn't.
102+
103+
The version-free result types carry every revision's fields, so such a key
104+
(e.g. 2026-07-28 `ttlMs`/`cacheScope` on a pre-2026 session) is outside the
105+
negotiated contract yet would still parse into the model and trip that later
106+
revision's constraints. Empty at the newest known revision.
107+
"""
108+
current = _methods.SERVER_RESULTS.get((method, version))
109+
if current is None or version not in KNOWN_PROTOCOL_VERSIONS:
110+
return frozenset()
111+
newer = KNOWN_PROTOCOL_VERSIONS[KNOWN_PROTOCOL_VERSIONS.index(version) + 1 :]
112+
later: set[str] = set()
113+
for revision in newer:
114+
row = _methods.SERVER_RESULTS.get((method, revision))
115+
if row is not None:
116+
later |= _wire_fields(row)
117+
return frozenset(later) - _wire_fields(current)
118+
119+
80120
def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool:
81121
"""JSON equality for two output schemas.
82122
@@ -558,6 +598,11 @@ async def send_request(
558598
_methods.validate_server_result(method, version, raw)
559599
except KeyError:
560600
pass
601+
# Drop a later revision's fields (e.g. 2026-07-28 cache hints on a pre-2026
602+
# session): they are outside the negotiated contract, and the version-free
603+
# result type would otherwise apply that revision's constraints to them.
604+
if not (foreign := _later_revision_fields(method, version)).isdisjoint(raw):
605+
raw = {key: value for key, value in raw.items() if key not in foreign}
561606
if isinstance(result_type, TypeAdapter):
562607
return result_type.validate_python(raw, by_name=False)
563608
return result_type.model_validate(raw, by_name=False)

‎tests/client/test_session.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1759,6 +1759,75 @@ async def test_a_boolean_inbound_ttl_is_not_clamped_only_coerced_by_validation(w
17591759
assert result.ttl_ms == int(wire_ttl)
17601760

17611761

1762+
_LEGACY_HINTED_RESULTS: list[tuple[str, dict[str, Any]]] = [
1763+
("list_tools", {"tools": []}),
1764+
("list_prompts", {"prompts": []}),
1765+
("list_resources", {"resources": []}),
1766+
("list_resource_templates", {"resourceTemplates": []}),
1767+
("read_resource", {"contents": []}),
1768+
]
1769+
1770+
_LEGACY_TAGGED_RESULTS: list[tuple[str, dict[str, Any]]] = [
1771+
("call_tool", {"content": []}),
1772+
("get_prompt", {"messages": []}),
1773+
]
1774+
1775+
1776+
def _legacy_init(version: str) -> dict[str, Any]:
1777+
return InitializeResult(
1778+
protocol_version=version,
1779+
capabilities=ServerCapabilities(),
1780+
server_info=Implementation(name="mock-server", version="0.1.0"),
1781+
).model_dump(by_alias=True, mode="json", exclude_none=True)
1782+
1783+
1784+
async def _call_legacy(session: ClientSession, verb: str) -> Any:
1785+
if verb == "read_resource":
1786+
return await session.read_resource("mem://x")
1787+
if verb == "call_tool":
1788+
return await session.call_tool("t", {})
1789+
if verb == "get_prompt":
1790+
return await session.get_prompt("p")
1791+
return await getattr(session, verb)()
1792+
1793+
1794+
@pytest.mark.anyio
1795+
@pytest.mark.parametrize("version", HANDSHAKE_PROTOCOL_VERSIONS)
1796+
@pytest.mark.parametrize(("verb", "body"), _LEGACY_HINTED_RESULTS)
1797+
async def test_cache_hints_from_a_legacy_server_never_reach_the_result(
1798+
version: str, verb: str, body: dict[str, Any]
1799+
) -> None:
1800+
"""SDK-defined: on any pre-2026 session the caching fields are outside the negotiated
1801+
schema, so whatever a server puts in them - even values the 2026-07-28 enum
1802+
rejects - is dropped and the model shows its conservative defaults."""
1803+
dispatcher = _ScriptedDispatcher(_legacy_init(version), {**body, "ttlMs": -1, "cacheScope": "session"})
1804+
with anyio.fail_after(5):
1805+
async with ClientSession(dispatcher=dispatcher) as session:
1806+
await session.initialize()
1807+
result = await _call_legacy(session, verb)
1808+
assert (result.ttl_ms, result.cache_scope) == (0, "private")
1809+
assert not {"ttl_ms", "cache_scope"} & result.model_fields_set
1810+
1811+
1812+
@pytest.mark.anyio
1813+
@pytest.mark.parametrize(("verb", "body"), _LEGACY_TAGGED_RESULTS)
1814+
async def test_a_2026_result_type_tag_from_a_legacy_server_never_reaches_the_result(
1815+
verb: str, body: dict[str, Any]
1816+
) -> None:
1817+
"""SDK-defined: `resultType` is 2026-07-28 vocabulary that also feeds result-union
1818+
routing, so a tag on a pre-2026 wire (even one no union arm claims) is dropped and
1819+
the plain result is returned rather than mis-routing or failing."""
1820+
# `call_tool` re-lists tools to validate structured content; that answer trails harmlessly for `get_prompt`.
1821+
dispatcher = _ScriptedDispatcher(
1822+
_legacy_init(LATEST_HANDSHAKE_VERSION), {**body, "resultType": "task"}, {"tools": []}
1823+
)
1824+
with anyio.fail_after(5):
1825+
async with ClientSession(dispatcher=dispatcher) as session:
1826+
await session.initialize()
1827+
result = await _call_legacy(session, verb)
1828+
assert result.result_type == "complete"
1829+
1830+
17621831
@pytest.mark.anyio
17631832
async def test_session_call_tool_returns_input_required_result_when_opted_in() -> None:
17641833
"""`ClientSession.call_tool(..., allow_input_required=True)` surfaces the

0 commit comments

Comments
 (0)