|
3 | 3 | import json |
4 | 4 | import logging |
5 | 5 | from collections.abc import Iterator |
6 | | -from typing import Any |
| 6 | +from typing import Any, Final |
7 | 7 | from unittest.mock import AsyncMock, patch |
8 | 8 |
|
9 | 9 | import anyio |
10 | 10 | import httpx2 |
11 | 11 | import pytest |
12 | 12 | from mcp_types import INVALID_REQUEST, ListToolsResult, PaginatedRequestParams |
| 13 | +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS |
13 | 14 | from starlette.types import Message, Receive, Scope, Send |
14 | 15 |
|
15 | 16 | from mcp import Client |
|
24 | 25 | StreamableHTTPSessionManager, |
25 | 26 | ) |
26 | 27 |
|
| 28 | +_INITIALIZE_BODY: Final[bytes] = json.dumps( |
| 29 | + { |
| 30 | + "jsonrpc": "2.0", |
| 31 | + "id": 1, |
| 32 | + "method": "initialize", |
| 33 | + "params": { |
| 34 | + "protocolVersion": HANDSHAKE_PROTOCOL_VERSIONS[-1], |
| 35 | + "capabilities": {}, |
| 36 | + "clientInfo": {"name": "test-client", "version": "1.0"}, |
| 37 | + }, |
| 38 | + } |
| 39 | +).encode() |
| 40 | + |
27 | 41 |
|
28 | 42 | @pytest.mark.anyio |
29 | 43 | async def test_run_can_only_be_called_once(): |
@@ -146,6 +160,67 @@ async def send(message: Message) -> None: |
146 | 160 | assert response_start["status"] == 413 |
147 | 161 |
|
148 | 162 |
|
| 163 | +@pytest.mark.anyio |
| 164 | +@pytest.mark.parametrize( |
| 165 | + ("method", "headers", "body", "expected_status"), |
| 166 | + [ |
| 167 | + pytest.param( |
| 168 | + "POST", |
| 169 | + [(b"content-type", b"application/json"), (b"accept", b"application/json, text/event-stream")], |
| 170 | + json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}}).encode(), |
| 171 | + 400, |
| 172 | + id="post-that-is-not-initialize", |
| 173 | + ), |
| 174 | + pytest.param( |
| 175 | + "POST", |
| 176 | + [(b"content-type", b"application/json"), (b"accept", b"application/json, text/event-stream")], |
| 177 | + b"{not valid json", |
| 178 | + 400, |
| 179 | + id="post-with-malformed-body", |
| 180 | + ), |
| 181 | + pytest.param( |
| 182 | + "POST", |
| 183 | + [(b"content-type", b"application/json"), (b"accept", b"text/plain")], |
| 184 | + _INITIALIZE_BODY, |
| 185 | + 406, |
| 186 | + id="post-with-unacceptable-accept", |
| 187 | + ), |
| 188 | + pytest.param("GET", [(b"accept", b"text/event-stream")], b"", 400, id="get-without-session"), |
| 189 | + pytest.param("DELETE", [], b"", 400, id="delete-without-session"), |
| 190 | + ], |
| 191 | +) |
| 192 | +async def test_refused_request_leaves_no_session_behind( |
| 193 | + method: str, headers: list[tuple[bytes, bytes]], body: bytes, expected_status: int |
| 194 | +) -> None: |
| 195 | + """SDK-defined: a request the transport refuses must not leave a registered session behind. |
| 196 | +
|
| 197 | + The session is minted before any validation runs -- Host, Accept, Content-Type, JSON parse, |
| 198 | + JSON-RPC shape and the "Missing session ID" check all live downstream in the transport -- so a |
| 199 | + refusal has to undo it. Otherwise a rejected request grows `_server_instances` forever and hands |
| 200 | + the caller a session id that later requests can still use. |
| 201 | +
|
| 202 | + This is the same property the suite already asserts by name for the 413 path in |
| 203 | + `test_oversized_content_length_is_rejected_before_body_read_or_session_creation`. |
| 204 | + """ |
| 205 | + manager = StreamableHTTPSessionManager(app=Server("test-refused-request")) |
| 206 | + sent_messages: list[Message] = [] |
| 207 | + |
| 208 | + async def mock_send(message: Message) -> None: |
| 209 | + sent_messages.append(message) |
| 210 | + |
| 211 | + async def mock_receive() -> Message: |
| 212 | + return {"type": "http.request", "body": body, "more_body": False} |
| 213 | + |
| 214 | + scope: Scope = {"type": "http", "method": method, "path": "/mcp", "headers": headers} |
| 215 | + |
| 216 | + async with manager.run(): |
| 217 | + await manager.handle_request(scope, mock_receive, mock_send) |
| 218 | + |
| 219 | + response_start = next(msg for msg in sent_messages if msg["type"] == "http.response.start") |
| 220 | + assert response_start["status"] == expected_status |
| 221 | + assert manager._server_instances == {} |
| 222 | + |
| 223 | + |
149 | 224 | @pytest.mark.anyio |
150 | 225 | async def test_client_disconnect_while_streaming_request_body_is_replayed() -> None: |
151 | 226 | """SDK-defined: raw ASGI is required to prove a disconnect before body completion reaches the transport.""" |
@@ -513,35 +588,13 @@ async def test_idle_session_is_reaped(caplog: pytest.LogCaptureFixture, request: |
513 | 588 | caplog.set_level(logging.INFO, logger=streamable_http_manager.__name__) |
514 | 589 |
|
515 | 590 | async with manager.run(): |
516 | | - sent_messages: list[Message] = [] |
517 | | - |
518 | | - async def mock_send(message: Message): |
519 | | - sent_messages.append(message) |
| 591 | + # Establish the session with a real `initialize`: a request the transport refuses no |
| 592 | + # longer leaves a session behind, so there would be nothing for the reaper to reap. |
| 593 | + session_id = await _open_session(manager, None) |
520 | 594 |
|
521 | | - scope = { |
522 | | - "type": "http", |
523 | | - "method": "POST", |
524 | | - "path": "/mcp", |
525 | | - "headers": [(b"content-type", b"application/json")], |
526 | | - } |
527 | | - |
528 | | - async def mock_receive(): |
| 595 | + async def mock_receive() -> Message: |
529 | 596 | return {"type": "http.request", "body": b"", "more_body": False} |
530 | 597 |
|
531 | | - await manager.handle_request(scope, mock_receive, mock_send) |
532 | | - |
533 | | - session_id = None |
534 | | - for msg in sent_messages: # pragma: no branch |
535 | | - if msg["type"] == "http.response.start": # pragma: no branch |
536 | | - for header_name, header_value in msg.get("headers", []): # pragma: no branch |
537 | | - if header_name.decode().lower() == MCP_SESSION_ID_HEADER.lower(): |
538 | | - session_id = header_value.decode() |
539 | | - break |
540 | | - if session_id: # pragma: no branch |
541 | | - break |
542 | | - |
543 | | - assert session_id is not None, "Session ID not found in response headers" |
544 | | - |
545 | 598 | # Wait for the 50ms idle timeout to fire and the session to be unregistered. Re-requesting |
546 | 599 | # the session to poll for the 404 would push its idle deadline forward and keep it alive. |
547 | 600 | with anyio.fail_after(5): |
@@ -613,18 +666,31 @@ def _request_scope( |
613 | 666 |
|
614 | 667 |
|
615 | 668 | async def _open_session(manager: StreamableHTTPSessionManager, user: AuthenticatedUser | None) -> str: |
616 | | - """Create a new session as `user` and return its session ID.""" |
| 669 | + """Create a new session as `user` and return its session ID. |
| 670 | +
|
| 671 | + Establishes the session the way a real client does, with an `initialize` request, because |
| 672 | + a request the transport refuses no longer leaves a session behind. The reply is an SSE |
| 673 | + stream, so the body is followed by a disconnect: that ends the stream and lets |
| 674 | + `handle_request` return, while the session itself lives on in the manager's task group. |
| 675 | + """ |
617 | 676 | sent_messages: list[Message] = [] |
| 677 | + body_sent = False |
618 | 678 |
|
619 | 679 | async def mock_send(message: Message) -> None: |
620 | 680 | sent_messages.append(message) |
621 | 681 |
|
622 | 682 | async def mock_receive() -> Message: |
623 | | - return {"type": "http.request", "body": b"", "more_body": False} |
| 683 | + nonlocal body_sent |
| 684 | + if body_sent: |
| 685 | + return {"type": "http.disconnect"} |
| 686 | + body_sent = True |
| 687 | + return {"type": "http.request", "body": _INITIALIZE_BODY, "more_body": False} |
624 | 688 |
|
625 | | - await manager.handle_request(_request_scope(user=user), mock_receive, mock_send) |
| 689 | + with anyio.fail_after(5): |
| 690 | + await manager.handle_request(_request_scope(user=user), mock_receive, mock_send) |
626 | 691 |
|
627 | 692 | response_start = next(msg for msg in sent_messages if msg["type"] == "http.response.start") |
| 693 | + assert response_start["status"] == 200, f"initialize was refused with {response_start['status']}" |
628 | 694 | headers = dict(response_start.get("headers", [])) |
629 | 695 | return headers[MCP_SESSION_ID_HEADER.encode()].decode() |
630 | 696 |
|
|
0 commit comments