|
| 1 | +""" |
| 2 | +Test for race condition fix in initialization flow. |
| 3 | +
|
| 4 | +This test verifies that requests can be processed immediately after |
| 5 | +responding to InitializeRequest, without waiting for InitializedNotification. |
| 6 | +
|
| 7 | +This is critical for HTTP transport where requests can arrive in any order. |
| 8 | +""" |
| 9 | + |
| 10 | +import anyio |
| 11 | +import pytest |
| 12 | + |
| 13 | +import mcp.types as types |
| 14 | +from mcp.server.models import InitializationOptions |
| 15 | +from mcp.server.session import ServerSession |
| 16 | +from mcp.shared.message import SessionMessage |
| 17 | +from mcp.shared.session import RequestResponder |
| 18 | +from mcp.types import ServerCapabilities, Tool |
| 19 | + |
| 20 | + |
| 21 | +@pytest.mark.anyio |
| 22 | +async def test_request_immediately_after_initialize_response(): |
| 23 | + """ |
| 24 | + Test that requests are accepted immediately after initialize response. |
| 25 | +
|
| 26 | + This reproduces the race condition in stateful HTTP mode where: |
| 27 | + 1. Client sends InitializeRequest |
| 28 | + 2. Server responds with InitializeResult |
| 29 | + 3. Client immediately sends tools/list (before server receives InitializedNotification) |
| 30 | + 4. Without fix: Server rejects with "Received request before initialization was complete" |
| 31 | + 5. With fix: Server accepts and processes the request |
| 32 | +
|
| 33 | + This test simulates the HTTP transport behavior where InitializedNotification |
| 34 | + may arrive in a separate POST request after other requests. |
| 35 | + """ |
| 36 | + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) |
| 37 | + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](10) |
| 38 | + |
| 39 | + tools_list_success = False |
| 40 | + error_received = None |
| 41 | + |
| 42 | + async def run_server(): |
| 43 | + nonlocal tools_list_success |
| 44 | + |
| 45 | + async with ServerSession( |
| 46 | + client_to_server_receive, |
| 47 | + server_to_client_send, |
| 48 | + InitializationOptions( |
| 49 | + server_name="test-server", |
| 50 | + server_version="1.0.0", |
| 51 | + capabilities=ServerCapabilities( |
| 52 | + tools=types.ToolsCapability(listChanged=False), |
| 53 | + ), |
| 54 | + ), |
| 55 | + ) as server_session: |
| 56 | + async for message in server_session.incoming_messages: |
| 57 | + if isinstance(message, Exception): |
| 58 | + raise message |
| 59 | + |
| 60 | + # Handle tools/list request |
| 61 | + if isinstance(message, RequestResponder): |
| 62 | + if isinstance(message.request.root, types.ListToolsRequest): |
| 63 | + tools_list_success = True |
| 64 | + # Respond with a tool list |
| 65 | + with message: |
| 66 | + await message.respond( |
| 67 | + types.ServerResult( |
| 68 | + types.ListToolsResult( |
| 69 | + tools=[ |
| 70 | + Tool( |
| 71 | + name="example_tool", |
| 72 | + description="An example tool", |
| 73 | + inputSchema={"type": "object", "properties": {}}, |
| 74 | + ) |
| 75 | + ] |
| 76 | + ) |
| 77 | + ) |
| 78 | + ) |
| 79 | + |
| 80 | + # Handle InitializedNotification |
| 81 | + if isinstance(message, types.ClientNotification): |
| 82 | + if isinstance(message.root, types.InitializedNotification): |
| 83 | + # Done - exit gracefully |
| 84 | + return |
| 85 | + |
| 86 | + async def mock_client(): |
| 87 | + nonlocal error_received |
| 88 | + |
| 89 | + # Step 1: Send InitializeRequest |
| 90 | + await client_to_server_send.send( |
| 91 | + SessionMessage( |
| 92 | + types.JSONRPCMessage( |
| 93 | + types.JSONRPCRequest( |
| 94 | + jsonrpc="2.0", |
| 95 | + id=1, |
| 96 | + method="initialize", |
| 97 | + params=types.InitializeRequestParams( |
| 98 | + protocolVersion=types.LATEST_PROTOCOL_VERSION, |
| 99 | + capabilities=types.ClientCapabilities(), |
| 100 | + clientInfo=types.Implementation(name="test-client", version="1.0.0"), |
| 101 | + ).model_dump(by_alias=True, mode="json", exclude_none=True), |
| 102 | + ) |
| 103 | + ) |
| 104 | + ) |
| 105 | + ) |
| 106 | + |
| 107 | + # Step 2: Wait for InitializeResult |
| 108 | + init_msg = await server_to_client_receive.receive() |
| 109 | + assert isinstance(init_msg.message.root, types.JSONRPCResponse) |
| 110 | + |
| 111 | + # Step 3: Immediately send tools/list BEFORE InitializedNotification |
| 112 | + # This is the race condition scenario |
| 113 | + await client_to_server_send.send( |
| 114 | + SessionMessage( |
| 115 | + types.JSONRPCMessage( |
| 116 | + types.JSONRPCRequest( |
| 117 | + jsonrpc="2.0", |
| 118 | + id=2, |
| 119 | + method="tools/list", |
| 120 | + ) |
| 121 | + ) |
| 122 | + ) |
| 123 | + ) |
| 124 | + |
| 125 | + # Step 4: Check the response |
| 126 | + tools_msg = await server_to_client_receive.receive() |
| 127 | + if isinstance(tools_msg.message.root, types.JSONRPCError): |
| 128 | + error_received = tools_msg.message.root.error.message |
| 129 | + |
| 130 | + # Step 5: Send InitializedNotification |
| 131 | + await client_to_server_send.send( |
| 132 | + SessionMessage( |
| 133 | + types.JSONRPCMessage( |
| 134 | + types.JSONRPCNotification( |
| 135 | + jsonrpc="2.0", |
| 136 | + method="notifications/initialized", |
| 137 | + ) |
| 138 | + ) |
| 139 | + ) |
| 140 | + ) |
| 141 | + |
| 142 | + async with ( |
| 143 | + client_to_server_send, |
| 144 | + client_to_server_receive, |
| 145 | + server_to_client_send, |
| 146 | + server_to_client_receive, |
| 147 | + anyio.create_task_group() as tg, |
| 148 | + ): |
| 149 | + tg.start_soon(run_server) |
| 150 | + tg.start_soon(mock_client) |
| 151 | + |
| 152 | + # With the PR fix: tools_list_success should be True, error_received should be None |
| 153 | + # Without the fix: error_received would contain "Received request before initialization was complete" |
| 154 | + assert tools_list_success, f"tools/list should have succeeded. Error received: {error_received}" |
| 155 | + assert error_received is None, f"Expected no error, but got: {error_received}" |
0 commit comments