Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions docs/client/session-groups.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: The new main() is never invoked, so copying this authentication example only defines the function and never connects to the server. Add an async runner such as asyncio.run(main()) and its import.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At docs/client/session-groups.md, line 82:

<comment>The new `main()` is never invoked, so copying this authentication example only defines the function and never connects to the server. Add an async runner such as `asyncio.run(main())` and its import.</comment>

<file context>
@@ -78,21 +78,23 @@ When connecting to HTTP servers using `StreamableHttpParameters` or `SseServerPa
-async with ClientSessionGroup() as group:
-    await group.connect_to_server(server_params)
+
+async def main() -> None:
+    server_auth = OAuthClientProvider(
+        server_url="https://api.example.com",
</file context>

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=`.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: When the same auth object is assigned to two parameter instances, ClientSessionGroup shares its mutable OAuth context and tokens. Tell users to construct a distinct auth instance per server before claiming these lifecycles are independent.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At docs/client/session-groups.md, line 98:

<comment>When the same auth object is assigned to two parameter instances, `ClientSessionGroup` shares its mutable OAuth context and tokens. Tell users to construct a distinct auth instance per server before claiming these lifecycles are independent.</comment>

<file context>
@@ -70,6 +70,33 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand
+    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
</file context>
Suggested change
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=`.
Because `auth` is configured per `ServerParameters` instance, give each server a distinct auth instance to maintain 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.
Expand Down
14 changes: 13 additions & 1 deletion src/mcp/client/session_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion src/mcp/shared/_httpx_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
223 changes: 223 additions & 0 deletions tests/client/test_session_group.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import contextlib
from typing import Any
from unittest import mock

import httpx2
Expand Down Expand Up @@ -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)
Expand All @@ -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()

Expand All @@ -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"
Loading