diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..3d55cc8ae7 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -4,6 +4,8 @@ A `Client` connects to one server. Real applications often want several (a searc **`ClientSessionGroup`** is one object that holds many connections and merges everything they expose into a single view. +The group only lists tools, resources, and prompts that each server advertises in its negotiated capabilities. + ## Two servers Start with two ordinary servers. They have nothing to do with each other, so both naturally called their tool `search`: diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index 13d8928ed4..760e7d7529 100644 --- a/src/mcp/client/session_group.py +++ b/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, Literal, Self, TypeAlias, overload +from typing import Any, Literal, Self, TypeAlias, cast, overload import anyio import httpx2 @@ -374,6 +374,8 @@ 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 = cast(types.ServerCapabilities, session.server_capabilities) + # Create a reverse index so we can find all prompts, resources, and # tools belonging to this session. Used for removing components from # the session group via self.disconnect_from_server. @@ -387,35 +389,38 @@ 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. - try: - prompts = (await session.list_prompts()).prompts - for prompt in prompts: - name = self._component_name(prompt.name, server_info) - prompts_temp[name] = prompt - component_names.prompts.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch prompts: {err}") + if capabilities.prompts is not None: + try: + prompts = (await session.list_prompts()).prompts + for prompt in prompts: + name = self._component_name(prompt.name, server_info) + prompts_temp[name] = prompt + component_names.prompts.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch prompts: {err}") # Query the server for its resources and aggregate to list. - try: - resources = (await session.list_resources()).resources - for resource in resources: - name = self._component_name(resource.name, server_info) - resources_temp[name] = resource - component_names.resources.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch resources: {err}") + if capabilities.resources is not None: + try: + resources = (await session.list_resources()).resources + for resource in resources: + name = self._component_name(resource.name, server_info) + resources_temp[name] = resource + component_names.resources.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch resources: {err}") # Query the server for its tools and aggregate to list. - try: - tools = (await session.list_tools()).tools - for tool in tools: - name = self._component_name(tool.name, server_info) - tools_temp[name] = tool - tool_to_session_temp[name] = session - component_names.tools.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch tools: {err}") + if capabilities.tools is not None: + try: + tools = (await session.list_tools()).tools + for tool in tools: + name = self._component_name(tool.name, server_info) + tools_temp[name] = tool + tool_to_session_temp[name] = session + component_names.tools.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch tools: {err}") # Clean up exit stack for session if we couldn't retrieve anything # from the server. diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index 922e1eebcb..2e04fa9d4c 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -1,4 +1,5 @@ import contextlib +import logging from unittest import mock import httpx2 @@ -6,6 +7,7 @@ import pytest import mcp +from mcp import Client from mcp.client.session_group import ( ClientSessionGroup, ClientSessionParameters, @@ -13,6 +15,7 @@ StreamableHttpParameters, ) from mcp.client.stdio import StdioServerParameters +from mcp.server import Server, ServerRequestContext from mcp.shared.exceptions import MCPError @@ -142,6 +145,83 @@ async def test_client_session_group_connect_to_server(mock_exit_stack: contextli mock_session.list_prompts.assert_awaited_once() +@pytest.mark.anyio +async def test_client_session_group_connect_with_session_respects_negotiated_capabilities( + caplog: pytest.LogCaptureFixture, +) -> None: + async def handle_list_tools( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListToolsResult: + return types.ListToolsResult( + tools=[ + types.Tool( + name="ping", + description="Ping", + input_schema={"type": "object", "properties": {}}, + ) + ] + ) + + async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult: + return types.CallToolResult(content=[types.TextContent(type="text", text="pong")]) + + server = Server( + "tools-only-server", + on_list_tools=handle_list_tools, + on_call_tool=handle_call_tool, + ) + + group = ClientSessionGroup() + + with caplog.at_level(logging.WARNING): + async with Client(server) as client: + assert isinstance(client.server_info, types.Implementation) + assert client.server_capabilities.prompts is None + assert client.server_capabilities.resources is None + + client.session.list_prompts = mock.AsyncMock(side_effect=AssertionError("list_prompts() was called")) + client.session.list_resources = mock.AsyncMock(side_effect=AssertionError("list_resources() was called")) + + await group.connect_with_session(client.server_info, client.session) + result = await group.call_tool("ping") + assert result.content == [types.TextContent(type="text", text="pong")] + + assert not caplog.records + + +@pytest.mark.anyio +async def test_client_session_group_skips_unadvertised_tools_and_resources( + caplog: pytest.LogCaptureFixture, +) -> None: + 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 isinstance(client.server_info, types.Implementation) + assert client.server_capabilities.tools is None + assert client.server_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.server_info, client.session) + assert group.prompts == {"hello": types.Prompt(name="hello", description="Hello", arguments=[])} + assert group.tools == {} + assert group.resources == {} + + assert not caplog.records + + @pytest.mark.anyio async def test_client_session_group_connect_to_server_with_name_hook(mock_exit_stack: contextlib.AsyncExitStack): """Test connecting with a component name hook."""