diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..dac513064d 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -70,6 +70,20 @@ 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 + +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. + +```python title="client.py" hl_lines="26" +--8<-- "docs_src/session_groups/tutorial005.py" +``` + +* 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`. +* `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. +* `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. + +`SseServerParameters` takes the same `auth`. + ## Recap * `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 * Names must be unique across the whole group; two servers with a `search` tool cannot coexist on their own. * `component_name_hook=` rewrites every registered name. The dict key changes, the wire name does not. * `connect_with_session` adds a session you already hold; `disconnect_from_server` removes one. +* HTTP parameters take `auth=`, an `httpx2.Auth`. It is not serialized. One instance per server. The handshake a group speaks (and the faster one a `Client` prefers) is the subject of **[Protocol versions](../protocol-versions.md)**. diff --git a/docs/client/transports.md b/docs/client/transports.md index b41fe51348..43ff0c799d 100644 --- a/docs/client/transports.md +++ b/docs/client/transports.md @@ -80,6 +80,8 @@ same option on `StreamableHttpParameters`. proxies, event hooks, retries and connection limits here. The SDK adds nothing on top and takes nothing away, except [redirect handling](#redirects). It is also where OAuth plugs in: `httpx2.AsyncClient(auth=OAuthClientProvider(...))`. That whole flow is **[OAuth clients](oauth-clients.md)**. + A `ClientSessionGroup` has no `http_client=` argument. Put that same object on `auth=` of + `StreamableHttpParameters` or `SseServerParameters` (**[Multiple servers](session-groups.md#authentication)**). ### Redirects diff --git a/docs_src/session_groups/tutorial005.py b/docs_src/session_groups/tutorial005.py new file mode 100644 index 0000000000..955e734f67 --- /dev/null +++ b/docs_src/session_groups/tutorial005.py @@ -0,0 +1,36 @@ +import asyncio +from collections.abc import Generator + +import httpx2 + +from mcp.client.session_group import ClientSessionGroup, StreamableHttpParameters + + +class BearerAuth(httpx2.Auth): + """Attach one static bearer token. + + A real OAuth flow uses `OAuthClientProvider` in this same `auth=` slot. + """ + + def __init__(self, token: str) -> None: + self.token = token + + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + request.headers["Authorization"] = f"Bearer {self.token}" + yield request + + +library = StreamableHttpParameters( + url="http://127.0.0.1:8000/mcp", + headers={"X-Trace-Id": "library"}, + auth=BearerAuth("library-token"), +) + + +async def main() -> None: + async with ClientSessionGroup() as group: + await group.connect_to_server(library) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index 13d8928ed4..1a2d55be6c 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 import mcp from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT @@ -31,6 +31,9 @@ class SseServerParameters(BaseModel): """Parameters for initializing an sse_client.""" + # httpx2.Auth is not a pydantic type. Omitted from dumps and repr: it can hold credentials. + model_config = ConfigDict(arbitrary_types_allowed=True) + # The endpoint URL. url: str @@ -43,10 +46,16 @@ class SseServerParameters(BaseModel): # Timeout for SSE read operations (in seconds). sse_read_timeout: float = 300.0 + # Optional httpx2 authentication handler, for example OAuthClientProvider. + auth: httpx2.Auth | None = Field(default=None, exclude=True, repr=False) + class StreamableHttpParameters(BaseModel): """Parameters for initializing a streamable_http_client.""" + # httpx2.Auth is not a pydantic type. Omitted from dumps and repr: it can hold credentials. + model_config = ConfigDict(arbitrary_types_allowed=True) + # The endpoint URL. url: str @@ -65,6 +74,9 @@ class StreamableHttpParameters(BaseModel): # Maximum bytes in one server-sent event. None disables the limit. max_sse_event_size: int | None = Field(default=DEFAULT_MAX_SSE_EVENT_SIZE, gt=0) + # Optional httpx2 authentication handler, for example OAuthClientProvider. + auth: httpx2.Auth | None = Field(default=None, exclude=True, repr=False) + ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters @@ -321,6 +333,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: @@ -330,6 +343,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 7107e5e535..b5d1bf0af5 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 922e1eebcb..c57161f9e0 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -1,11 +1,16 @@ import contextlib +from collections.abc import Generator +from typing import Any from unittest import mock +import anyio import httpx2 import mcp_types as types import pytest +from starlette.types import ASGIApp, Receive, Scope, Send import mcp +import mcp.shared._httpx_utils as httpx_utils from mcp.client.session_group import ( ClientSessionGroup, ClientSessionParameters, @@ -13,7 +18,11 @@ StreamableHttpParameters, ) from mcp.client.stdio import StdioServerParameters +from mcp.server import MCPServer +from mcp.server.transport_security import TransportSecuritySettings from mcp.shared.exceptions import MCPError +from tests.interaction._connect import BASE_URL, build_sse_app +from tests.interaction.transports import StreamingASGITransport @pytest.fixture @@ -374,6 +383,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) @@ -405,3 +415,182 @@ 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 _BearerAuth(httpx2.Auth): + """Test auth handler. `__repr__` includes the token so a leak into `repr(params)` fails the test.""" + + def __init__(self, token: str) -> None: + self.token = token + + def __repr__(self) -> str: + return f"_BearerAuth({self.token})" + + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + request.headers["Authorization"] = f"Bearer {self.token}" + yield request + + +def _by_host(apps: dict[str, ASGIApp]) -> ASGIApp: + """One ASGI app that forwards to the server for the request's host. + + Each HTTP client then owns a single streaming transport. One client entering several + transports exits their cancel scopes from the wrong task. + """ + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + host = next(value.decode("latin-1") for name, value in scope["headers"] if name == b"host") + await apps[host.split(":", 1)[0]](scope, receive, send) + + return app + + +def _recording_app(app: ASGIApp, authorizations: list[str | None], traces: list[str | None]) -> ASGIApp: + async def wrapped(scope: Scope, receive: Receive, send: Send) -> None: + headers = {name.decode("latin-1").lower(): value.decode("latin-1") for name, value in scope["headers"]} + authorizations.append(headers.get("authorization")) + traces.append(headers.get("x-trace-id")) + await app(scope, receive, send) + + return wrapped + + +def _server(name: str, tool_name: str) -> MCPServer: + server = MCPServer(name) + + @server.tool(name=tool_name) + def _call() -> str: + """Return the server name.""" + return name + + return server + + +def _streamable_app(server: MCPServer, authorizations: list[str | None], traces: list[str | None]) -> ASGIApp: + app = server.streamable_http_app( + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False) + ) + return _recording_app(app, authorizations, traces) + + +def _patch_in_process_clients(monkeypatch: pytest.MonkeyPatch, app: ASGIApp, *, cancel_on_close: bool) -> None: + """Make `create_mcp_http_client` talk to `app`. The group builds that client itself.""" + + class InProcessClient(httpx2.AsyncClient): + def __init__(self, *args: Any, **kwargs: Any) -> None: + kwargs["transport"] = StreamingASGITransport(app, cancel_on_close=cancel_on_close) + super().__init__(*args, **kwargs) + + monkeypatch.setattr(httpx_utils.httpx2, "AsyncClient", InProcessClient) + + +def test_http_parameters_omit_auth_from_serialization() -> None: + """`auth` stays on the object and off `model_dump`, `model_dump_json`, and `repr`. + + SDK-defined: the handler is not JSON and can hold a credential. `model_copy` keeps the same instance. + """ + auth = _BearerAuth("secret-token") + sse = SseServerParameters(url="http://example.test/sse", auth=auth) + stream = StreamableHttpParameters(url="http://example.test/mcp", auth=auth) + + assert SseServerParameters(url="http://example.test/sse").auth is None + assert StreamableHttpParameters(url="http://example.test/mcp").auth is None + # The token is visible on the handler itself, so a leak into the parameters repr would fail below. + assert repr(auth) == "_BearerAuth(secret-token)" + for params in (sse, stream): + assert params.auth is auth + assert params.model_copy().auth is auth + assert "auth" not in params.model_dump() + assert "secret-token" not in params.model_dump_json() + assert "secret-token" not in repr(params) + + +@pytest.mark.anyio +async def test_streamable_http_group_sends_auth_and_keeps_other_headers(monkeypatch: pytest.MonkeyPatch) -> None: + """A streamable HTTP server sees the auth handler's bearer token, not a manual `Authorization` header. + + SDK-defined. `X-Trace-Id` from `headers=` is still sent. The group builds the HTTP client, so this + test routes that client at an in-process app. + """ + server = _server("Library", "search") + authorizations: list[str | None] = [] + traces: list[str | None] = [] + app = _streamable_app(server, authorizations, traces) + _patch_in_process_clients(monkeypatch, app, cancel_on_close=True) + params = StreamableHttpParameters( + url=f"{BASE_URL}/mcp", + headers={"X-Trace-Id": "trace-1", "Authorization": "Bearer manual"}, + auth=_BearerAuth("from-auth"), + ) + + with anyio.fail_after(5): + async with server.session_manager.run(), contextlib.AsyncExitStack() as stack: + group = ClientSessionGroup(exit_stack=stack) + await group.connect_to_server(params) + result = await group.call_tool("search") + + assert result.structured_content == {"result": "Library"} + assert authorizations + assert set(authorizations) == {"Bearer from-auth"} + assert set(traces) == {"trace-1"} + + +@pytest.mark.anyio +async def test_sse_group_sends_the_auth_handlers_bearer_token(monkeypatch: pytest.MonkeyPatch) -> None: + """An SSE server sees the bearer token from `SseServerParameters.auth`. + + SDK-defined. Same in-process route as the streamable HTTP test. `cancel_on_close` is false + because the legacy SSE server only finishes the session after the GET observes a disconnect. + """ + server = _server("Legacy", "search") + app, _sse = build_sse_app(server) + authorizations: list[str | None] = [] + traces: list[str | None] = [] + _patch_in_process_clients(monkeypatch, _recording_app(app, authorizations, traces), cancel_on_close=False) + params = SseServerParameters(url=f"{BASE_URL}/sse", auth=_BearerAuth("sse-token")) + + with anyio.fail_after(5): + async with contextlib.AsyncExitStack() as stack: + group = ClientSessionGroup(exit_stack=stack) + await group.connect_to_server(params) + result = await group.call_tool("search") + + assert result.structured_content == {"result": "Legacy"} + assert authorizations + assert set(authorizations) == {"Bearer sse-token"} + + +@pytest.mark.anyio +async def test_each_server_in_a_group_uses_its_own_auth(monkeypatch: pytest.MonkeyPatch) -> None: + """Two streamable HTTP servers in one group each see only the auth object they were given. + + SDK-defined. The handlers are distinct instances, so a token on one server is not the other's. + """ + library = _server("Library", "library_search") + web = _server("Web", "web_search") + library_auth: list[str | None] = [] + web_auth: list[str | None] = [] + app = _by_host( + { + "library.example": _streamable_app(library, library_auth, []), + "web.example": _streamable_app(web, web_auth, []), + } + ) + _patch_in_process_clients(monkeypatch, app, cancel_on_close=True) + + with anyio.fail_after(5): + async with library.session_manager.run(), web.session_manager.run(), contextlib.AsyncExitStack() as stack: + group = ClientSessionGroup(exit_stack=stack) + await group.connect_to_server( + StreamableHttpParameters(url="http://library.example/mcp", auth=_BearerAuth("library-token")) + ) + await group.connect_to_server( + StreamableHttpParameters(url="http://web.example/mcp", auth=_BearerAuth("web-token")) + ) + library_result = await group.call_tool("library_search") + web_result = await group.call_tool("web_search") + + assert library_result.structured_content == {"result": "Library"} + assert web_result.structured_content == {"result": "Web"} + assert set(library_auth) == {"Bearer library-token"} + assert set(web_auth) == {"Bearer web-token"} diff --git a/tests/docs_src/test_session_groups.py b/tests/docs_src/test_session_groups.py index b8d54db176..13a5c3d0fc 100644 --- a/tests/docs_src/test_session_groups.py +++ b/tests/docs_src/test_session_groups.py @@ -6,10 +6,11 @@ import traceback +import httpx2 import pytest from mcp_types import INVALID_PARAMS, Implementation -from docs_src.session_groups import tutorial001, tutorial002, tutorial004 +from docs_src.session_groups import tutorial001, tutorial002, tutorial004, tutorial005 from mcp import Client, ClientSessionGroup, MCPError # 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: assert library_result.structured_content == {"result": "3 books match 'dune'."} +def test_auth_stays_on_the_parameters_and_out_of_a_dump() -> None: + """tutorial005: `auth` is the handler, headers survive a dump, and the token does not.""" + params = tutorial005.library + assert isinstance(params.auth, tutorial005.BearerAuth) + request = httpx2.Request("POST", params.url) + assert next(params.auth.auth_flow(request)) is request + assert request.headers["Authorization"] == "Bearer library-token" + assert params.model_dump()["headers"] == {"X-Trace-Id": "library"} + assert "auth" not in params.model_dump() + assert "library-token" not in params.model_dump_json() + assert "library-token" not in repr(params) + + async def test_disconnect_removes_every_component_of_that_server() -> None: """tutorial004: `disconnect_from_server` takes the session back out of all three dicts.""" async with Client(tutorial001.mcp) as library, Client(tutorial002.mcp) as web: