Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 48 additions & 3 deletions src/mcp/client/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@
import logging
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import reduce
from functools import cache, reduce
from operator import or_
from types import TracebackType
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, overload
from types import TracebackType, UnionType
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, get_args, overload

import anyio
import anyio.abc
Expand All @@ -30,6 +30,7 @@
from mcp_types import methods as _methods
from mcp_types.version import (
HANDSHAKE_PROTOCOL_VERSIONS,
KNOWN_PROTOCOL_VERSIONS,
LATEST_HANDSHAKE_VERSION,
LATEST_MODERN_VERSION,
MODERN_PROTOCOL_VERSIONS,
Expand Down Expand Up @@ -77,6 +78,45 @@ def _clamp_inbound_ttl(raw: dict[str, Any]) -> None:
raw["ttlMs"] = 0


@cache
def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]:
"""Top-level wire keys `target` declares (its members', for a union).

A `RootModel` row (e.g. an empty result carried as `RootModel[Result]`)
reports its wrapped type's keys, not the pydantic-internal `root`.
"""
members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,)
models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)]
fields: set[str] = set()
for model in models:
if getattr(model, "__pydantic_root_model__", False): # a RootModel wrapper row
fields |= _wire_fields(model.model_fields["root"].annotation)
else:
fields.update(field.alias or name for name, field in model.model_fields.items())
return frozenset(fields)


@cache
def _later_revision_fields(method: str, version: str) -> frozenset[str]:
"""Result keys a revision newer than `version` declares for `method` but `version` doesn't.

The version-free result types carry every revision's fields, so such a key
(e.g. 2026-07-28 `ttlMs`/`cacheScope` on a pre-2026 session) is outside the
negotiated contract yet would still parse into the model and trip that later
revision's constraints. Empty at the newest known revision.
"""
current = _methods.SERVER_RESULTS.get((method, version))
if current is None or version not in KNOWN_PROTOCOL_VERSIONS:
return frozenset()
newer = KNOWN_PROTOCOL_VERSIONS[KNOWN_PROTOCOL_VERSIONS.index(version) + 1 :]
later: set[str] = set()
for revision in newer:
row = _methods.SERVER_RESULTS.get((method, revision))
if row is not None:
later |= _wire_fields(row)
return frozenset(later) - _wire_fields(current)


def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool:
"""JSON equality for two output schemas.

Expand Down Expand Up @@ -558,6 +598,11 @@ async def send_request(
_methods.validate_server_result(method, version, raw)
except KeyError:
pass
# Drop a later revision's fields (e.g. 2026-07-28 cache hints on a pre-2026
# session): they are outside the negotiated contract, and the version-free
# result type would otherwise apply that revision's constraints to them.
if not (foreign := _later_revision_fields(method, version)).isdisjoint(raw):
raw = {key: value for key, value in raw.items() if key not in foreign}
if isinstance(result_type, TypeAdapter):
return result_type.validate_python(raw, by_name=False)
return result_type.model_validate(raw, by_name=False)
Expand Down
69 changes: 69 additions & 0 deletions tests/client/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -1759,6 +1759,75 @@ async def test_a_boolean_inbound_ttl_is_not_clamped_only_coerced_by_validation(w
assert result.ttl_ms == int(wire_ttl)


_LEGACY_HINTED_RESULTS: list[tuple[str, dict[str, Any]]] = [
("list_tools", {"tools": []}),
("list_prompts", {"prompts": []}),
("list_resources", {"resources": []}),
("list_resource_templates", {"resourceTemplates": []}),
("read_resource", {"contents": []}),
]

_LEGACY_TAGGED_RESULTS: list[tuple[str, dict[str, Any]]] = [
("call_tool", {"content": []}),
("get_prompt", {"messages": []}),
]


def _legacy_init(version: str) -> dict[str, Any]:
return InitializeResult(
protocol_version=version,
capabilities=ServerCapabilities(),
server_info=Implementation(name="mock-server", version="0.1.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True)


async def _call_legacy(session: ClientSession, verb: str) -> Any:
if verb == "read_resource":
return await session.read_resource("mem://x")
if verb == "call_tool":
return await session.call_tool("t", {})
if verb == "get_prompt":
return await session.get_prompt("p")
return await getattr(session, verb)()


@pytest.mark.anyio
@pytest.mark.parametrize("version", HANDSHAKE_PROTOCOL_VERSIONS)
@pytest.mark.parametrize(("verb", "body"), _LEGACY_HINTED_RESULTS)
async def test_cache_hints_from_a_legacy_server_never_reach_the_result(
version: str, verb: str, body: dict[str, Any]
) -> None:
"""SDK-defined: on any pre-2026 session the caching fields are outside the negotiated
schema, so whatever a server puts in them - even values the 2026-07-28 enum
rejects - is dropped and the model shows its conservative defaults."""
dispatcher = _ScriptedDispatcher(_legacy_init(version), {**body, "ttlMs": -1, "cacheScope": "session"})
with anyio.fail_after(5):
async with ClientSession(dispatcher=dispatcher) as session:
await session.initialize()
result = await _call_legacy(session, verb)
assert (result.ttl_ms, result.cache_scope) == (0, "private")
assert not {"ttl_ms", "cache_scope"} & result.model_fields_set


@pytest.mark.anyio
@pytest.mark.parametrize(("verb", "body"), _LEGACY_TAGGED_RESULTS)
async def test_a_2026_result_type_tag_from_a_legacy_server_never_reaches_the_result(
verb: str, body: dict[str, Any]
) -> None:
"""SDK-defined: `resultType` is 2026-07-28 vocabulary that also feeds result-union
routing, so a tag on a pre-2026 wire (even one no union arm claims) is dropped and
the plain result is returned rather than mis-routing or failing."""
# `call_tool` re-lists tools to validate structured content; that answer trails harmlessly for `get_prompt`.
dispatcher = _ScriptedDispatcher(
_legacy_init(LATEST_HANDSHAKE_VERSION), {**body, "resultType": "task"}, {"tools": []}
)
with anyio.fail_after(5):
async with ClientSession(dispatcher=dispatcher) as session:
await session.initialize()
result = await _call_legacy(session, verb)
assert result.result_type == "complete"


@pytest.mark.anyio
async def test_session_call_tool_returns_input_required_result_when_opted_in() -> None:
"""`ClientSession.call_tool(..., allow_input_required=True)` surfaces the
Expand Down
Loading