Skip to content

Commit bbef3cf

Browse files
committed
fix(client): validate protocol_version_override and fix in-proc auto-mode connect
cubic's review found two real gaps in the protocol_version_override feature: - Client(server, mode="auto", protocol_version_override=...) against an in-process Server/MCPServer failed to connect. negotiate_auto's override path calls session.initialize() directly, but _connect_inproc picked the handshake-less DirectDispatcher for every non-"legacy" mode, so there was no JSON-RPC dispatcher for initialize() to run on. _build_session now collapses this combination to the legacy connector. - protocol_version_override accepted any string uncritically: a modern version silently misbehaved instead of being rejected, and setting it alongside a version-pinned mode was silently ignored rather than erroring. Both are now validated in __post_init__, mirroring the existing mode check. Also pinned the untested default (None) branch's exact initialize() call shape in test_session_group.py, per the same review.
1 parent ce45982 commit bbef3cf

3 files changed

Lines changed: 57 additions & 4 deletions

File tree

‎src/mcp/client/client.py‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -367,7 +367,12 @@ async def main():
367367

368368
_entered: bool = field(init=False, default=False)
369369
protocol_version_override: str | None = None
370-
"""The protocol version to request during initialization. Defaults to the latest version."""
370+
"""Pin the legacy `initialize` handshake to a specific handshake-era version.
371+
372+
Only meaningful with `mode='legacy'` or `mode='auto'` (where it skips `server/discover`
373+
and negotiates directly); raises at construction with any other `mode`, since a version
374+
pin already fixes the negotiated version. Must be a member of `HANDSHAKE_PROTOCOL_VERSIONS`.
375+
`None` (the default) negotiates the latest version each `mode` would otherwise pick."""
371376
_session: ClientSession | None = field(init=False, default=None)
372377
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
373378
_connect: _Connector = field(init=False, repr=False, compare=False)
@@ -385,6 +390,23 @@ def __post_init__(self) -> None:
385390
f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}"
386391
)
387392

393+
if self.protocol_version_override is not None:
394+
if self.protocol_version_override not in HANDSHAKE_PROTOCOL_VERSIONS:
395+
hint = (
396+
f" ({self.protocol_version_override!r} is a modern version; mode='auto' already negotiates it)"
397+
if self.protocol_version_override in MODERN_PROTOCOL_VERSIONS
398+
else ""
399+
)
400+
raise ValueError(
401+
"protocol_version_override must be one of "
402+
f"{list(HANDSHAKE_PROTOCOL_VERSIONS)}; got {self.protocol_version_override!r}{hint}"
403+
)
404+
if self.mode not in ("legacy", "auto"):
405+
raise ValueError(
406+
f"protocol_version_override has no effect with mode={self.mode!r} "
407+
"(a version pin already fixes the negotiated version); use mode='legacy' or mode='auto'"
408+
)
409+
388410
self._folded_extensions = _fold_extensions(self.extensions)
389411

390412
srv = self.server
@@ -426,7 +448,12 @@ def __post_init__(self) -> None:
426448

427449
async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession:
428450
"""Enter the resolved connector and return an un-entered ClientSession."""
429-
dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions)
451+
# An override on mode='auto' skips discovery and drives `initialize()` directly
452+
# (see `negotiate_auto`), so the in-proc connector must hand back the legacy,
453+
# stream-backed dispatcher for this combination too, not the handshake-less
454+
# DirectDispatcher it otherwise picks for every non-'legacy' mode.
455+
connect_mode = "legacy" if self.mode == "auto" and self.protocol_version_override is not None else self.mode
456+
dispatcher = await self._connect(exit_stack, connect_mode, self.raise_exceptions)
430457
message_handler = self.message_handler
431458
if self._response_cache is not None:
432459
message_handler = _evicting_message_handler(self._response_cache, self.message_handler)

‎tests/client/test_client.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
Tool,
3434
ToolsCapability,
3535
)
36-
from mcp_types.version import LATEST_HANDSHAKE_VERSION
36+
from mcp_types.version import LATEST_HANDSHAKE_VERSION, LATEST_MODERN_VERSION
3737
from pydantic import FileUrl
3838

3939
from mcp import MCPDeprecationWarning, MCPError, StdioServerParameters
@@ -138,6 +138,32 @@ async def test_client_custom_protocol_version(app: MCPServer):
138138
assert client.server_info.name == "test"
139139

140140

141+
async def test_client_auto_mode_with_override_against_in_process_server(app: MCPServer):
142+
"""Regression: `mode='auto'` with `protocol_version_override` against an in-process
143+
`Server`/`MCPServer` used to always get the handshake-less `DirectDispatcher` (every
144+
non-'legacy' mode picked it), so `negotiate_auto`'s direct `initialize()` call for the
145+
override case had no JSON-RPC dispatcher to run on and the connect failed.
146+
"""
147+
async with Client(app, mode="auto", protocol_version_override="2024-11-05") as client:
148+
assert client.protocol_version == "2024-11-05"
149+
assert client.server_info is not None
150+
assert client.server_info.name == "test"
151+
152+
153+
def test_client_rejects_modern_protocol_version_override(app: MCPServer):
154+
"""`protocol_version_override` only pins the legacy handshake; a modern version string
155+
is a construction-time error rather than a confusing failure once connected."""
156+
with pytest.raises(ValueError, match="protocol_version_override must be one of"):
157+
Client(app, mode="auto", protocol_version_override=LATEST_MODERN_VERSION)
158+
159+
160+
def test_client_rejects_protocol_version_override_with_a_version_pin_mode(app: MCPServer):
161+
"""`protocol_version_override` has no effect once `mode` already pins a version, so it's
162+
rejected at construction instead of being silently ignored."""
163+
with pytest.raises(ValueError, match="protocol_version_override has no effect with mode="):
164+
Client(app, mode=LATEST_MODERN_VERSION, protocol_version_override="2024-11-05")
165+
166+
141167
async def test_client_with_simple_server(simple_server: Server):
142168
"""Test that from_server works with a basic Server instance."""
143169
async with Client(simple_server) as client:

‎tests/client/test_session_group.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -397,7 +397,7 @@ async def test_client_session_group_establish_session_parameterized(
397397
client_info=None,
398398
)
399399
mock_raw_session_cm.__aenter__.assert_awaited_once()
400-
mock_entered_session.initialize.assert_awaited_once()
400+
mock_entered_session.initialize.assert_awaited_once_with()
401401

402402
# 3. Assert returned values
403403
assert returned_server_info is mock_initialize_result.server_info

0 commit comments

Comments
 (0)