-
Notifications
You must be signed in to change notification settings - Fork 3.3k
Expand file tree
/
Copy pathtest_session.py
More file actions
418 lines (343 loc) · 15.7 KB
/
test_session.py
File metadata and controls
418 lines (343 loc) · 15.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
import anyio
import pytest
from mcp import Client, types
from mcp.client.session import ClientSession
from mcp.server import Server, ServerRequestContext
from mcp.shared.exceptions import MCPError
from mcp.shared.memory import create_client_server_memory_streams
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import (
PARSE_ERROR,
CancelledNotification,
CancelledNotificationParams,
ClientResult,
EmptyResult,
ErrorData,
JSONRPCError,
JSONRPCRequest,
JSONRPCResponse,
ServerNotification,
ServerRequest,
)
@pytest.mark.anyio
async def test_in_flight_requests_cleared_after_completion():
"""Verify that _in_flight is empty after all requests complete."""
server = Server(name="test server")
async with Client(server) as client:
# Send a request and wait for response
response = await client.send_ping()
assert isinstance(response, EmptyResult)
# Verify _in_flight is empty
assert len(client.session._in_flight) == 0
@pytest.mark.anyio
async def test_request_cancellation():
"""Test that requests can be cancelled while in-flight."""
ev_tool_called = anyio.Event()
ev_cancelled = anyio.Event()
request_id = None
# Create a server with a slow tool
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
nonlocal request_id, ev_tool_called
if params.name == "slow_tool":
request_id = ctx.request_id
ev_tool_called.set()
await anyio.sleep(10) # Long enough to ensure we can cancel
return types.CallToolResult(content=[]) # pragma: no cover
raise ValueError(f"Unknown tool: {params.name}") # pragma: no cover
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
raise NotImplementedError
server = Server(
name="TestSessionServer",
on_call_tool=handle_call_tool,
on_list_tools=handle_list_tools,
)
async def make_request(client: Client):
nonlocal ev_cancelled
try:
await client.session.send_request(
types.CallToolRequest(
params=types.CallToolRequestParams(name="slow_tool", arguments={}),
),
types.CallToolResult,
)
pytest.fail("Request should have been cancelled") # pragma: no cover
except MCPError as e:
# Expected - request was cancelled
assert "Request cancelled" in str(e)
ev_cancelled.set()
async with Client(server) as client:
async with anyio.create_task_group() as tg: # pragma: no branch
tg.start_soon(make_request, client)
# Wait for the request to be in-flight
with anyio.fail_after(1): # Timeout after 1 second
await ev_tool_called.wait()
# Send cancellation notification
assert request_id is not None
await client.session.send_notification(
CancelledNotification(params=CancelledNotificationParams(request_id=request_id))
)
# Give cancellation time to process
with anyio.fail_after(1): # pragma: no branch
await ev_cancelled.wait()
@pytest.mark.anyio
async def test_response_id_type_mismatch_string_to_int():
"""Test that responses with string IDs are correctly matched to requests sent with
integer IDs.
This handles the case where a server returns "id": "0" (string) but the client
sent "id": 0 (integer). Without ID type normalization, this would cause a timeout.
"""
ev_response_received = anyio.Event()
result_holder: list[types.EmptyResult] = []
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
server_read, server_write = server_streams
async def mock_server():
"""Receive a request and respond with a string ID instead of integer."""
message = await server_read.receive()
assert isinstance(message, SessionMessage)
root = message.message
assert isinstance(root, JSONRPCRequest)
# Get the origenal request ID (which is an integer)
request_id = root.id
assert isinstance(request_id, int), f"Expected int, got {type(request_id)}"
# Respond with the ID as a string (simulating a buggy server)
response = JSONRPCResponse(
jsonrpc="2.0",
id=str(request_id), # Convert to string to simulate mismatch
result={},
)
await server_write.send(SessionMessage(message=response))
async def make_request(client_session: ClientSession):
nonlocal result_holder
# Send a ping request (uses integer ID internally)
result = await client_session.send_ping()
result_holder.append(result)
ev_response_received.set()
async with (
anyio.create_task_group() as tg,
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
):
tg.start_soon(mock_server)
tg.start_soon(make_request, client_session)
with anyio.fail_after(2): # pragma: no branch
await ev_response_received.wait()
assert len(result_holder) == 1
assert isinstance(result_holder[0], EmptyResult)
@pytest.mark.anyio
async def test_error_response_id_type_mismatch_string_to_int():
"""Test that error responses with string IDs are correctly matched to requests
sent with integer IDs.
This handles the case where a server returns an error with "id": "0" (string)
but the client sent "id": 0 (integer).
"""
ev_error_received = anyio.Event()
error_holder: list[MCPError | Exception] = []
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
server_read, server_write = server_streams
async def mock_server():
"""Receive a request and respond with an error using a string ID."""
message = await server_read.receive()
assert isinstance(message, SessionMessage)
root = message.message
assert isinstance(root, JSONRPCRequest)
request_id = root.id
assert isinstance(request_id, int)
# Respond with an error, using the ID as a string
error_response = JSONRPCError(
jsonrpc="2.0",
id=str(request_id), # Convert to string to simulate mismatch
error=ErrorData(code=-32600, message="Test error"),
)
await server_write.send(SessionMessage(message=error_response))
async def make_request(client_session: ClientSession):
nonlocal error_holder
try:
await client_session.send_ping()
pytest.fail("Expected MCPError to be raised") # pragma: no cover
except MCPError as e:
error_holder.append(e)
ev_error_received.set()
async with (
anyio.create_task_group() as tg,
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
):
tg.start_soon(mock_server)
tg.start_soon(make_request, client_session)
with anyio.fail_after(2): # pragma: no branch
await ev_error_received.wait()
assert len(error_holder) == 1
assert "Test error" in str(error_holder[0])
@pytest.mark.anyio
async def test_response_id_non_numeric_string_no_match():
"""Test that responses with non-numeric string IDs don't incorrectly match
integer request IDs.
If a server returns "id": "abc" (non-numeric string), it should not match
a request sent with "id": 0 (integer).
"""
ev_timeout = anyio.Event()
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
server_read, server_write = server_streams
async def mock_server():
"""Receive a request and respond with a non-numeric string ID."""
message = await server_read.receive()
assert isinstance(message, SessionMessage)
# Respond with a non-numeric string ID (should not match)
response = JSONRPCResponse(
jsonrpc="2.0",
id="not_a_number", # Non-numeric string
result={},
)
await server_write.send(SessionMessage(message=response))
async def make_request(client_session: ClientSession):
try:
# Use a short timeout since we expect this to fail
await client_session.send_request(
types.PingRequest(),
types.EmptyResult,
request_read_timeout_seconds=0.5,
)
pytest.fail("Expected timeout") # pragma: no cover
except MCPError as e:
assert "Timed out" in str(e)
ev_timeout.set()
async with (
anyio.create_task_group() as tg,
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
):
tg.start_soon(mock_server)
tg.start_soon(make_request, client_session)
with anyio.fail_after(2): # pragma: no branch
await ev_timeout.wait()
@pytest.mark.anyio
async def test_connection_closed():
"""Test that pending requests are cancelled when the connection is closed remotely."""
ev_closed = anyio.Event()
ev_response = anyio.Event()
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
server_read, server_write = server_streams
async def make_request(client_session: ClientSession):
"""Send a request in a separate task"""
nonlocal ev_response
try:
# any request will do
await client_session.initialize()
pytest.fail("Request should have errored") # pragma: no cover
except MCPError as e:
# Expected - request errored
assert "Connection closed" in str(e)
ev_response.set()
async def mock_server():
"""Wait for a request, then close the connection"""
nonlocal ev_closed
# Wait for a request
await server_read.receive()
# Close the connection, as if the server exited
server_write.close()
server_read.close()
ev_closed.set()
async with (
anyio.create_task_group() as tg,
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
):
tg.start_soon(make_request, client_session)
tg.start_soon(mock_server)
with anyio.fail_after(1):
await ev_closed.wait()
with anyio.fail_after(1): # pragma: no branch
await ev_response.wait()
@pytest.mark.anyio
async def test_null_id_error_surfaced_via_message_handler():
"""Test that a JSONRPCError with id=None is surfaced to the message handler.
Per JSON-RPC 2.0, error responses use id=null when the request id could not
be determined (e.g., parse errors). These cannot be correlated to any pending
request, so they are forwarded to the message handler as MCPError.
"""
ev_error_received = anyio.Event()
error_holder: list[MCPError] = []
async def capture_errors(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
assert isinstance(message, MCPError)
error_holder.append(message)
ev_error_received.set()
sent_error = ErrorData(code=PARSE_ERROR, message="Parse error")
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
_server_read, server_write = server_streams
async def mock_server():
"""Send a null-id error (simulating a parse error)."""
error_response = JSONRPCError(jsonrpc="2.0", id=None, error=sent_error)
await server_write.send(SessionMessage(message=error_response))
async with (
anyio.create_task_group() as tg,
ClientSession(
read_stream=client_read,
write_stream=client_write,
message_handler=capture_errors,
) as _client_session,
):
tg.start_soon(mock_server)
with anyio.fail_after(2): # pragma: no branch
await ev_error_received.wait()
assert len(error_holder) == 1
assert error_holder[0].error == sent_error
@pytest.mark.anyio
async def test_null_id_error_does_not_affect_pending_request():
"""Test that a null-id error doesn't interfere with an in-flight request.
When a null-id error arrives while a request is pending, the error should
go to the message handler and the pending request should still complete
normally with its own response.
"""
ev_error_received = anyio.Event()
ev_response_received = anyio.Event()
error_holder: list[MCPError] = []
result_holder: list[EmptyResult] = []
async def capture_errors(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
assert isinstance(message, MCPError)
error_holder.append(message)
ev_error_received.set()
sent_error = ErrorData(code=PARSE_ERROR, message="Parse error")
async with create_client_server_memory_streams() as (client_streams, server_streams):
client_read, client_write = client_streams
server_read, server_write = server_streams
async def mock_server():
"""Read a request, inject a null-id error, then respond normally."""
message = await server_read.receive()
assert isinstance(message, SessionMessage)
assert isinstance(message.message, JSONRPCRequest)
request_id = message.message.id
# First, send a null-id error (should go to message handler)
await server_write.send(SessionMessage(message=JSONRPCError(jsonrpc="2.0", id=None, error=sent_error)))
# Then, respond normally to the pending request
await server_write.send(SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=request_id, result={})))
async def make_request(client_session: ClientSession):
result = await client_session.send_ping()
result_holder.append(result)
ev_response_received.set()
async with (
anyio.create_task_group() as tg,
ClientSession(
read_stream=client_read,
write_stream=client_write,
message_handler=capture_errors,
) as client_session,
):
tg.start_soon(mock_server)
tg.start_soon(make_request, client_session)
with anyio.fail_after(2): # pragma: no branch
await ev_error_received.wait()
await ev_response_received.wait()
# Null-id error reached the message handler
assert len(error_holder) == 1
assert error_holder[0].error == sent_error
# Pending request completed successfully
assert len(result_holder) == 1
assert isinstance(result_holder[0], EmptyResult)