Skip to content
Open
Show file tree
Hide file tree
Changes from 2 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
45 changes: 42 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,39 @@ 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)."""
models: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,)
return frozenset(
field.alias or name
for model in models
if isinstance(model, type) and issubclass(model, BaseModel)
for name, field in model.model_fields.items()
)


@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 +592,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
31 changes: 31 additions & 0 deletions tests/client/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -1759,6 +1759,37 @@ 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": []}),
]


@pytest.mark.anyio
@pytest.mark.parametrize(("verb", "body"), _LEGACY_HINTED_RESULTS)
async def test_cache_hints_from_a_legacy_server_never_reach_the_result(verb: str, body: dict[str, Any]) -> None:
"""SDK-defined: on a 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."""
init_result = InitializeResult(
protocol_version=LATEST_HANDSHAKE_VERSION,
capabilities=ServerCapabilities(),
server_info=Implementation(name="mock-server", version="0.1.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True)
hinted = {**body, "ttlMs": -1, "cacheScope": "session"}
dispatcher = _ScriptedDispatcher(init_result, hinted)
with anyio.fail_after(5):
async with ClientSession(dispatcher=dispatcher) as session:
await session.initialize()
call = getattr(session, verb)
result = await (call("mem://x") if verb == "read_resource" else call())
assert (result.ttl_ms, result.cache_scope) == (0, "private")
assert not {"ttl_ms", "cache_scope"} & result.model_fields_set


@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