|
6 | 6 | import pytest |
7 | 7 |
|
8 | 8 | import mcp |
| 9 | +from mcp.client import Client |
9 | 10 | from mcp.client.session_group import ( |
10 | 11 | ClientSessionGroup, |
11 | 12 | ClientSessionParameters, |
12 | 13 | SseServerParameters, |
13 | 14 | StreamableHttpParameters, |
14 | 15 | ) |
15 | 16 | from mcp.client.stdio import StdioServerParameters |
| 17 | +from mcp.server.mcpserver import MCPServer |
16 | 18 | from mcp.shared.exceptions import MCPError |
17 | 19 |
|
18 | 20 |
|
@@ -405,3 +407,93 @@ async def test_client_session_group_establish_session_parameterized( |
405 | 407 | # 3. Assert returned values |
406 | 408 | assert returned_server_info is mock_initialize_result.server_info |
407 | 409 | assert returned_session is mock_entered_session |
| 410 | + |
| 411 | + |
| 412 | +@pytest.mark.anyio |
| 413 | +async def test_connect_with_session_accepts_server_with_no_components(): |
| 414 | + """A server advertising no prompts, resources or tools still joins the group. |
| 415 | +
|
| 416 | + Regression test for #3384: aggregation used to delete an exit stack that |
| 417 | + `connect_with_session` never registered, raising a bare `KeyError`. |
| 418 | + """ |
| 419 | + server = MCPServer("no-components") |
| 420 | + server_info = types.Implementation(name="no-components", version="1.0") |
| 421 | + |
| 422 | + async with Client(server) as client: |
| 423 | + async with ClientSessionGroup() as group: |
| 424 | + await group.connect_with_session(server_info, client.session) |
| 425 | + |
| 426 | + assert client.session in group._sessions |
| 427 | + assert group.tools == {} |
| 428 | + assert group.resources == {} |
| 429 | + assert group.prompts == {} |
| 430 | + |
| 431 | + |
| 432 | +@pytest.mark.anyio |
| 433 | +async def test_connect_with_session_accepts_server_exposing_only_resource_templates(): |
| 434 | + """A server whose entire surface is a resource template still joins the group. |
| 435 | +
|
| 436 | + Regression test for #3384. The group never calls `list_resource_templates`, |
| 437 | + so a templates-only server looked componentless to aggregation even though |
| 438 | + it advertises a real surface — `list_resources` is empty while |
| 439 | + `list_resource_templates` is not. |
| 440 | + """ |
| 441 | + server = MCPServer("templates-only") |
| 442 | + |
| 443 | + @server.resource("doc://{doc_id}") |
| 444 | + def get_doc(doc_id: str) -> str: |
| 445 | + return f"contents of {doc_id}" |
| 446 | + |
| 447 | + server_info = types.Implementation(name="templates-only", version="1.0") |
| 448 | + |
| 449 | + async with Client(server) as client: |
| 450 | + listed = await client.list_resource_templates() |
| 451 | + assert [t.uri_template for t in listed.resource_templates] == ["doc://{doc_id}"] |
| 452 | + assert (await client.list_resources()).resources == [] |
| 453 | + |
| 454 | + async with ClientSessionGroup() as group: |
| 455 | + await group.connect_with_session(server_info, client.session) |
| 456 | + |
| 457 | + assert client.session in group._sessions |
| 458 | + |
| 459 | + |
| 460 | +@pytest.mark.anyio |
| 461 | +async def test_disconnect_from_server_closes_transport_of_componentless_server(): |
| 462 | + """Disconnecting a componentless server closes the transport its connection opened. |
| 463 | +
|
| 464 | + Regression test for #3384: aggregation dropped the session's exit stack |
| 465 | + without closing it, so `disconnect_from_server` had nothing left to tear |
| 466 | + down and the transport stayed open until the whole group exited. |
| 467 | + """ |
| 468 | + server_info = mock.Mock(spec=types.Implementation) |
| 469 | + server_info.name = "no-components" |
| 470 | + session = mock.AsyncMock(spec=mcp.ClientSession) |
| 471 | + session.list_tools.return_value = mock.AsyncMock(tools=[]) |
| 472 | + session.list_resources.return_value = mock.AsyncMock(resources=[]) |
| 473 | + session.list_prompts.return_value = mock.AsyncMock(prompts=[]) |
| 474 | + |
| 475 | + closed = False |
| 476 | + |
| 477 | + async def record_close() -> None: |
| 478 | + nonlocal closed |
| 479 | + closed = True |
| 480 | + |
| 481 | + group = ClientSessionGroup() |
| 482 | + |
| 483 | + async def establish_session(*args: object, **kwargs: object): |
| 484 | + # Mirrors the real `_establish_session`: the session owns a stack |
| 485 | + # holding its transport, registered before components are aggregated. |
| 486 | + session_stack = contextlib.AsyncExitStack() |
| 487 | + session_stack.push_async_callback(record_close) |
| 488 | + group._session_exit_stacks[session] = session_stack |
| 489 | + return server_info, session |
| 490 | + |
| 491 | + async with group: |
| 492 | + with mock.patch.object(group, "_establish_session", establish_session): |
| 493 | + await group.connect_to_server(StdioServerParameters(command="test")) |
| 494 | + |
| 495 | + assert not closed |
| 496 | + |
| 497 | + await group.disconnect_from_server(session) |
| 498 | + |
| 499 | + assert closed |
0 commit comments