Skip to content

Commit 0478dd4

Browse files
committed
feat(client): support HTTP auth in ClientSessionGroup
1 parent a525a3d commit 0478dd4

3 files changed

Lines changed: 251 additions & 1 deletion

File tree

‎docs/client/session-groups.md‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,33 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand
7070

7171
`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.
7272

73+
## Authentication
74+
75+
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:
76+
77+
```python
78+
from mcp.client.auth import OAuthClientProvider
79+
from mcp.client.session_group import ClientSessionGroup, StreamableHttpParameters
80+
81+
server_auth = OAuthClientProvider(
82+
server_url="https://api.example.com",
83+
client_metadata=client_metadata,
84+
storage=token_storage,
85+
redirect_handler=redirect_handler,
86+
callback_handler=callback_handler,
87+
)
88+
89+
server_params = StreamableHttpParameters(
90+
url="https://api.example.com/mcp",
91+
auth=server_auth,
92+
)
93+
94+
async with ClientSessionGroup() as group:
95+
await group.connect_to_server(server_params)
96+
```
97+
98+
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=`.
99+
73100
## Recap
74101

75102
* `ClientSessionGroup` holds many server connections and merges their tools, resources, and prompts into one `dict` each.

‎src/mcp/client/session_group.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import anyio
1717
import httpx2
1818
import mcp_types as types
19-
from pydantic import BaseModel, Field
19+
from pydantic import BaseModel, ConfigDict, Field
2020

2121
import mcp
2222
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
@@ -31,6 +31,8 @@
3131
class SseServerParameters(BaseModel):
3232
"""Parameters for initializing an sse_client."""
3333

34+
model_config = ConfigDict(arbitrary_types_allowed=True)
35+
3436
# The endpoint URL.
3537
url: str
3638

@@ -43,10 +45,15 @@ class SseServerParameters(BaseModel):
4345
# Timeout for SSE read operations (in seconds).
4446
sse_read_timeout: float = 300.0
4547

48+
# Optional HTTPX authentication handler.
49+
auth: httpx2.Auth | None = Field(default=None, description="Optional HTTPX authentication handler.", exclude=True)
50+
4651

4752
class StreamableHttpParameters(BaseModel):
4853
"""Parameters for initializing a streamable_http_client."""
4954

55+
model_config = ConfigDict(arbitrary_types_allowed=True)
56+
5057
# The endpoint URL.
5158
url: str
5259

@@ -65,6 +72,10 @@ class StreamableHttpParameters(BaseModel):
6572
# Maximum bytes in one server-sent event. None disables the limit.
6673
max_sse_event_size: int | None = Field(default=DEFAULT_MAX_SSE_EVENT_SIZE, gt=0)
6774

75+
# Optional HTTPX authentication handler.
76+
auth: httpx2.Auth | None = Field(default=None, description="Optional HTTPX authentication handler.", exclude=True)
77+
78+
6879

6980
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
7081

@@ -321,6 +332,7 @@ async def _establish_session(
321332
headers=server_params.headers,
322333
timeout=server_params.timeout,
323334
sse_read_timeout=server_params.sse_read_timeout,
335+
auth=server_params.auth,
324336
)
325337
read, write = await session_stack.enter_async_context(client)
326338
else:
@@ -330,6 +342,7 @@ async def _establish_session(
330342
server_params.timeout,
331343
read=server_params.sse_read_timeout,
332344
),
345+
auth=server_params.auth,
333346
)
334347
await session_stack.enter_async_context(httpx_client)
335348

‎tests/client/test_session_group.py‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import contextlib
2+
from typing import Any
23
from unittest import mock
34

45
import httpx2
@@ -374,6 +375,7 @@ async def test_client_session_group_establish_session_parameterized(
374375
headers=server_params_instance.headers,
375376
timeout=server_params_instance.timeout,
376377
sse_read_timeout=server_params_instance.sse_read_timeout,
378+
auth=server_params_instance.auth,
377379
)
378380
elif client_type_name == "streamablehttp": # pragma: no branch
379381
assert isinstance(server_params_instance, StreamableHttpParameters)
@@ -384,6 +386,7 @@ async def test_client_session_group_establish_session_parameterized(
384386
assert call_args.kwargs["terminate_on_close"] == server_params_instance.terminate_on_close
385387
assert call_args.kwargs["max_sse_event_size"] == server_params_instance.max_sse_event_size
386388
assert isinstance(call_args.kwargs["http_client"], httpx2.AsyncClient)
389+
assert call_args.kwargs["http_client"].auth is server_params_instance.auth
387390

388391
mock_client_cm_instance.__aenter__.assert_awaited_once()
389392

@@ -405,3 +408,210 @@ async def test_client_session_group_establish_session_parameterized(
405408
# 3. Assert returned values
406409
assert returned_server_info is mock_initialize_result.server_info
407410
assert returned_session is mock_entered_session
411+
412+
413+
class _FakeBearerAuth(httpx2.Auth):
414+
"""Simple test auth implementation decorating requests with a bearer token."""
415+
416+
def __init__(self, token: str) -> None:
417+
self.token = token
418+
self.call_count = 0
419+
420+
def auth_flow(self, request: httpx2.Request):
421+
self.call_count += 1
422+
request.headers["Authorization"] = f"Bearer {self.token}"
423+
yield request
424+
425+
426+
@pytest.mark.anyio
427+
async def test_establish_session_sse_passes_auth():
428+
"""_establish_session should pass auth to sse_client for SseServerParameters."""
429+
mock_auth = mock.Mock(spec=httpx2.Auth)
430+
server_params = SseServerParameters(url="http://test.com/sse", auth=mock_auth)
431+
432+
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class:
433+
with mock.patch("mcp.client.session_group.sse_client") as mock_sse_client:
434+
mock_client_cm = mock.AsyncMock()
435+
mock_read = mock.AsyncMock()
436+
mock_write = mock.AsyncMock()
437+
mock_client_cm.__aenter__.return_value = (mock_read, mock_write)
438+
mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None)
439+
mock_sse_client.return_value = mock_client_cm
440+
441+
mock_session_cm = mock.AsyncMock()
442+
mock_ClientSession_class.return_value = mock_session_cm
443+
mock_session = mock.AsyncMock()
444+
mock_session_cm.__aenter__.return_value = mock_session
445+
mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None)
446+
447+
mock_result = mock.AsyncMock()
448+
mock_result.server_info = types.Implementation(name="test", version="1")
449+
mock_session.initialize.return_value = mock_result
450+
451+
group = ClientSessionGroup()
452+
async with contextlib.AsyncExitStack() as stack:
453+
group._exit_stack = stack
454+
await group._establish_session(server_params, ClientSessionParameters())
455+
456+
mock_sse_client.assert_called_once_with(
457+
url="http://test.com/sse",
458+
headers=None,
459+
timeout=5.0,
460+
sse_read_timeout=300.0,
461+
auth=mock_auth,
462+
)
463+
464+
465+
@pytest.mark.anyio
466+
async def test_establish_session_streamable_http_passes_auth():
467+
"""_establish_session should pass auth to create_mcp_http_client for StreamableHttpParameters."""
468+
mock_auth = mock.Mock(spec=httpx2.Auth)
469+
server_params = StreamableHttpParameters(url="http://test.com/stream", auth=mock_auth)
470+
471+
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class:
472+
with mock.patch("mcp.client.session_group.streamable_http_client") as mock_streamable_client:
473+
with mock.patch("mcp.client.session_group.create_mcp_http_client") as mock_create_client:
474+
mock_httpx_client = mock.AsyncMock(spec=httpx2.AsyncClient)
475+
mock_httpx_client.__aenter__ = mock.AsyncMock(return_value=mock_httpx_client)
476+
mock_httpx_client.__aexit__ = mock.AsyncMock(return_value=None)
477+
mock_create_client.return_value = mock_httpx_client
478+
479+
mock_client_cm = mock.AsyncMock()
480+
mock_read = mock.AsyncMock()
481+
mock_write = mock.AsyncMock()
482+
mock_client_cm.__aenter__.return_value = (mock_read, mock_write)
483+
mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None)
484+
mock_streamable_client.return_value = mock_client_cm
485+
486+
mock_session_cm = mock.AsyncMock()
487+
mock_ClientSession_class.return_value = mock_session_cm
488+
mock_session = mock.AsyncMock()
489+
mock_session_cm.__aenter__.return_value = mock_session
490+
mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None)
491+
492+
mock_result = mock.AsyncMock()
493+
mock_result.server_info = types.Implementation(name="test", version="1")
494+
mock_session.initialize.return_value = mock_result
495+
496+
group = ClientSessionGroup()
497+
async with contextlib.AsyncExitStack() as stack:
498+
group._exit_stack = stack
499+
await group._establish_session(server_params, ClientSessionParameters())
500+
501+
mock_create_client.assert_called_once()
502+
call_kwargs = mock_create_client.call_args.kwargs
503+
assert call_kwargs["auth"] is mock_auth
504+
505+
506+
def test_server_parameters_auth_model_config_and_serialization():
507+
"""Verify that auth works with Pydantic arbitrary types and is excluded on dump."""
508+
fake_auth = _FakeBearerAuth(token="secret-token-123")
509+
sse_params = SseServerParameters(url="http://test.com/sse", auth=fake_auth)
510+
assert sse_params.auth is fake_auth
511+
dumped_sse = sse_params.model_dump()
512+
assert "auth" not in dumped_sse
513+
json_sse = sse_params.model_dump_json()
514+
assert "secret-token-123" not in json_sse
515+
516+
stream_params = StreamableHttpParameters(url="http://test.com/stream", auth=fake_auth)
517+
assert stream_params.auth is fake_auth
518+
dumped_stream = stream_params.model_dump()
519+
assert "auth" not in dumped_stream
520+
json_stream = stream_params.model_dump_json()
521+
assert "secret-token-123" not in json_stream
522+
523+
524+
@pytest.mark.anyio
525+
async def test_session_group_multiple_servers_isolated_auth():
526+
"""Verify that multiple servers in a group receive independent auth instances."""
527+
auth_a = _FakeBearerAuth(token="token-server-a")
528+
auth_b = _FakeBearerAuth(token="token-server-b")
529+
530+
params_a = StreamableHttpParameters(url="http://server-a.com/stream", auth=auth_a)
531+
params_b = StreamableHttpParameters(url="http://server-b.com/stream", auth=auth_b)
532+
533+
recorded_clients: list[httpx2.AsyncClient] = []
534+
535+
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_session_cls:
536+
with mock.patch("mcp.client.session_group.streamable_http_client") as mock_stream_client:
537+
mock_client_cm = mock.AsyncMock()
538+
mock_client_cm.__aenter__.return_value = (mock.AsyncMock(), mock.AsyncMock())
539+
mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None)
540+
541+
def record_stream_client(*args: Any, **kwargs: Any) -> mock.AsyncMock:
542+
client = kwargs.get("http_client")
543+
assert isinstance(client, httpx2.AsyncClient)
544+
recorded_clients.append(client)
545+
return mock_client_cm
546+
547+
mock_stream_client.side_effect = record_stream_client
548+
549+
mock_session = mock.AsyncMock()
550+
mock_session.initialize.return_value = mock.AsyncMock(
551+
server_info=types.Implementation(name="srv", version="1")
552+
)
553+
mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[])
554+
mock_session.list_resources.return_value = mock.AsyncMock(resources=[])
555+
mock_session.list_tools.return_value = mock.AsyncMock(tools=[])
556+
557+
mock_session_cm = mock.AsyncMock()
558+
mock_session_cm.__aenter__.return_value = mock_session
559+
mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None)
560+
mock_session_cls.return_value = mock_session_cm
561+
562+
group = ClientSessionGroup()
563+
async with group:
564+
await group.connect_to_server(params_a)
565+
await group.connect_to_server(params_b)
566+
567+
assert len(recorded_clients) == 2
568+
assert recorded_clients[0].auth is auth_a
569+
assert recorded_clients[1].auth is auth_b
570+
assert recorded_clients[0].auth is not recorded_clients[1].auth
571+
572+
573+
@pytest.mark.anyio
574+
async def test_session_group_headers_and_auth_coexistence():
575+
"""Verify that custom headers and auth handler coexist on created HTTP clients."""
576+
custom_auth = _FakeBearerAuth(token="auth-bearer-token")
577+
custom_headers = {"X-Trace-Id": "trace-999", "Authorization": "Bearer manual"}
578+
579+
params = StreamableHttpParameters(
580+
url="http://test.com/stream",
581+
headers=custom_headers,
582+
auth=custom_auth,
583+
)
584+
585+
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_session_cls:
586+
with mock.patch("mcp.client.session_group.streamable_http_client") as mock_stream_client:
587+
created_client: list[httpx2.AsyncClient] = []
588+
mock_client_cm = mock.AsyncMock()
589+
mock_client_cm.__aenter__.return_value = (mock.AsyncMock(), mock.AsyncMock())
590+
mock_client_cm.__aexit__ = mock.AsyncMock(return_value=None)
591+
592+
def record_client(*args: Any, **kwargs: Any) -> mock.AsyncMock:
593+
client = kwargs.get("http_client")
594+
assert isinstance(client, httpx2.AsyncClient)
595+
created_client.append(client)
596+
return mock_client_cm
597+
598+
mock_stream_client.side_effect = record_client
599+
600+
mock_session = mock.AsyncMock()
601+
mock_session.initialize.return_value = mock.AsyncMock(
602+
server_info=types.Implementation(name="srv", version="1")
603+
)
604+
mock_session_cm = mock.AsyncMock()
605+
mock_session_cm.__aenter__.return_value = mock_session
606+
mock_session_cm.__aexit__ = mock.AsyncMock(return_value=None)
607+
mock_session_cls.return_value = mock_session_cm
608+
609+
group = ClientSessionGroup()
610+
async with contextlib.AsyncExitStack() as stack:
611+
group._exit_stack = stack
612+
await group._establish_session(params, ClientSessionParameters())
613+
614+
assert len(created_client) == 1
615+
client = created_client[0]
616+
assert client.auth is custom_auth
617+
assert client.headers.get("X-Trace-Id") == "trace-999"

0 commit comments

Comments
 (0)