test(client): cover unadvertised capability branches

Author: Epochex

src/mcp/client/session_group.py

@@ -11,7 +11,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, TypeAlias
+from typing import Any, TypeAlias, cast
 
 import anyio
 import httpx
@@ -332,7 +332,7 @@ async def _establish_session(
     async def _aggregate_components(self, server_info: types.Implementation, session: mcp.ClientSession) -> None:
         """Aggregates prompts, resources, and tools from a given session."""
 
-        capabilities = session.initialize_result.capabilities if session.initialize_result else None
+        capabilities = cast(types.InitializeResult, session.initialize_result).capabilities
 
         # Create a reverse index so we can find all prompts, resources, and
         # tools belonging to this session. Used for removing components from
@@ -347,7 +347,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session
         tool_to_session_temp: dict[str, mcp.ClientSession] = {}
 
         # Query the server for its prompts and aggregate to list.
-        if capabilities is None or capabilities.prompts is not None:
+        if capabilities.prompts is not None:
             try:
                 prompts = (await session.list_prompts()).prompts
                 for prompt in prompts:
@@ -358,7 +358,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session
                 logging.warning(f"Could not fetch prompts: {err}")
 
         # Query the server for its resources and aggregate to list.
-        if capabilities is None or capabilities.resources is not None:
+        if capabilities.resources is not None:
             try:
                 resources = (await session.list_resources()).resources
                 for resource in resources:
@@ -369,7 +369,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session
                 logging.warning(f"Could not fetch resources: {err}")
 
         # Query the server for its tools and aggregate to list.
-        if capabilities is None or capabilities.tools is not None:
+        if capabilities.tools is not None:
             try:
                 tools = (await session.list_tools()).tools
                 for tool in tools:

tests/client/test_session_group.py

@@ -166,6 +166,39 @@ async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequ
             client.session.list_resources = mock.AsyncMock(side_effect=AssertionError("list_resources() was called"))
 
             await group.connect_with_session(client.initialize_result.server_info, client.session)
+            await group.call_tool("ping")
+
+    assert not caplog.records
+
+
+@pytest.mark.anyio
+async def test_client_session_group_skips_unadvertised_tools_and_resources(
+    caplog: pytest.LogCaptureFixture,
+):
+    from mcp import Client
+    from mcp.server import Server, ServerRequestContext
+
+    async def handle_list_prompts(
+        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
+    ) -> types.ListPromptsResult:
+        return types.ListPromptsResult(prompts=[types.Prompt(name="hello", description="Hello", arguments=[])])
+
+    server = Server(
+        "prompts-only-server",
+        on_list_prompts=handle_list_prompts,
+    )
+
+    group = ClientSessionGroup()
+
+    with caplog.at_level(logging.WARNING):
+        async with Client(server) as client:
+            assert client.initialize_result.capabilities.tools is None
+            assert client.initialize_result.capabilities.resources is None
+
+            client.session.list_tools = mock.AsyncMock(side_effect=AssertionError("list_tools() was called"))
+            client.session.list_resources = mock.AsyncMock(side_effect=AssertionError("list_resources() was called"))
+
+            await group.connect_with_session(client.initialize_result.server_info, client.session)
 
     assert not caplog.records