|
2 | 2 | import pytest
|
3 | 3 |
|
4 | 4 | import mcp.types as types
|
5 |
| -from mcp.client.session import ClientSession |
| 5 | +from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession |
6 | 6 | from mcp.shared.session import RequestResponder
|
7 | 7 | from mcp.types import (
|
8 | 8 | LATEST_PROTOCOL_VERSION,
|
@@ -111,3 +111,131 @@ async def message_handler(
|
111 | 111 | # Check that the client sent the initialized notification
|
112 | 112 | assert initialized_notification
|
113 | 113 | assert isinstance(initialized_notification.root, InitializedNotification)
|
| 114 | + |
| 115 | + |
| 116 | +@pytest.mark.anyio |
| 117 | +async def test_client_session_custom_client_info(): |
| 118 | + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[ |
| 119 | + JSONRPCMessage |
| 120 | + ](1) |
| 121 | + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[ |
| 122 | + JSONRPCMessage |
| 123 | + ](1) |
| 124 | + |
| 125 | + custom_client_info = Implementation(name="test-client", version="1.2.3") |
| 126 | + received_client_info = None |
| 127 | + |
| 128 | + async def mock_server(): |
| 129 | + nonlocal received_client_info |
| 130 | + |
| 131 | + jsonrpc_request = await client_to_server_receive.receive() |
| 132 | + assert isinstance(jsonrpc_request.root, JSONRPCRequest) |
| 133 | + request = ClientRequest.model_validate( |
| 134 | + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) |
| 135 | + ) |
| 136 | + assert isinstance(request.root, InitializeRequest) |
| 137 | + received_client_info = request.root.params.clientInfo |
| 138 | + |
| 139 | + result = ServerResult( |
| 140 | + InitializeResult( |
| 141 | + protocolVersion=LATEST_PROTOCOL_VERSION, |
| 142 | + capabilities=ServerCapabilities(), |
| 143 | + serverInfo=Implementation(name="mock-server", version="0.1.0"), |
| 144 | + ) |
| 145 | + ) |
| 146 | + |
| 147 | + async with server_to_client_send: |
| 148 | + await server_to_client_send.send( |
| 149 | + JSONRPCMessage( |
| 150 | + JSONRPCResponse( |
| 151 | + jsonrpc="2.0", |
| 152 | + id=jsonrpc_request.root.id, |
| 153 | + result=result.model_dump( |
| 154 | + by_alias=True, mode="json", exclude_none=True |
| 155 | + ), |
| 156 | + ) |
| 157 | + ) |
| 158 | + ) |
| 159 | + # Receive initialized notification |
| 160 | + await client_to_server_receive.receive() |
| 161 | + |
| 162 | + async with ( |
| 163 | + ClientSession( |
| 164 | + server_to_client_receive, |
| 165 | + client_to_server_send, |
| 166 | + client_info=custom_client_info, |
| 167 | + ) as session, |
| 168 | + anyio.create_task_group() as tg, |
| 169 | + client_to_server_send, |
| 170 | + client_to_server_receive, |
| 171 | + server_to_client_send, |
| 172 | + server_to_client_receive, |
| 173 | + ): |
| 174 | + tg.start_soon(mock_server) |
| 175 | + await session.initialize() |
| 176 | + |
| 177 | + # Assert that the custom client info was sent |
| 178 | + assert received_client_info == custom_client_info |
| 179 | + |
| 180 | + |
| 181 | +@pytest.mark.anyio |
| 182 | +async def test_client_session_default_client_info(): |
| 183 | + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[ |
| 184 | + JSONRPCMessage |
| 185 | + ](1) |
| 186 | + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[ |
| 187 | + JSONRPCMessage |
| 188 | + ](1) |
| 189 | + |
| 190 | + received_client_info = None |
| 191 | + |
| 192 | + async def mock_server(): |
| 193 | + nonlocal received_client_info |
| 194 | + |
| 195 | + jsonrpc_request = await client_to_server_receive.receive() |
| 196 | + assert isinstance(jsonrpc_request.root, JSONRPCRequest) |
| 197 | + request = ClientRequest.model_validate( |
| 198 | + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) |
| 199 | + ) |
| 200 | + assert isinstance(request.root, InitializeRequest) |
| 201 | + received_client_info = request.root.params.clientInfo |
| 202 | + |
| 203 | + result = ServerResult( |
| 204 | + InitializeResult( |
| 205 | + protocolVersion=LATEST_PROTOCOL_VERSION, |
| 206 | + capabilities=ServerCapabilities(), |
| 207 | + serverInfo=Implementation(name="mock-server", version="0.1.0"), |
| 208 | + ) |
| 209 | + ) |
| 210 | + |
| 211 | + async with server_to_client_send: |
| 212 | + await server_to_client_send.send( |
| 213 | + JSONRPCMessage( |
| 214 | + JSONRPCResponse( |
| 215 | + jsonrpc="2.0", |
| 216 | + id=jsonrpc_request.root.id, |
| 217 | + result=result.model_dump( |
| 218 | + by_alias=True, mode="json", exclude_none=True |
| 219 | + ), |
| 220 | + ) |
| 221 | + ) |
| 222 | + ) |
| 223 | + # Receive initialized notification |
| 224 | + await client_to_server_receive.receive() |
| 225 | + |
| 226 | + async with ( |
| 227 | + ClientSession( |
| 228 | + server_to_client_receive, |
| 229 | + client_to_server_send, |
| 230 | + ) as session, |
| 231 | + anyio.create_task_group() as tg, |
| 232 | + client_to_server_send, |
| 233 | + client_to_server_receive, |
| 234 | + server_to_client_send, |
| 235 | + server_to_client_receive, |
| 236 | + ): |
| 237 | + tg.start_soon(mock_server) |
| 238 | + await session.initialize() |
| 239 | + |
| 240 | + # Assert that the default client info was sent |
| 241 | + assert received_client_info == DEFAULT_CLIENT_INFO |
0 commit comments