Skip to content

Commit ff78bb6

Browse files
Pass httpx auth through ClientSessionGroup HTTP parameters
SseServerParameters and StreamableHttpParameters can carry an httpx2.Auth, and connect_to_server forwards it to the HTTP client. That is what OAuthClientProvider needs. The handler is left out of serialization. Co-authored-by: Michael <265398295+lafalce-assistant@users.noreply.github.com>
1 parent 77001c9 commit ff78bb6

7 files changed

Lines changed: 274 additions & 3 deletions

File tree

‎docs/client/session-groups.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,20 @@ 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+
An HTTP server takes the same `httpx2.Auth` you would put on an `httpx2.AsyncClient`. `OAuthClientProvider` is one. Build it as in **[OAuth clients](oauth-clients.md)**, then pass it on the parameters. The group builds the HTTP client, so there is no `http_client=` argument here.
76+
77+
```python title="client.py" hl_lines="26"
78+
--8<-- "docs_src/session_groups/tutorial005.py"
79+
```
80+
81+
* One auth object per server. `OAuthClientProvider` keeps its tokens on that object, so two servers that share one instance share one refresh cycle. Construct a new one for each `StreamableHttpParameters` or `SseServerParameters`.
82+
* `headers=` is still sent. A header other than `Authorization` arrives as you set it. `Authorization` does not, once the handler has a token: the handler writes that header itself. Before it has a token, a manual `Authorization` header is what goes out.
83+
* `auth` is left out of `model_dump()`, `model_dump_json()`, and `repr()`. It can hold a credential, and it is not JSON. Build it again when you build the parameters again.
84+
85+
`SseServerParameters` takes the same `auth`.
86+
7387
## Recap
7488

7589
* `ClientSessionGroup` holds many server connections and merges their tools, resources, and prompts into one `dict` each.
@@ -78,5 +92,6 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand
7892
* Names must be unique across the whole group; two servers with a `search` tool cannot coexist on their own.
7993
* `component_name_hook=` rewrites every registered name. The dict key changes, the wire name does not.
8094
* `connect_with_session` adds a session you already hold; `disconnect_from_server` removes one.
95+
* HTTP parameters take `auth=`, an `httpx2.Auth`. It is not serialized. One instance per server.
8196

8297
The handshake a group speaks (and the faster one a `Client` prefers) is the subject of **[Protocol versions](../protocol-versions.md)**.

‎docs/client/transports.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ same option on `StreamableHttpParameters`.
8080
proxies, event hooks, retries and connection limits here. The SDK adds nothing on top and takes
8181
nothing away, except [redirect handling](#redirects). It is also where OAuth plugs in:
8282
`httpx2.AsyncClient(auth=OAuthClientProvider(...))`. That whole flow is **[OAuth clients](oauth-clients.md)**.
83+
A `ClientSessionGroup` has no `http_client=` argument. Put that same object on `auth=` of
84+
`StreamableHttpParameters` or `SseServerParameters` (**[Multiple servers](session-groups.md#authentication)**).
8385

8486
### Redirects
8587

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import asyncio
2+
from collections.abc import Generator
3+
4+
import httpx2
5+
6+
from mcp.client.session_group import ClientSessionGroup, StreamableHttpParameters
7+
8+
9+
class BearerAuth(httpx2.Auth):
10+
"""Attach one static bearer token.
11+
12+
A real OAuth flow uses `OAuthClientProvider` in this same `auth=` slot.
13+
"""
14+
15+
def __init__(self, token: str) -> None:
16+
self.token = token
17+
18+
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
19+
request.headers["Authorization"] = f"Bearer {self.token}"
20+
yield request
21+
22+
23+
library = StreamableHttpParameters(
24+
url="http://127.0.0.1:8000/mcp",
25+
headers={"X-Trace-Id": "library"},
26+
auth=BearerAuth("library-token"),
27+
)
28+
29+
30+
async def main() -> None:
31+
async with ClientSessionGroup() as group:
32+
await group.connect_to_server(library)
33+
34+
35+
if __name__ == "__main__":
36+
asyncio.run(main())

‎src/mcp/client/session_group.py‎

Lines changed: 15 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,9 @@
3131
class SseServerParameters(BaseModel):
3232
"""Parameters for initializing an sse_client."""
3333

34+
# httpx2.Auth is not a pydantic type. Omitted from dumps and repr: it can hold credentials.
35+
model_config = ConfigDict(arbitrary_types_allowed=True)
36+
3437
# The endpoint URL.
3538
url: str
3639

@@ -43,10 +46,16 @@ class SseServerParameters(BaseModel):
4346
# Timeout for SSE read operations (in seconds).
4447
sse_read_timeout: float = 300.0
4548

49+
# Optional httpx2 authentication handler, for example OAuthClientProvider.
50+
auth: httpx2.Auth | None = Field(default=None, exclude=True, repr=False)
51+
4652

4753
class StreamableHttpParameters(BaseModel):
4854
"""Parameters for initializing a streamable_http_client."""
4955

56+
# httpx2.Auth is not a pydantic type. Omitted from dumps and repr: it can hold credentials.
57+
model_config = ConfigDict(arbitrary_types_allowed=True)
58+
5059
# The endpoint URL.
5160
url: str
5261

@@ -65,6 +74,9 @@ class StreamableHttpParameters(BaseModel):
6574
# Maximum bytes in one server-sent event. None disables the limit.
6675
max_sse_event_size: int | None = Field(default=DEFAULT_MAX_SSE_EVENT_SIZE, gt=0)
6776

77+
# Optional httpx2 authentication handler, for example OAuthClientProvider.
78+
auth: httpx2.Auth | None = Field(default=None, exclude=True, repr=False)
79+
6880

6981
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
7082

@@ -321,6 +333,7 @@ async def _establish_session(
321333
headers=server_params.headers,
322334
timeout=server_params.timeout,
323335
sse_read_timeout=server_params.sse_read_timeout,
336+
auth=server_params.auth,
324337
)
325338
read, write = await session_stack.enter_async_context(client)
326339
else:
@@ -330,6 +343,7 @@ async def _establish_session(
330343
server_params.timeout,
331344
read=server_params.sse_read_timeout,
332345
),
346+
auth=server_params.auth,
333347
)
334348
await session_stack.enter_async_context(httpx_client)
335349

‎src/mcp/shared/_httpx_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def create_mcp_http_client(
5959
kwargs: dict[str, Any] = {"timeout": timeout}
6060
if headers is not None:
6161
kwargs["headers"] = headers
62-
if auth is not None: # pragma: no cover
62+
if auth is not None:
6363
kwargs["auth"] = auth
6464
return httpx2.AsyncClient(**kwargs)
6565

‎tests/client/test_session_group.py‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,28 @@
11
import contextlib
2+
from collections.abc import Generator
3+
from typing import Any
24
from unittest import mock
35

6+
import anyio
47
import httpx2
58
import mcp_types as types
69
import pytest
10+
from starlette.types import ASGIApp, Receive, Scope, Send
711

812
import mcp
13+
import mcp.shared._httpx_utils as httpx_utils
914
from mcp.client.session_group import (
1015
ClientSessionGroup,
1116
ClientSessionParameters,
1217
SseServerParameters,
1318
StreamableHttpParameters,
1419
)
1520
from mcp.client.stdio import StdioServerParameters
21+
from mcp.server import MCPServer
22+
from mcp.server.transport_security import TransportSecuritySettings
1623
from mcp.shared.exceptions import MCPError
24+
from tests.interaction._connect import BASE_URL, build_sse_app
25+
from tests.interaction.transports import StreamingASGITransport
1726

1827

1928
@pytest.fixture
@@ -374,6 +383,7 @@ async def test_client_session_group_establish_session_parameterized(
374383
headers=server_params_instance.headers,
375384
timeout=server_params_instance.timeout,
376385
sse_read_timeout=server_params_instance.sse_read_timeout,
386+
auth=server_params_instance.auth,
377387
)
378388
elif client_type_name == "streamablehttp": # pragma: no branch
379389
assert isinstance(server_params_instance, StreamableHttpParameters)
@@ -405,3 +415,183 @@ async def test_client_session_group_establish_session_parameterized(
405415
# 3. Assert returned values
406416
assert returned_server_info is mock_initialize_result.server_info
407417
assert returned_session is mock_entered_session
418+
419+
420+
class _BearerAuth(httpx2.Auth):
421+
"""Test auth handler. `__repr__` includes the token so a leak into `repr(params)` fails the test."""
422+
423+
def __init__(self, token: str) -> None:
424+
self.token = token
425+
426+
def __repr__(self) -> str:
427+
return f"_BearerAuth({self.token})"
428+
429+
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
430+
request.headers["Authorization"] = f"Bearer {self.token}"
431+
yield request
432+
433+
434+
def _by_host(apps: dict[str, ASGIApp]) -> ASGIApp:
435+
"""One ASGI app that forwards to the server for the request's host.
436+
437+
Each HTTP client then owns a single streaming transport. One client entering several
438+
transports exits their cancel scopes from the wrong task.
439+
"""
440+
441+
async def app(scope: Scope, receive: Receive, send: Send) -> None:
442+
if scope["type"] != "http":
443+
return
444+
host = next(value.decode("latin-1") for name, value in scope["headers"] if name == b"host")
445+
await apps[host.split(":", 1)[0]](scope, receive, send)
446+
447+
return app
448+
449+
450+
def _recording_app(app: ASGIApp, authorizations: list[str | None], traces: list[str | None]) -> ASGIApp:
451+
async def wrapped(scope: Scope, receive: Receive, send: Send) -> None:
452+
if scope["type"] == "http":
453+
headers = {name.decode("latin-1").lower(): value.decode("latin-1") for name, value in scope["headers"]}
454+
authorizations.append(headers.get("authorization"))
455+
traces.append(headers.get("x-trace-id"))
456+
await app(scope, receive, send)
457+
458+
return wrapped
459+
460+
461+
def _server(name: str, tool_name: str) -> MCPServer:
462+
server = MCPServer(name)
463+
464+
@server.tool(name=tool_name)
465+
def _call() -> str:
466+
"""Return the server name."""
467+
return name
468+
469+
return server
470+
471+
472+
def _streamable_app(server: MCPServer, authorizations: list[str | None], traces: list[str | None]) -> ASGIApp:
473+
app = server.streamable_http_app(
474+
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False)
475+
)
476+
return _recording_app(app, authorizations, traces)
477+
478+
479+
def _patch_in_process_clients(monkeypatch: pytest.MonkeyPatch, app: ASGIApp, *, cancel_on_close: bool) -> None:
480+
"""Make `create_mcp_http_client` talk to `app`. The group builds that client itself."""
481+
482+
class InProcessClient(httpx2.AsyncClient):
483+
def __init__(self, *args: Any, **kwargs: Any) -> None:
484+
kwargs["transport"] = StreamingASGITransport(app, cancel_on_close=cancel_on_close)
485+
super().__init__(*args, **kwargs)
486+
487+
monkeypatch.setattr(httpx_utils.httpx2, "AsyncClient", InProcessClient)
488+
489+
490+
def test_http_parameters_omit_auth_from_serialization() -> None:
491+
"""`auth` stays on the object and off `model_dump`, `model_dump_json`, and `repr`.
492+
493+
SDK-defined: the handler is not JSON and can hold a credential. `model_copy` keeps the same instance.
494+
"""
495+
auth = _BearerAuth("secret-token")
496+
sse = SseServerParameters(url="http://example.test/sse", auth=auth)
497+
stream = StreamableHttpParameters(url="http://example.test/mcp", auth=auth)
498+
499+
assert SseServerParameters(url="http://example.test/sse").auth is None
500+
assert StreamableHttpParameters(url="http://example.test/mcp").auth is None
501+
for params in (sse, stream):
502+
assert params.auth is auth
503+
assert params.model_copy().auth is auth
504+
assert "auth" not in params.model_dump()
505+
assert "secret-token" not in params.model_dump_json()
506+
assert "secret-token" not in repr(params)
507+
508+
509+
@pytest.mark.anyio
510+
async def test_streamable_http_group_sends_auth_and_keeps_other_headers(monkeypatch: pytest.MonkeyPatch) -> None:
511+
"""A streamable HTTP server sees the auth handler's bearer token, not a manual `Authorization` header.
512+
513+
SDK-defined. `X-Trace-Id` from `headers=` is still sent. The group builds the HTTP client, so this
514+
test routes that client at an in-process app.
515+
"""
516+
server = _server("Library", "search")
517+
authorizations: list[str | None] = []
518+
traces: list[str | None] = []
519+
app = _streamable_app(server, authorizations, traces)
520+
_patch_in_process_clients(monkeypatch, app, cancel_on_close=True)
521+
params = StreamableHttpParameters(
522+
url=f"{BASE_URL}/mcp",
523+
headers={"X-Trace-Id": "trace-1", "Authorization": "Bearer manual"},
524+
auth=_BearerAuth("from-auth"),
525+
)
526+
527+
with anyio.fail_after(5):
528+
async with server.session_manager.run(), contextlib.AsyncExitStack() as stack:
529+
group = ClientSessionGroup(exit_stack=stack)
530+
await group.connect_to_server(params)
531+
result = await group.call_tool("search")
532+
533+
assert result.structured_content == {"result": "Library"}
534+
assert authorizations
535+
assert set(authorizations) == {"Bearer from-auth"}
536+
assert set(traces) == {"trace-1"}
537+
538+
539+
@pytest.mark.anyio
540+
async def test_sse_group_sends_the_auth_handlers_bearer_token(monkeypatch: pytest.MonkeyPatch) -> None:
541+
"""An SSE server sees the bearer token from `SseServerParameters.auth`.
542+
543+
SDK-defined. Same in-process route as the streamable HTTP test. `cancel_on_close` is false
544+
because the legacy SSE server only finishes the session after the GET observes a disconnect.
545+
"""
546+
server = _server("Legacy", "search")
547+
app, _sse = build_sse_app(server)
548+
authorizations: list[str | None] = []
549+
traces: list[str | None] = []
550+
_patch_in_process_clients(monkeypatch, _recording_app(app, authorizations, traces), cancel_on_close=False)
551+
params = SseServerParameters(url=f"{BASE_URL}/sse", auth=_BearerAuth("sse-token"))
552+
553+
with anyio.fail_after(5):
554+
async with contextlib.AsyncExitStack() as stack:
555+
group = ClientSessionGroup(exit_stack=stack)
556+
await group.connect_to_server(params)
557+
result = await group.call_tool("search")
558+
559+
assert result.structured_content == {"result": "Legacy"}
560+
assert authorizations
561+
assert set(authorizations) == {"Bearer sse-token"}
562+
563+
564+
@pytest.mark.anyio
565+
async def test_each_server_in_a_group_uses_its_own_auth(monkeypatch: pytest.MonkeyPatch) -> None:
566+
"""Two streamable HTTP servers in one group each see only the auth object they were given.
567+
568+
SDK-defined. The handlers are distinct instances, so a token on one server is not the other's.
569+
"""
570+
library = _server("Library", "library_search")
571+
web = _server("Web", "web_search")
572+
library_auth: list[str | None] = []
573+
web_auth: list[str | None] = []
574+
app = _by_host(
575+
{
576+
"library.example": _streamable_app(library, library_auth, []),
577+
"web.example": _streamable_app(web, web_auth, []),
578+
}
579+
)
580+
_patch_in_process_clients(monkeypatch, app, cancel_on_close=True)
581+
582+
with anyio.fail_after(5):
583+
async with library.session_manager.run(), web.session_manager.run(), contextlib.AsyncExitStack() as stack:
584+
group = ClientSessionGroup(exit_stack=stack)
585+
await group.connect_to_server(
586+
StreamableHttpParameters(url="http://library.example/mcp", auth=_BearerAuth("library-token"))
587+
)
588+
await group.connect_to_server(
589+
StreamableHttpParameters(url="http://web.example/mcp", auth=_BearerAuth("web-token"))
590+
)
591+
library_result = await group.call_tool("library_search")
592+
web_result = await group.call_tool("web_search")
593+
594+
assert library_result.structured_content == {"result": "Library"}
595+
assert web_result.structured_content == {"result": "Web"}
596+
assert set(library_auth) == {"Bearer library-token"}
597+
assert set(web_auth) == {"Bearer web-token"}

‎tests/docs_src/test_session_groups.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@
66

77
import traceback
88

9+
import httpx2
910
import pytest
1011
from mcp_types import INVALID_PARAMS, Implementation
1112

12-
from docs_src.session_groups import tutorial001, tutorial002, tutorial004
13+
from docs_src.session_groups import tutorial001, tutorial002, tutorial004, tutorial005
1314
from mcp import Client, ClientSessionGroup, MCPError
1415

1516
# See test_index.py for why this is a per-module mark and not a conftest hook.
@@ -93,6 +94,19 @@ async def test_call_tool_routes_to_the_owning_server() -> None:
9394
assert library_result.structured_content == {"result": "3 books match 'dune'."}
9495

9596

97+
def test_auth_stays_on_the_parameters_and_out_of_a_dump() -> None:
98+
"""tutorial005: `auth` is the handler, headers survive a dump, and the token does not."""
99+
params = tutorial005.library
100+
assert isinstance(params.auth, tutorial005.BearerAuth)
101+
request = httpx2.Request("POST", params.url)
102+
assert next(params.auth.auth_flow(request)) is request
103+
assert request.headers["Authorization"] == "Bearer library-token"
104+
assert params.model_dump()["headers"] == {"X-Trace-Id": "library"}
105+
assert "auth" not in params.model_dump()
106+
assert "library-token" not in params.model_dump_json()
107+
assert "library-token" not in repr(params)
108+
109+
96110
async def test_disconnect_removes_every_component_of_that_server() -> None:
97111
"""tutorial004: `disconnect_from_server` takes the session back out of all three dicts."""
98112
async with Client(tutorial001.mcp) as library, Client(tutorial002.mcp) as web:

0 commit comments

Comments
 (0)