|
2 | 2 | import multiprocessing |
3 | 3 | import socket |
4 | 4 | import time |
5 | | -from collections.abc import AsyncGenerator, Generator |
| 5 | +from collections.abc import AsyncGenerator, Callable, Coroutine, Generator |
6 | 6 | from typing import Any |
7 | 7 | from unittest.mock import Mock |
8 | 8 |
|
|
24 | 24 | from mcp.server import Server |
25 | 25 | from mcp.server.sse import SseServerTransport |
26 | 26 | from mcp.server.transport_security import TransportSecuritySettings |
| 27 | +from mcp.shared._httpx_utils import McpHttpClientFactory |
| 28 | +from mcp.shared.context import RequestContext |
27 | 29 | from mcp.shared.exceptions import McpError |
28 | 30 | from mcp.shared.message import SessionMessage |
29 | 31 | from mcp.types import ( |
@@ -269,6 +271,127 @@ async def test_sse_client_exception_handling( |
269 | 271 | await session.read_resource(uri=AnyUrl("xxx://will-not-work")) |
270 | 272 |
|
271 | 273 |
|
| 274 | +class _FailingPostTransport(httpx.AsyncBaseTransport): |
| 275 | + """Passes requests on to the test server, except the message POSTs of one JSON-RPC method, which |
| 276 | + fail with `failure` (a status to answer with, or an exception to raise) the way a rate limiter, |
| 277 | + proxy or network in front of the server would, without reaching it.""" |
| 278 | + |
| 279 | + def __init__(self, method: str, failure: int | Exception) -> None: |
| 280 | + self._inner = httpx.AsyncHTTPTransport() |
| 281 | + self._method = method |
| 282 | + self._failure = failure |
| 283 | + |
| 284 | + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: |
| 285 | + if request.method == "POST" and json.loads(request.content).get("method") == self._method: |
| 286 | + if isinstance(self._failure, Exception): |
| 287 | + raise self._failure |
| 288 | + # A redirect names where to go, so the transport has a redirect to decide on. |
| 289 | + headers = {"location": "/elsewhere"} if 300 <= self._failure < 400 else None |
| 290 | + return httpx.Response(self._failure, headers=headers) |
| 291 | + return await self._inner.handle_async_request(request) |
| 292 | + |
| 293 | + async def aclose(self) -> None: |
| 294 | + await self._inner.aclose() |
| 295 | + |
| 296 | + |
| 297 | +def _failing_post_client_factory(method: str, failure: int | Exception) -> McpHttpClientFactory: |
| 298 | + def factory( |
| 299 | + headers: dict[str, str] | None = None, |
| 300 | + timeout: httpx.Timeout | None = None, |
| 301 | + auth: httpx.Auth | None = None, |
| 302 | + ) -> httpx.AsyncClient: |
| 303 | + return httpx.AsyncClient( |
| 304 | + transport=_FailingPostTransport(method, failure), headers=headers, timeout=timeout, auth=auth |
| 305 | + ) |
| 306 | + |
| 307 | + return factory |
| 308 | + |
| 309 | + |
| 310 | +@pytest.mark.anyio |
| 311 | +@pytest.mark.parametrize("status_code", [302, 401, 429, 500]) |
| 312 | +async def test_sse_client_request_post_http_error_reaches_caller_and_session_survives( |
| 313 | + server: None, server_url: str, status_code: int |
| 314 | +) -> None: |
| 315 | + """SDK-defined: a non-2xx answer to a request's message POST reaches the waiting caller promptly as a |
| 316 | + JSON-RPC error for that request, and the session stays usable. It used to be swallowed inside the |
| 317 | + transport, leaving the caller waiting forever and every later message unsent (#2110). An |
| 318 | + unfollowed redirect counts: the message never reached the server.""" |
| 319 | + factory = _failing_post_client_factory("resources/read", status_code) |
| 320 | + with anyio.fail_after(5): |
| 321 | + async with ( |
| 322 | + sse_client(server_url + "/sse", httpx_client_factory=factory) as streams, |
| 323 | + ClientSession(*streams) as session, |
| 324 | + ): |
| 325 | + await session.initialize() |
| 326 | + |
| 327 | + with pytest.raises(McpError) as exc_info: |
| 328 | + await session.read_resource(uri=AnyUrl("foobar://should-work")) |
| 329 | + assert exc_info.value.error.code == types.INTERNAL_ERROR |
| 330 | + assert exc_info.value.error.message == "Server returned an error response" |
| 331 | + |
| 332 | + # The session survived the failed POST: the next request round-trips. |
| 333 | + assert isinstance(await session.send_ping(), EmptyResult) |
| 334 | + |
| 335 | + |
| 336 | +@pytest.mark.anyio |
| 337 | +async def test_sse_client_request_post_404_reports_session_terminated(server: None, server_url: str) -> None: |
| 338 | + """SDK-defined: a 404 answer to a request's message POST, while the endpoint URL carries a session id, |
| 339 | + means the server no longer knows the session: the caller gets "Session terminated" |
| 340 | + (INVALID_REQUEST) and can reconnect.""" |
| 341 | + factory = _failing_post_client_factory("resources/read", 404) |
| 342 | + with anyio.fail_after(5): |
| 343 | + async with ( |
| 344 | + sse_client(server_url + "/sse", httpx_client_factory=factory) as streams, |
| 345 | + ClientSession(*streams) as session, |
| 346 | + ): |
| 347 | + await session.initialize() |
| 348 | + |
| 349 | + with pytest.raises(McpError) as exc_info: |
| 350 | + await session.read_resource(uri=AnyUrl("foobar://should-work")) |
| 351 | + assert exc_info.value.error.code == types.INVALID_REQUEST |
| 352 | + assert exc_info.value.error.message == "Session terminated" |
| 353 | + |
| 354 | + |
| 355 | +@pytest.mark.anyio |
| 356 | +async def test_sse_client_request_post_network_error_reaches_caller_and_session_survives( |
| 357 | + server: None, server_url: str |
| 358 | +) -> None: |
| 359 | + """SDK-defined: a request's message POST that fails without a response reaches the waiting caller |
| 360 | + promptly as a CONNECTION_CLOSED error for that request, and the session stays usable (#2110).""" |
| 361 | + factory = _failing_post_client_factory("resources/read", httpx.ConnectError("connection refused")) |
| 362 | + with anyio.fail_after(5): |
| 363 | + async with ( |
| 364 | + sse_client(server_url + "/sse", httpx_client_factory=factory) as streams, |
| 365 | + ClientSession(*streams) as session, |
| 366 | + ): |
| 367 | + await session.initialize() |
| 368 | + |
| 369 | + with pytest.raises(McpError) as exc_info: |
| 370 | + await session.read_resource(uri=AnyUrl("foobar://should-work")) |
| 371 | + assert exc_info.value.error.code == types.CONNECTION_CLOSED |
| 372 | + assert exc_info.value.error.message == "Failed to send message: connection refused" |
| 373 | + |
| 374 | + # The session survived the failed POST: the next request round-trips. |
| 375 | + assert isinstance(await session.send_ping(), EmptyResult) |
| 376 | + |
| 377 | + |
| 378 | +@pytest.mark.anyio |
| 379 | +async def test_sse_client_notification_post_failure_keeps_session_usable(server: None, server_url: str) -> None: |
| 380 | + """SDK-defined: a failed message POST of a notification has no waiter to resolve, so it is only |
| 381 | + logged, and the session stays usable (#2110: it used to leave every later message unsent).""" |
| 382 | + factory = _failing_post_client_factory("notifications/roots/list_changed", 500) |
| 383 | + with anyio.fail_after(5): |
| 384 | + async with ( |
| 385 | + sse_client(server_url + "/sse", httpx_client_factory=factory) as streams, |
| 386 | + ClientSession(*streams) as session, |
| 387 | + ): |
| 388 | + await session.initialize() |
| 389 | + |
| 390 | + await session.send_roots_list_changed() |
| 391 | + |
| 392 | + assert isinstance(await session.send_ping(), EmptyResult) |
| 393 | + |
| 394 | + |
272 | 395 | @pytest.mark.anyio |
273 | 396 | @pytest.mark.skip("this test highlights a possible bug in SSE read timeout exception handling") |
274 | 397 | async def test_sse_client_timeout( # pragma: no cover |
@@ -640,3 +763,128 @@ def factory( |
640 | 763 |
|
641 | 764 | assert exc_info.group_contains(httpx.HTTPStatusError, match="307 Temporary Redirect") |
642 | 765 | assert received == ["http://test/sse"] |
| 766 | + |
| 767 | + |
| 768 | +def _mock_client_factory( |
| 769 | + serve: Callable[[httpx.Request], Coroutine[None, None, httpx.Response]], |
| 770 | +) -> McpHttpClientFactory: |
| 771 | + def factory( |
| 772 | + headers: dict[str, str] | None = None, |
| 773 | + timeout: httpx.Timeout | None = None, |
| 774 | + auth: httpx.Auth | None = None, |
| 775 | + ) -> httpx.AsyncClient: |
| 776 | + return httpx.AsyncClient(transport=httpx.MockTransport(serve)) |
| 777 | + |
| 778 | + return factory |
| 779 | + |
| 780 | + |
| 781 | +@pytest.mark.anyio |
| 782 | +async def test_sse_client_post_404_without_session_endpoint_keeps_generic_error() -> None: |
| 783 | + """SDK-defined: a 404 answer to a request's message POST when the endpoint URL carries no session id |
| 784 | + keeps the generic error: with no session to expire, "Session terminated" would be wrong. The |
| 785 | + endpoint is scripted because `SseServerTransport` always issues a session id.""" |
| 786 | + |
| 787 | + async def event_stream() -> AsyncGenerator[bytes, None]: |
| 788 | + yield b"event: endpoint\ndata: /messages/\n\n" |
| 789 | + await anyio.sleep_forever() |
| 790 | + |
| 791 | + async def serve(request: httpx.Request) -> httpx.Response: |
| 792 | + if request.method == "GET": |
| 793 | + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=event_stream()) |
| 794 | + return httpx.Response(404) |
| 795 | + |
| 796 | + with anyio.fail_after(5): |
| 797 | + async with ( |
| 798 | + sse_client("http://test/sse", httpx_client_factory=_mock_client_factory(serve)) as streams, |
| 799 | + ClientSession(*streams) as session, |
| 800 | + ): |
| 801 | + with pytest.raises(McpError) as exc_info: |
| 802 | + await session.send_ping() |
| 803 | + assert exc_info.value.error.code == types.INTERNAL_ERROR |
| 804 | + assert exc_info.value.error.message == "Server returned an error response" |
| 805 | + |
| 806 | + |
| 807 | +@pytest.mark.anyio |
| 808 | +async def test_sse_client_failed_request_post_does_not_hold_up_reply_to_server_request() -> None: |
| 809 | + """SDK-defined: the error for a request's failed message POST is delivered without holding up the |
| 810 | + messages after it. ClientSession answers a server request inline in its receive loop, so its reply |
| 811 | + can be waiting on the transport exactly when that error needs the receive loop: delivering it from |
| 812 | + the loop that sends messages would leave the two waiting on each other forever.""" |
| 813 | + list_tools_posted = anyio.Event() |
| 814 | + roots_listed = anyio.Event() |
| 815 | + reply_posted = anyio.Event() |
| 816 | + |
| 817 | + async def event_stream() -> AsyncGenerator[bytes, None]: |
| 818 | + yield b"event: endpoint\ndata: /messages/?session_id=abc123\n\n" |
| 819 | + await list_tools_posted.wait() |
| 820 | + # A server request arrives while the client's tools/list POST is in flight. |
| 821 | + server_request = types.JSONRPCRequest(jsonrpc="2.0", id="server-1", method="roots/list") |
| 822 | + yield f"event: message\ndata: {server_request.model_dump_json()}\n\n".encode() |
| 823 | + await anyio.sleep_forever() |
| 824 | + |
| 825 | + async def serve(request: httpx.Request) -> httpx.Response: |
| 826 | + if request.method == "GET": |
| 827 | + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=event_stream()) |
| 828 | + message = json.loads(request.content) |
| 829 | + if message.get("method") == "tools/list": |
| 830 | + list_tools_posted.set() |
| 831 | + # Fail once the client is answering the server request: its reply then waits behind this POST. |
| 832 | + await roots_listed.wait() |
| 833 | + return httpx.Response(429) |
| 834 | + assert message == {"jsonrpc": "2.0", "id": "server-1", "result": {"roots": []}} |
| 835 | + reply_posted.set() |
| 836 | + return httpx.Response(202) |
| 837 | + |
| 838 | + async def list_roots(context: RequestContext[ClientSession, Any]) -> types.ListRootsResult: |
| 839 | + roots_listed.set() |
| 840 | + return types.ListRootsResult(roots=[]) |
| 841 | + |
| 842 | + with anyio.fail_after(5): |
| 843 | + async with ( |
| 844 | + sse_client("http://test/sse", httpx_client_factory=_mock_client_factory(serve)) as streams, |
| 845 | + ClientSession(*streams, list_roots_callback=list_roots) as session, |
| 846 | + ): |
| 847 | + with pytest.raises(McpError) as exc_info: |
| 848 | + await session.list_tools() |
| 849 | + assert exc_info.value.error.message == "Server returned an error response" |
| 850 | + |
| 851 | + await reply_posted.wait() |
| 852 | + |
| 853 | + |
| 854 | +@pytest.mark.anyio |
| 855 | +async def test_sse_client_post_error_after_reader_closed_is_contained() -> None: |
| 856 | + """SDK-defined: a failed request POST whose error can no longer be delivered, because the server's |
| 857 | + SSE stream already ended and closed the read stream, is contained: later messages are still sent |
| 858 | + and the transport closes cleanly. Raw streams, because a ClientSession stops sending once its read |
| 859 | + stream closes.""" |
| 860 | + stream_done = anyio.Event() |
| 861 | + release_post = anyio.Event() |
| 862 | + second_post = anyio.Event() |
| 863 | + |
| 864 | + async def event_stream() -> AsyncGenerator[bytes, None]: |
| 865 | + yield b"event: endpoint\ndata: /messages/?session_id=abc123\n\n" |
| 866 | + await stream_done.wait() |
| 867 | + |
| 868 | + async def serve(request: httpx.Request) -> httpx.Response: |
| 869 | + if request.method == "GET": |
| 870 | + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=event_stream()) |
| 871 | + if json.loads(request.content).get("method") == "ping": |
| 872 | + await release_post.wait() |
| 873 | + return httpx.Response(500) |
| 874 | + second_post.set() |
| 875 | + return httpx.Response(202) |
| 876 | + |
| 877 | + ping = types.JSONRPCRequest(jsonrpc="2.0", id=1, method="ping") |
| 878 | + initialized = types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") |
| 879 | + with anyio.fail_after(5): |
| 880 | + async with sse_client("http://test/sse", httpx_client_factory=_mock_client_factory(serve)) as ( |
| 881 | + read_stream, |
| 882 | + write_stream, |
| 883 | + ): |
| 884 | + await write_stream.send(SessionMessage(types.JSONRPCMessage(ping))) |
| 885 | + stream_done.set() |
| 886 | + assert [message async for message in read_stream] == [] |
| 887 | + |
| 888 | + release_post.set() |
| 889 | + await write_stream.send(SessionMessage(types.JSONRPCMessage(initialized))) |
| 890 | + await second_post.wait() |
0 commit comments