diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..4bfc49f3af 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -70,6 +70,35 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand `ClientSessionGroup` is built on `ClientSession`, not on `Client`. Each `connect_to_server` runs the classic `initialize` handshake. It never sends the `server/discover` probe described in **[Protocol versions](../protocol-versions.md)**. Every MCP server understands that handshake, so this costs you compatibility with nothing; it only means a group takes the older, slower path to a server that could do better. +## Authentication + +When connecting to HTTP servers using `StreamableHttpParameters` or `SseServerParameters`, you can configure an authentication provider (such as an OAuth 2.0 `OAuthClientProvider` or custom `httpx2.Auth` handler) via the `auth=` parameter: + +```python +from mcp.client.auth import OAuthClientProvider +from mcp.client.session_group import ClientSessionGroup, StreamableHttpParameters + + +async def main() -> None: + server_auth = OAuthClientProvider( + server_url="https://api.example.com", + client_metadata=client_metadata, + storage=token_storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + ) + + server_params = StreamableHttpParameters( + url="https://api.example.com/mcp", + auth=server_auth, + ) + + async with ClientSessionGroup() as group: + await group.connect_to_server(server_params) +``` + +Because `auth` is configured per `ServerParameters` instance, each server in the session group maintains independent authentication context, scopes, and token-refresh lifecycle. Custom headers can still be supplied alongside `auth` via `headers=`. + ## Recap * `ClientSessionGroup` holds many server connections and merges their tools, resources, and prompts into one `dict` each. diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index a544cecbe8..92b299e844 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -16,7 +16,7 @@ import anyio import httpx2 import mcp_types as types -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import Self import mcp @@ -32,6 +32,8 @@ class SseServerParameters(BaseModel): """Parameters for initializing an sse_client.""" + model_config = ConfigDict(arbitrary_types_allowed=True) + # The endpoint URL. url: str @@ -44,10 +46,15 @@ class SseServerParameters(BaseModel): # Timeout for SSE read operations (in seconds). sse_read_timeout: float = 300.0 + # Optional HTTPX authentication handler. + auth: httpx2.Auth | None = Field(default=None, description="Optional HTTPX authentication handler.", exclude=True) + class StreamableHttpParameters(BaseModel): """Parameters for initializing a streamable_http_client.""" + model_config = ConfigDict(arbitrary_types_allowed=True) + # The endpoint URL. url: str @@ -63,6 +70,9 @@ class StreamableHttpParameters(BaseModel): # Close the client session when the transport closes. terminate_on_close: bool = True + # Optional HTTPX authentication handler. + auth: httpx2.Auth | None = Field(default=None, description="Optional HTTPX authentication handler.", exclude=True) + ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters @@ -319,6 +329,7 @@ async def _establish_session( headers=server_params.headers, timeout=server_params.timeout, sse_read_timeout=server_params.sse_read_timeout, + auth=server_params.auth, ) read, write = await session_stack.enter_async_context(client) else: @@ -328,6 +339,7 @@ async def _establish_session( server_params.timeout, read=server_params.sse_read_timeout, ), + auth=server_params.auth, ) await session_stack.enter_async_context(httpx_client) diff --git a/src/mcp/shared/_httpx_utils.py b/src/mcp/shared/_httpx_utils.py index 940b9f08cc..d3e38f02e5 100644 --- a/src/mcp/shared/_httpx_utils.py +++ b/src/mcp/shared/_httpx_utils.py @@ -59,7 +59,7 @@ def create_mcp_http_client( kwargs: dict[str, Any] = {"timeout": timeout} if headers is not None: kwargs["headers"] = headers - if auth is not None: # pragma: no cover + if auth is not None: kwargs["auth"] = auth return httpx2.AsyncClient(**kwargs) diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index b75d22b7a0..fdd571991a 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -1,4 +1,5 @@ import contextlib +from typing import Any from unittest import mock import httpx2 @@ -372,6 +373,7 @@ async def test_client_session_group_establish_session_parameterized( headers=server_params_instance.headers, timeout=server_params_instance.timeout, sse_read_timeout=server_params_instance.sse_read_timeout, + auth=server_params_instance.auth, ) elif client_type_name == "streamablehttp": # pragma: no branch assert isinstance(server_params_instance, StreamableHttpParameters) @@ -381,6 +383,7 @@ async def test_client_session_group_establish_session_parameterized( assert call_args.kwargs["url"] == server_params_instance.url assert call_args.kwargs["terminate_on_close"] == server_params_instance.terminate_on_close assert isinstance(call_args.kwargs["http_client"], httpx2.AsyncClient) + assert call_args.kwargs["http_client"].auth is server_params_instance.auth mock_client_cm_instance.__aenter__.assert_awaited_once() @@ -402,3 +405,223 @@ async def test_client_session_group_establish_session_parameterized( # 3. Assert returned values assert returned_server_info is mock_initialize_result.server_info assert returned_session is mock_entered_session + + +class _FakeBearerAuth(httpx2.Auth): + """Simple test auth implementation decorating requests with a bearer token.""" + + def __init__(self, token: str) -> None: + self.token = token + self.call_count = 0 + + def auth_flow(self, request: httpx2.Request): + self.call_count += 1 + request.headers["Authorization"] = f"Bearer {self.token}" + yield request + + +@pytest.mark.anyio +@mock.patch("mcp.client.session_group.sse_client") +@mock.patch("mcp.client.session_group.mcp.ClientSession") +async def test_establish_session_sse_passes_auth( + mock_ClientSession_class: mock.MagicMock, + mock_sse_client: mock.MagicMock, +): + """_establish_session should pass auth to sse_client for SseServerParameters.""" + mock_auth = mock.Mock(spec=httpx2.Auth) + server_params = SseServerParameters(url="http://test.com/sse", auth=mock_auth) + + mock_client_cm = mock.AsyncMock() + mock_read = mock.AsyncMock() + mock_write = mock.AsyncMock() + mock_client_cm.__aenter__.return_value = (mock_read, mock_write) + mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None) + mock_sse_client.return_value = mock_client_cm + + mock_session_cm = mock.AsyncMock() + mock_ClientSession_class.return_value = mock_session_cm + mock_session = mock.AsyncMock() + mock_session_cm.__aenter__.return_value = mock_session + mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None) + + mock_result = mock.AsyncMock() + mock_result.server_info = types.Implementation(name="test", version="1") + mock_session.initialize.return_value = mock_result + + group = ClientSessionGroup() + async with contextlib.AsyncExitStack() as stack: + group._exit_stack = stack + await group._establish_session(server_params, ClientSessionParameters()) + + mock_sse_client.assert_called_once_with( + url="http://test.com/sse", + headers=None, + timeout=5.0, + sse_read_timeout=300.0, + auth=mock_auth, + ) + + +@pytest.mark.anyio +@mock.patch("mcp.client.session_group.create_mcp_http_client") +@mock.patch("mcp.client.session_group.streamable_http_client") +@mock.patch("mcp.client.session_group.mcp.ClientSession") +async def test_establish_session_streamable_http_passes_auth( + mock_ClientSession_class: mock.MagicMock, + mock_streamable_client: mock.MagicMock, + mock_create_client: mock.MagicMock, +): + """_establish_session should pass auth to create_mcp_http_client for StreamableHttpParameters.""" + mock_auth = mock.Mock(spec=httpx2.Auth) + server_params = StreamableHttpParameters(url="http://test.com/stream", auth=mock_auth) + + mock_httpx_client = mock.AsyncMock(spec=httpx2.AsyncClient) + mock_httpx_client.__aenter__ = mock.AsyncMock(return_value=mock_httpx_client) + mock_httpx_client.__aexit__ = mock.AsyncMock(return_value=None) + mock_create_client.return_value = mock_httpx_client + + mock_client_cm = mock.AsyncMock() + mock_read = mock.AsyncMock() + mock_write = mock.AsyncMock() + mock_client_cm.__aenter__.return_value = (mock_read, mock_write) + mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None) + mock_streamable_client.return_value = mock_client_cm + + mock_session_cm = mock.AsyncMock() + mock_ClientSession_class.return_value = mock_session_cm + mock_session = mock.AsyncMock() + mock_session_cm.__aenter__.return_value = mock_session + mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None) + + mock_result = mock.AsyncMock() + mock_result.server_info = types.Implementation(name="test", version="1") + mock_session.initialize.return_value = mock_result + + group = ClientSessionGroup() + async with contextlib.AsyncExitStack() as stack: + group._exit_stack = stack + await group._establish_session(server_params, ClientSessionParameters()) + + mock_create_client.assert_called_once() + call_kwargs = mock_create_client.call_args.kwargs + assert call_kwargs["auth"] is mock_auth + + +def test_server_parameters_auth_model_config_and_serialization(): + """Verify that auth works with Pydantic arbitrary types and is excluded on dump.""" + fake_auth = _FakeBearerAuth(token="secret-token-123") + req = httpx2.Request("GET", "http://test.com") + list(fake_auth.auth_flow(req)) + assert fake_auth.call_count == 1 + assert req.headers["Authorization"] == "Bearer secret-token-123" + + sse_params = SseServerParameters(url="http://test.com/sse", auth=fake_auth) + assert sse_params.auth is fake_auth + dumped_sse = sse_params.model_dump() + assert "auth" not in dumped_sse + json_sse = sse_params.model_dump_json() + assert "secret-token-123" not in json_sse + + stream_params = StreamableHttpParameters(url="http://test.com/stream", auth=fake_auth) + assert stream_params.auth is fake_auth + dumped_stream = stream_params.model_dump() + assert "auth" not in dumped_stream + json_stream = stream_params.model_dump_json() + assert "secret-token-123" not in json_stream + + +@pytest.mark.anyio +async def test_session_group_multiple_servers_isolated_auth(): + """Verify that multiple servers in a group receive independent auth instances.""" + auth_a = _FakeBearerAuth(token="token-server-a") + auth_b = _FakeBearerAuth(token="token-server-b") + + params_a = StreamableHttpParameters(url="http://server-a.com/stream", auth=auth_a) + params_b = StreamableHttpParameters(url="http://server-b.com/stream", auth=auth_b) + + recorded_clients: list[httpx2.AsyncClient] = [] + + with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_session_cls: + with mock.patch("mcp.client.session_group.streamable_http_client") as mock_stream_client: + mock_client_cm = mock.AsyncMock() + mock_client_cm.__aenter__.return_value = (mock.AsyncMock(), mock.AsyncMock()) + mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None) + + def record_stream_client(*args: Any, **kwargs: Any) -> mock.AsyncMock: + client = kwargs.get("http_client") + assert isinstance(client, httpx2.AsyncClient) + recorded_clients.append(client) + return mock_client_cm + + mock_stream_client.side_effect = record_stream_client + + mock_session = mock.AsyncMock() + mock_session.initialize.return_value = mock.AsyncMock( + server_info=types.Implementation(name="srv", version="1") + ) + mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[]) + mock_session.list_resources.return_value = mock.AsyncMock(resources=[]) + mock_session.list_tools.return_value = mock.AsyncMock(tools=[]) + + mock_session_cm = mock.AsyncMock() + mock_session_cm.__aenter__.return_value = mock_session + mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None) + mock_session_cls.return_value = mock_session_cm + + group = ClientSessionGroup() + async with contextlib.AsyncExitStack() as stack: + group._exit_stack = stack + await group._establish_session(params_a, ClientSessionParameters()) + await group._establish_session(params_b, ClientSessionParameters()) + + assert len(recorded_clients) == 2 + assert recorded_clients[0].auth is auth_a + assert recorded_clients[1].auth is auth_b + assert recorded_clients[0].auth is not recorded_clients[1].auth + + +@pytest.mark.anyio +async def test_session_group_headers_and_auth_coexistence(): + """Verify that custom headers and auth handler coexist on created HTTP clients.""" + custom_auth = _FakeBearerAuth(token="auth-bearer-token") + custom_headers = {"X-Trace-Id": "trace-999", "Authorization": "Bearer manual"} + + params = StreamableHttpParameters( + url="http://test.com/stream", + headers=custom_headers, + auth=custom_auth, + ) + + with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_session_cls: + with mock.patch("mcp.client.session_group.streamable_http_client") as mock_stream_client: + created_client: list[httpx2.AsyncClient] = [] + mock_client_cm = mock.AsyncMock() + mock_client_cm.__aenter__.return_value = (mock.AsyncMock(), mock.AsyncMock()) + mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None) + + def record_client(*args: Any, **kwargs: Any) -> mock.AsyncMock: + client = kwargs.get("http_client") + assert isinstance(client, httpx2.AsyncClient) + created_client.append(client) + return mock_client_cm + + mock_stream_client.side_effect = record_client + + mock_session = mock.AsyncMock() + mock_session.initialize.return_value = mock.AsyncMock( + server_info=types.Implementation(name="srv", version="1") + ) + mock_session_cm = mock.AsyncMock() + mock_session_cm.__aenter__.return_value = mock_session + mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None) + mock_session_cls.return_value = mock_session_cm + + group = ClientSessionGroup() + async with contextlib.AsyncExitStack() as stack: + group._exit_stack = stack + await group._establish_session(params, ClientSessionParameters()) + + assert len(created_client) == 1 + client = created_client[0] + assert client.auth is custom_auth + assert client.headers.get("X-Trace-Id") == "trace-999"