Skip to content

Commit 15ff3c2

Browse files
committed
Fix core-arm sentinel collision and close binding queues on every exit
A claim tagged 'core' could collide with the adapter's internal routing sentinel and hijack ordinary tools/call parsing; the sentinel is now derived to never equal a claimed tag. Binding queues close in finally blocks so a raising task-group exit cannot leak them. The mapping-form migration error now also covers an empty dict; claim dedup keys by resultType alone (the verb is single-valued at construction, and the activation map already keys by tag); the one-time private-spelling tree-grep test is gone; docstring backticks and a stray __future__ import brought in line with repo conventions.
1 parent 616e81b commit 15ff3c2

6 files changed

Lines changed: 65 additions & 49 deletions

File tree

src/mcp/client/client.py

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -207,22 +207,22 @@ def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExt
207207
Mirrors the server's consumption-time posture (`MCPServer._apply_extension`): a
208208
per-instance identifier is validated here because no class attribute existed to
209209
validate at definition time. `settings()` is read exactly once per extension and
210-
the returned dict is held by reference. Duplicate `(method, resultType)` claims
211-
and duplicate notification methods are rejected here, where both owning extensions
210+
the returned dict is held by reference. Duplicate resultType claims and
211+
duplicate notification methods are rejected here, where both owning extensions
212212
can be named — the session's own duplicate checks know only methods and tags.
213213
"""
214-
if not extensions:
215-
return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={})
216214
if isinstance(extensions, Mapping):
217215
raise TypeError(
218216
"extensions= takes a sequence of ClientExtension instances; the mapping form was "
219217
"replaced — use advertise(identifier, settings) for advertise-only entries"
220218
)
219+
if not extensions:
220+
return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={})
221221
ad: dict[str, dict[str, Any]] = {}
222222
claims: dict[str, tuple[ResultClaim[Any], ...]] = {}
223223
bindings: list[NotificationBinding[Any]] = []
224224
by_model: dict[type[Result], ResultClaim[Any]] = {}
225-
claim_owners: dict[tuple[str, str], str] = {}
225+
claim_owners: dict[str, str] = {}
226226
binding_owners: dict[str, str] = {}
227227
for extension in extensions:
228228
identifier = getattr(extension, "identifier", None)
@@ -237,20 +237,18 @@ def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExt
237237
ad[identifier] = extension.settings()
238238
extension_claims = tuple(extension.claims())
239239
for claim in extension_claims:
240-
key = (claim.method, claim.result_type)
241-
if key in claim_owners:
242-
owner = claim_owners[key]
240+
tag = claim.result_type
241+
if tag in claim_owners:
242+
owner = claim_owners[tag]
243243
both = (
244244
f"extension {identifier!r} claims"
245245
if owner == identifier
246246
else (f"extensions {owner!r} and {identifier!r} both claim")
247247
)
248-
raise ValueError(
249-
f"{both} {claim.method!r} resultType {claim.result_type!r}; a wire tag can have only one resolver"
250-
)
251-
claim_owners[key] = identifier
252-
# Collision-free by construction: a model's `result_type` Literal pins it to
253-
# exactly one tag, and each (method, tag) pair has exactly one owner.
248+
raise ValueError(f"{both} resultType {tag!r}; a wire tag can have only one resolver")
249+
claim_owners[tag] = identifier
250+
# One model, one tag: the model's result_type Literal is pinned to exactly
251+
# this tag at claim construction, so the type-keyed index cannot collide.
254252
by_model[claim.model] = claim
255253
if extension_claims:
256254
claims[identifier] = extension_claims

src/mcp/client/session.py

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,9 @@ def _build_call_tool_adapter(
241241
if not active:
242242
return _CallToolResultAdapter
243243
tags = frozenset(active)
244+
core_arm = "core"
245+
while core_arm in tags: # the routing sentinel must never collide with a claimed tag
246+
core_arm += "-"
244247

245248
def _route(value: Any) -> str:
246249
# pydantic hands the discriminator either the raw inbound dict or an
@@ -251,9 +254,9 @@ def _route(value: Any) -> str:
251254
tag = cast("dict[str, Any]", value).get("resultType")
252255
else:
253256
tag = getattr(value, "result_type", None)
254-
return tag if isinstance(tag, str) and tag in tags else "core"
257+
return tag if isinstance(tag, str) and tag in tags else core_arm
255258

256-
arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag("core")]]
259+
arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag(core_arm)]]
257260
arms += [Annotated[claim.model, Tag(tag)] for tag, claim in active.items()]
258261
# reduce(or_, ...) builds the Union dynamically; PEP-646 star-unpack needs py3.11+.
259262
return TypeAdapter(Annotated[reduce(or_, arms), Discriminator(_route)])
@@ -270,7 +273,7 @@ def _index_claims(
270273
advertised extension and no wire tag may be claimed twice.
271274
"""
272275
indexed: dict[str, tuple[ResultClaim[Any], ...]] = {}
273-
seen: set[tuple[str, str]] = set()
276+
seen: set[str] = set()
274277
for identifier, claims in (result_claims or {}).items():
275278
if extensions is None or identifier not in extensions:
276279
raise ValueError(
@@ -283,10 +286,9 @@ def _index_claims(
283286
"extension from the capability ad at every version — omit the key instead"
284287
)
285288
for claim in claims:
286-
key = (claim.method, claim.result_type)
287-
if key in seen:
288-
raise ValueError(f"duplicate result claim for {claim.method!r} resultType {claim.result_type!r}")
289-
seen.add(key)
289+
if claim.result_type in seen:
290+
raise ValueError(f"duplicate result claim for resultType {claim.result_type!r}")
291+
seen.add(claim.result_type)
290292
indexed[identifier] = tuple(claims)
291293
return indexed
292294

@@ -411,8 +413,10 @@ async def __aenter__(self) -> Self:
411413
# Shield the group's own scope (a new one would break LIFO exit)
412414
# so a pending outer cancellation cannot re-fire inside __aexit__.
413415
task_group.cancel_scope.shield = True
414-
await task_group.__aexit__(None, None, None)
415-
self._close_binding_queues()
416+
try:
417+
await task_group.__aexit__(None, None, None)
418+
finally:
419+
self._close_binding_queues()
416420
raise
417421
return self
418422

@@ -425,8 +429,10 @@ async def __aexit__(
425429
# Exit must not block: cancel the dispatcher, binding consumers, and in-flight callbacks.
426430
assert self._task_group is not None
427431
self._task_group.cancel_scope.cancel()
428-
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
429-
self._close_binding_queues()
432+
try:
433+
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
434+
finally:
435+
self._close_binding_queues()
430436
await resync_tracer()
431437
return result
432438

@@ -975,15 +981,15 @@ async def call_tool(
975981
allow_input_required: When ``False`` (default), an `InputRequiredResult`
976982
from the server raises `RuntimeError`; when ``True``, it is returned
977983
so the caller can resolve the requests and retry.
978-
allow_claimed: When ``False`` (default), a claimed (extension) result
979-
shape raises `UnexpectedClaimedResult`; when ``True``, the parsed
984+
allow_claimed: When `False` (default), a claimed (extension) result
985+
shape raises `UnexpectedClaimedResult`; when `True`, the parsed
980986
claim model is returned for the caller to handle.
981987
982988
Raises:
983989
RuntimeError: If the server returns an `InputRequiredResult` and
984990
``allow_input_required`` is ``False``.
985991
UnexpectedClaimedResult: If a claimed result shape parses and
986-
``allow_claimed`` is ``False``; carries the parsed value.
992+
`allow_claimed` is `False`; carries the parsed value.
987993
"""
988994
result = await self.send_request(
989995
types.CallToolRequest(

tests/client/test_client_extensions.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ def test_one_extension_claiming_a_tag_twice_reads_as_one_owner() -> None:
185185
Client(_add_server(), extensions=[_SelfConflictingClaims()])
186186

187187
assert str(exc_info.value) == snapshot(
188-
"extension 'com.example/twice' claims 'tools/call' resultType 'twice'; a wire tag can have only one resolver"
188+
"extension 'com.example/twice' claims resultType 'twice'; a wire tag can have only one resolver"
189189
)
190190

191191

@@ -252,8 +252,8 @@ def test_conflicting_claims_across_extensions_name_both_owners() -> None:
252252
Client(_add_server(), extensions=[_VoucherExtension(_unreachable_resolve), _RivalVoucherExtension()])
253253

254254
assert str(exc_info.value) == snapshot(
255-
"extensions 'com.example/voucher' and 'com.example/rival' both claim 'tools/call' "
256-
"resultType 'voucher'; a wire tag can have only one resolver"
255+
"extensions 'com.example/voucher' and 'com.example/rival' both claim resultType "
256+
"'voucher'; a wire tag can have only one resolver"
257257
)
258258

259259

tests/client/test_send_request_mcp_name.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
`send_request` typing: a `Request[...]` subclass passes without a cast.
88
"""
99

10-
from __future__ import annotations
1110

1211
from collections.abc import Mapping
1312
from typing import Any, Literal

tests/client/test_session_claims.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ def test_duplicate_claim_tag_across_extensions_rejected() -> None:
139139
result_claims={_TASKS_EXT: [_task_claim()], _AD_ONLY_EXT: [_task_claim()]},
140140
)
141141

142-
assert str(exc_info.value) == snapshot("duplicate result claim for 'tools/call' resultType 'task'")
142+
assert str(exc_info.value) == snapshot("duplicate result claim for resultType 'task'")
143143

144144

145145
def test_claims_keyed_to_unadvertised_extension_rejected() -> None:
@@ -351,6 +351,35 @@ async def test_discover_probe_ad_drops_claim_identifiers_at_a_legacy_probe_versi
351351
assert "extensions" not in capabilities
352352

353353

354+
class _CoreTaggedResult(Result):
355+
"""A claim whose wire tag collides with the adapter's internal routing sentinel."""
356+
357+
result_type: Literal["core"] = "core"
358+
payload: str = ""
359+
360+
361+
async def _resolve_core_tagged(result: _CoreTaggedResult, ctx: ClaimContext) -> CallToolResult:
362+
raise NotImplementedError
363+
364+
365+
@pytest.mark.anyio
366+
async def test_claim_tagged_core_cannot_hijack_core_parsing() -> None:
367+
"""SDK-defined: "core" is not protocol vocabulary, so a claim may use it as a wire
368+
tag — and the adapter's internal routing sentinel must not collide: ordinary tool
369+
results still parse as core results, and a claimed `core` raw routes to the model."""
370+
claim = ResultClaim(result_type="core", model=_CoreTaggedResult, resolve=_resolve_core_tagged)
371+
dispatcher = _RecordingDispatcher(tool_result={"resultType": "core", "payload": "p-1"})
372+
session = ClientSession(dispatcher=dispatcher, extensions={_TASKS_EXT: {}}, result_claims={_TASKS_EXT: [claim]})
373+
with anyio.fail_after(5):
374+
async with session:
375+
_adopt_modern(session)
376+
ordinary = session._call_tool_adapter.validate_python(_COMPLETE_TOOL_RESULT)
377+
claimed = await session.call_tool("t", {}, allow_claimed=True)
378+
379+
assert isinstance(ordinary, CallToolResult)
380+
assert isinstance(claimed, _CoreTaggedResult)
381+
382+
354383
# ── Routing through the adopt-built adapter ─────────────────────────────────
355384

356385

tests/client/test_session_promotions.py

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
"""`dispatch_input_request` and `validate_tool_result` are public `ClientSession` API."""
22

3-
import re
4-
from pathlib import Path
53

64
import mcp_types as types
75
import pytest
@@ -68,17 +66,3 @@ async def test_validate_tool_result_raises_on_schema_mismatch() -> None:
6866
# Stable SDK prefix only: the message tail is jsonschema text that shifts with the dependency.
6967
with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
7068
await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": "no"}))
71-
72-
73-
def _spell_private(name: str) -> str:
74-
return f"_{name}"
75-
76-
77-
def test_no_private_spelling_references_remain() -> None:
78-
"""The promotions are renames, not aliases — the old private names are gone from `src/`."""
79-
pattern = re.compile(f"{_spell_private('dispatch_input_request')}|{_spell_private('validate_tool_result')}")
80-
src = Path(__file__).resolve().parents[2] / "src"
81-
offenders = [
82-
(path.name, match) for path in sorted(src.rglob("*.py")) for match in pattern.findall(path.read_text())
83-
]
84-
assert not offenders

0 commit comments

Comments
 (0)