|
1 | 1 | import contextlib |
| 2 | +from collections.abc import Generator |
| 3 | +from typing import Any |
2 | 4 | from unittest import mock |
3 | 5 |
|
| 6 | +import anyio |
4 | 7 | import httpx2 |
5 | 8 | import mcp_types as types |
6 | 9 | import pytest |
| 10 | +from starlette.types import ASGIApp, Receive, Scope, Send |
7 | 11 |
|
8 | 12 | import mcp |
| 13 | +import mcp.shared._httpx_utils as httpx_utils |
9 | 14 | from mcp.client.session_group import ( |
10 | 15 | ClientSessionGroup, |
11 | 16 | ClientSessionParameters, |
12 | 17 | SseServerParameters, |
13 | 18 | StreamableHttpParameters, |
14 | 19 | ) |
15 | 20 | from mcp.client.stdio import StdioServerParameters |
| 21 | +from mcp.server import MCPServer |
| 22 | +from mcp.server.transport_security import TransportSecuritySettings |
16 | 23 | from mcp.shared.exceptions import MCPError |
| 24 | +from tests.interaction._connect import BASE_URL, build_sse_app |
| 25 | +from tests.interaction.transports import StreamingASGITransport |
17 | 26 |
|
18 | 27 |
|
19 | 28 | @pytest.fixture |
@@ -374,6 +383,7 @@ async def test_client_session_group_establish_session_parameterized( |
374 | 383 | headers=server_params_instance.headers, |
375 | 384 | timeout=server_params_instance.timeout, |
376 | 385 | sse_read_timeout=server_params_instance.sse_read_timeout, |
| 386 | + auth=server_params_instance.auth, |
377 | 387 | ) |
378 | 388 | elif client_type_name == "streamablehttp": # pragma: no branch |
379 | 389 | assert isinstance(server_params_instance, StreamableHttpParameters) |
@@ -405,3 +415,183 @@ async def test_client_session_group_establish_session_parameterized( |
405 | 415 | # 3. Assert returned values |
406 | 416 | assert returned_server_info is mock_initialize_result.server_info |
407 | 417 | 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"} |
0 commit comments