Address review feedback on identity stripping and stale tool-map pruning

Author: maxisbey

src/mcp/client/client.py

@@ -9,7 +9,6 @@
 from contextlib import AsyncExitStack
 from dataclasses import KW_ONLY, dataclass, field
 from typing import Any, Literal, TypeVar, cast
-from urllib.parse import urlsplit
 
 import anyio
 import anyio.lowlevel
@@ -132,11 +131,19 @@ def _strip_userinfo(url: str) -> str:
 
     Credentials must not enter cache-key material; any further normalization could merge distinct servers.
     """
-    netloc = urlsplit(url).netloc  # raw authority bytes (urlsplit case-folds only `.scheme`), so slicing is exact
-    if "@" not in netloc:
+    # Pure text, no urlsplit: it strips embedded tab/CR/LF before parsing, which would misalign slices.
+    sep = url.find("//")
+    if sep == -1:
         return url
-    start = url.index("//") + 2
-    return url[:start] + netloc.rpartition("@")[2] + url[start + len(netloc) :]
+    start = sep + 2
+    end = len(url)
+    for delimiter in "/?#":
+        if (found := url.find(delimiter, start)) != -1:
+            end = min(end, found)
+    authority = url[start:end]
+    if "@" not in authority:
+        return url
+    return url[:start] + authority.rpartition("@")[2] + url[end:]
 
 
 def _evicting_message_handler(cache: ClientResponseCache, user_handler: MessageHandlerFnT | None) -> MessageHandlerFnT:
@@ -746,9 +753,12 @@ async def list_tools(
             meta=meta,
             cache_mode=cache_mode,
             send=lambda: self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta)),
-            # A cache hit skips session.list_tools, so the session re-absorbs the
-            # served listing to rebuild its derived per-tool state.
-            absorb=self.session._absorb_tool_listing,  # pyright: ignore[reportPrivateUsage]
+            # A cache hit skips session.list_tools, so the session re-absorbs the served
+            # listing to rebuild its derived per-tool state. Hits are cursorless, but a
+            # cached page 1 can carry next_cursor - never prune on a partial listing.
+            absorb=lambda hit: self.session._absorb_tool_listing(  # pyright: ignore[reportPrivateUsage]
+                hit, complete=hit.next_cursor is None
+            ),
         )
 
     @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)

src/mcp/client/session.py

@@ -906,12 +906,14 @@ async def list_tools(self, *, params: types.PaginatedRequestParams | None = None
             types.ListToolsRequest(params=params),
             types.ListToolsResult,
         )
-        return self._absorb_tool_listing(result)
+        complete = (params is None or params.cursor is None) and result.next_cursor is None
+        return self._absorb_tool_listing(result, complete=complete)
 
-    def _absorb_tool_listing(self, result: types.ListToolsResult) -> types.ListToolsResult:
+    def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) -> types.ListToolsResult:
         """Filter the listing per the 2026 x-mcp-header MUST and rebuild derived per-tool state, in place.
 
         Idempotent: cached values are already post-filter, so the response cache can re-absorb a served listing.
+        `complete` (an uncursored single-page listing) prunes per-tool state down to the listing's tools.
         """
         if self._negotiated_version in MODERN_PROTOCOL_VERSIONS:
             # 2026-07-28: clients MUST drop tools whose x-mcp-header annotations are invalid.
@@ -928,11 +930,17 @@ def _absorb_tool_listing(self, result: types.ListToolsResult) -> types.ListTools
                 kept.append(tool)
             result.tools = kept
 
-        # Cache tool output schemas for future validation
-        # Note: don't clear the cache, as we may be using a cursor
+        # Cache tool output schemas for future validation; cursor pages only ever add.
         for tool in result.tools:
             self._tool_output_schemas[tool.name] = tool.output_schema
 
+        if complete:
+            # The listing is the full tool universe, so state for unlisted tools is stale
+            # (the server dropped them, or a shared-cache writer's filter did).
+            names = {tool.name for tool in result.tools}
+            self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names}
+            self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names}
+
         return result
 
     @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)

tests/client/test_caching.py

@@ -432,6 +432,8 @@ async def clear(self) -> None:
         raise NotImplementedError
 
 
+# The lax pragmas here and in the wedged-store tests: 3.11's settrace-based coverage loses
+# tracing in frames resumed after the coordinator's bounded-shield cleanup cancellation.
 class _WedgingDeleteStore:
     """Once `wedged` flips, every `delete` blocks forever (an Event nothing sets),
     modelling a remote store with no socket timeout of its own."""
@@ -449,7 +451,7 @@ async def get(self, key: CacheKey) -> CacheEntry | None:
 
     async def set(self, key: CacheKey, entry: CacheEntry) -> None:
         await self.before_set_commits()
-        await self.inner.set(key, entry)
+        await self.inner.set(key, entry)  # pragma: lax no cover
 
     async def delete(self, key: CacheKey) -> None:
         self.deletes_started += 1
@@ -805,8 +807,10 @@ async def test_evict_key_with_a_wedged_store_delete_returns_at_the_cleanup_bound
     cache = _coordinator(store, store_cleanup_timeout=0.01)
     with caplog.at_level(logging.WARNING, logger="mcp.client.caching"), anyio.fail_after(5):
         await cache.evict_key("tools/list", "")
-    assert store.deletes_started == 1  # the second arm's delete was abandoned with the first
-    assert caplog.messages == snapshot(["Response cache store delete timed out; the entry will age out by TTL"])
+    assert store.deletes_started == 1  # pragma: lax no cover  # the second arm's delete was abandoned with the first
+    assert caplog.messages == snapshot(  # pragma: lax no cover
+        ["Response cache store delete timed out; the entry will age out by TTL"]
+    )
 
 
 async def test_a_refresh_purge_with_a_wedged_store_delete_returns_at_the_cleanup_bound() -> None:
@@ -815,7 +819,7 @@ async def test_a_refresh_purge_with_a_wedged_store_delete_returns_at_the_cleanup
     gen = cache.capture("tools/list", "")
     with anyio.fail_after(5):
         await cache.write("tools/list", "", _wire_result(ttl_ms=0), gen, "refresh")
-    assert store.deletes_started == 1
+    assert store.deletes_started == 1  # pragma: lax no cover
 
 
 async def test_an_eviction_mid_set_with_a_wedged_store_delete_returns_at_the_cleanup_bound() -> None:
@@ -833,9 +837,9 @@ async def wedge_then_evict() -> None:
     with anyio.fail_after(5):
         await cache.write("tools/list", "", _wire_result(ttl_ms=60_000), gen, "use")
     # Opposite-arm delete, the eviction's first delete, the compensating delete.
-    assert store.deletes_started == 3
+    assert store.deletes_started == 3  # pragma: lax no cover
     # The accepted degradation: the unreaped entry stays until its TTL expires.
-    assert await store.inner.get(CacheKey("tools/list", "", _private_arm())) is not None
+    assert await store.inner.get(CacheKey("tools/list", "", _private_arm())) is not None  # pragma: lax no cover
 
 
 # --- Coordinator: store error discipline ---

tests/client/test_client.py

@@ -506,6 +506,100 @@ async def on_list_tools(
         assert [t.name for t in result.tools] == ["ok", "dropme"]
 
 
+_RETIRED_TOOL = Tool(
+    name="retired",
+    input_schema={"type": "object", "properties": {"region": {"type": "string", "x-mcp-header": "Region"}}},
+    output_schema={"type": "object"},
+)
+_SURVIVOR_TOOL = Tool(name="survivor", input_schema={"type": "object"})
+
+
+def _scripted_listing_server(listings: list[ListToolsResult]) -> Server:
+    """Serves the given listings in order, one per tools/list request."""
+
+    async def on_list_tools(ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> ListToolsResult:
+        return listings.pop(0)
+
+    return Server("test", on_list_tools=on_list_tools)
+
+
+async def test_a_complete_listing_prunes_per_tool_state_for_tools_it_no_longer_contains() -> None:
+    """SDK-defined: a complete (uncursored, cursorless) listing is the full tool universe, so the
+    header map and output schema derived from an earlier listing of a now-absent tool are dropped."""
+    server = _scripted_listing_server(
+        [
+            ListToolsResult(tools=[_RETIRED_TOOL, _SURVIVOR_TOOL]),
+            ListToolsResult(tools=[_SURVIVOR_TOOL]),
+        ]
+    )
+
+    with anyio.fail_after(5):
+        async with Client(server) as client:
+            await client.session.list_tools()
+            assert set(client.session._x_mcp_header_maps) == {"retired", "survivor"}
+            assert set(client.session._tool_output_schemas) == {"retired", "survivor"}
+
+            await client.session.list_tools()
+            assert set(client.session._x_mcp_header_maps) == {"survivor"}
+            assert set(client.session._tool_output_schemas) == {"survivor"}
+
+
+async def test_a_complete_listing_prunes_output_schemas_on_a_legacy_session_too() -> None:
+    """SDK-defined: the prune is era-independent -- legacy sessions cache output schemas the same
+    way (their header-map dict just stays empty, since the x-mcp-header filter is 2026-only)."""
+    server = _scripted_listing_server(
+        [
+            ListToolsResult(tools=[_RETIRED_TOOL, _SURVIVOR_TOOL]),
+            ListToolsResult(tools=[_SURVIVOR_TOOL]),
+        ]
+    )
+
+    with anyio.fail_after(5):
+        async with Client(server, mode="legacy") as client:
+            await client.session.list_tools()
+            assert set(client.session._tool_output_schemas) == {"retired", "survivor"}
+            assert client.session._x_mcp_header_maps == {}
+
+            await client.session.list_tools()
+            assert set(client.session._tool_output_schemas) == {"survivor"}
+
+
+async def test_a_listing_with_a_next_cursor_prunes_no_per_tool_state() -> None:
+    """SDK-defined: a first page carrying next_cursor is not the full universe -- state for tools
+    expected on later pages must survive it."""
+    server = _scripted_listing_server(
+        [
+            ListToolsResult(tools=[_RETIRED_TOOL, _SURVIVOR_TOOL]),
+            ListToolsResult(tools=[_SURVIVOR_TOOL], next_cursor="2"),
+        ]
+    )
+
+    with anyio.fail_after(5):
+        async with Client(server) as client:
+            await client.session.list_tools()
+            await client.session.list_tools()
+            assert set(client.session._x_mcp_header_maps) == {"retired", "survivor"}
+            assert set(client.session._tool_output_schemas) == {"retired", "survivor"}
+
+
+async def test_a_cursor_page_fetch_prunes_no_per_tool_state() -> None:
+    """SDK-defined: a continuation page is partial even when it ends the pagination (no
+    next_cursor) -- only an uncursored single-page listing prunes."""
+    server = _scripted_listing_server(
+        [
+            ListToolsResult(tools=[_RETIRED_TOOL, _SURVIVOR_TOOL]),
+            ListToolsResult(tools=[_SURVIVOR_TOOL]),
+        ]
+    )
+
+    with anyio.fail_after(5):
+        async with Client(server) as client:
+            await client.session.list_tools()
+            await client.session.list_tools(params=types.PaginatedRequestParams(cursor="2"))
+            assert set(client.session._x_mcp_header_maps) == {"retired", "survivor"}
+            assert set(client.session._tool_output_schemas) == {"retired", "survivor"}
+
+
 def test_client_rejects_handshake_era_mode_at_construction() -> None:
     """A handshake-era protocol-version string passed as `mode=` is rejected by
     `__post_init__` with a hint to use `mode='legacy'` — the version-pin path is

tests/client/test_client_caching.py

@@ -139,16 +139,26 @@ def test_userinfo_variants_of_a_server_url_share_one_cache_identity() -> None:
         ("HTTPS://a@X.example/mcp", "HTTPS://X.example/mcp"),
         ("https://u@h/p?", "https://h/p?"),
         ("https://u@h/p#", "https://h/p#"),
+        ("https://u\tser:p@h.example/p", "https://h.example/p"),
+        ("https://u:p@h.example/pa\tth", "https://h.example/pa\tth"),
     ],
-    ids=["scheme-case", "empty-query", "empty-fragment"],
+    ids=["scheme-case", "empty-query", "empty-fragment", "tab-in-userinfo", "tab-in-path"],
 )
 def test_stripping_userinfo_changes_no_other_byte_of_the_url(with_userinfo: str, bare: str) -> None:
     """The removed `userinfo@` is the only byte difference: no scheme case-folding, no dropped
-    empty `?`/`#` delimiters. A userinfo-free URL passes through untouched, so arm equality
-    proves the stripped form is byte-identical to the bare URL."""
+    empty `?`/`#` delimiters, and control characters - which urlsplit would silently strip,
+    misaligning any parser-derived slice - stay byte-exact outside the removed span. A
+    userinfo-free URL passes through untouched, so arm equality proves the stripped form is
+    byte-identical to the bare URL."""
     assert _private_arm(Client(with_userinfo)) == _private_arm(Client(bare))
 
 
+def test_a_url_without_an_authority_passes_through_unchanged() -> None:
+    """No `//` means no authority span, so an `@` elsewhere strips nothing."""
+    arm_id = hashlib.sha256(b"mailto:a@b").hexdigest()
+    assert _private_arm(Client("mailto:a@b")) == json.dumps(["private", None, arm_id, ""])
+
+
 def test_the_server_url_is_sha256_hashed_before_it_enters_key_material() -> None:
     """Pins the docs' secrets-never-in-keys claim: a query-string secret never appears in store keys."""
     client = Client("https://user:pass@example.com/mcp?api_key=SECRET")
@@ -863,6 +873,33 @@ async def on_request(request: httpx.Request) -> None:
             assert posts[-1].headers["mcp-param-region"] == "us-west1"
 
 
+async def test_a_shared_store_hit_prunes_a_header_map_the_writers_filter_dropped() -> None:
+    """Cached listings are post-filter: when another client's refresh wrote a listing whose
+    filter dropped tool `x` (its annotation went invalid), a hit on that entry must prune the
+    reader's stale arg-to-header map, or it would keep emitting Mcp-Param-* headers for `x`."""
+    valid = {"type": "object", "properties": {"region": {"type": "string", "x-mcp-header": "Region"}}}
+    invalid = {"type": "object", "properties": {"region": {"type": "string", "x-mcp-header": "bad name"}}}
+    schema = valid
+
+    async def list_tools(ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> ListToolsResult:
+        return ListToolsResult(tools=[Tool(name="x", input_schema=schema)])
+
+    server = Server("filtering", on_list_tools=list_tools, cache_hints={"tools/list": CacheHint(ttl_ms=60_000)})
+    config = CacheConfig(store=InMemoryResponseCacheStore(), partition="p", target_id="svc", clock=_ManualClock())
+
+    with anyio.fail_after(5):
+        async with Client(server, cache=config) as reader, Client(server, cache=config) as writer:
+            await reader.list_tools()  # fetches while `x` is valid; the reader holds its header map
+            assert "x" in reader.session._x_mcp_header_maps
+
+            schema = invalid
+            await writer.list_tools(cache_mode="refresh")  # the writer's filter drops `x`; the entry is replaced
+
+            served = await reader.list_tools()  # hit on the writer's entry
+            assert served.tools == []
+            assert "x" not in reader.session._x_mcp_header_maps
+
+
 async def test_a_tools_list_changed_notification_makes_the_next_list_refetch() -> None:
     """Spec SHOULD: list_changed invalidates the cached listing. Legacy session +
     `default_ttl_ms` entry: eviction is era-independent."""

tests/interaction/transports/test_hosting_http_modern.py

@@ -511,7 +511,8 @@ async def test_modern_client_stops_mirroring_after_a_re_list_drops_the_tool() ->
     bad_schema = {"type": "object", "properties": {"a": {"type": "string", "x-mcp-header": "bad name"}}}
     valid = Tool(name="run", input_schema=schema)
     invalid = Tool(name="run", input_schema=bad_schema)
-    listings = iter([valid, invalid])
+    # Three pages: the call after the drop re-lists once because the prune also cleared `run`'s schema entry.
+    listings = iter([valid, invalid, invalid])
 
     async def list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
         return ListToolsResult(tools=[next(listings)], ttl_ms=0, cache_scope="public")