|
1 | 1 | import contextlib |
2 | 2 | from unittest import mock |
3 | 3 |
|
| 4 | +import anyio |
4 | 5 | import httpx2 |
5 | 6 | import mcp_types as types |
6 | 7 | import pytest |
|
13 | 14 | StreamableHttpParameters, |
14 | 15 | ) |
15 | 16 | from mcp.client.stdio import StdioServerParameters |
| 17 | +from mcp.server import Server, ServerRequestContext |
16 | 18 | from mcp.shared.exceptions import MCPError |
17 | 19 |
|
18 | 20 |
|
@@ -408,8 +410,50 @@ async def test_client_session_group_establish_session_parameterized( |
408 | 410 |
|
409 | 411 |
|
410 | 412 | @pytest.mark.anyio |
| 413 | +@pytest.mark.filterwarnings("ignore:Passing 'msg' argument.*:DeprecationWarning") |
411 | 414 | async def test_disconnect_removes_all_tools_after_same_session_reregistration() -> None: |
412 | | - """Verify disconnecting after re-registering a session removes all tools.""" |
| 415 | + """Reproduce the SDK cleanup gap after a caller re-registers a session with changed tools.""" |
| 416 | + catalog = [types.Tool(name="old", input_schema={"type": "object"})] |
| 417 | + |
| 418 | + async def list_tools( |
| 419 | + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None |
| 420 | + ) -> types.ListToolsResult: |
| 421 | + return types.ListToolsResult(tools=list(catalog)) |
| 422 | + |
| 423 | + async def list_resources( |
| 424 | + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None |
| 425 | + ) -> types.ListResourcesResult: |
| 426 | + return types.ListResourcesResult(resources=[]) |
| 427 | + |
| 428 | + async def list_prompts( |
| 429 | + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None |
| 430 | + ) -> types.ListPromptsResult: |
| 431 | + return types.ListPromptsResult(prompts=[]) |
| 432 | + |
| 433 | + server = Server( |
| 434 | + "mutable-catalog", |
| 435 | + on_list_tools=list_tools, |
| 436 | + on_list_resources=list_resources, |
| 437 | + on_list_prompts=list_prompts, |
| 438 | + ) |
| 439 | + group = ClientSessionGroup() |
| 440 | + with anyio.fail_after(5): |
| 441 | + async with mcp.Client(server) as client: |
| 442 | + server_info = types.Implementation(name="mutable-catalog", version="1") |
| 443 | + await group.connect_with_session(server_info, client.session) |
| 444 | + |
| 445 | + catalog[:] = [types.Tool(name="new", input_schema={"type": "object"})] |
| 446 | + await group.connect_with_session(server_info, client.session) |
| 447 | + |
| 448 | + await group.disconnect_from_server(client.session) |
| 449 | + |
| 450 | + assert group.sessions == [] |
| 451 | + assert group.tools == {} |
| 452 | + |
| 453 | + |
| 454 | +@pytest.mark.anyio |
| 455 | +async def test_disconnect_removes_all_tools_after_same_session_reregistration_mock() -> None: |
| 456 | + """Verify disconnecting after re-registering a session removes all tools (mock-based).""" |
413 | 457 | server_info = types.Implementation(name="TestServer", version="1.0") |
414 | 458 | mock_session = mock.AsyncMock(spec=mcp.ClientSession) |
415 | 459 |
|
@@ -669,21 +713,151 @@ async def test_connect_with_session_empty_catalog() -> None: |
669 | 713 |
|
670 | 714 | @pytest.mark.anyio |
671 | 715 | async def test_connect_to_server_empty_catalog(mock_exit_stack: contextlib.AsyncExitStack) -> None: |
672 | | - """Verify group-owned connect_to_server with empty catalog cleans up exit stack without error.""" |
| 716 | + """Verify group-owned connect_to_server with empty catalog retains exit stack and cleans up on disconnect.""" |
673 | 717 | mock_server_info = mock.Mock(spec=types.Implementation) |
674 | 718 | mock_server_info.name = "EmptyServer" |
675 | 719 | mock_session = mock.AsyncMock(spec=mcp.ClientSession) |
676 | 720 | mock_session.list_tools.return_value = types.ListToolsResult(tools=[]) |
677 | 721 | mock_session.list_resources.return_value = types.ListResourcesResult(resources=[]) |
678 | 722 | mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) |
679 | 723 |
|
| 724 | + session_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack) |
| 725 | + |
680 | 726 | async def fake_establish(*args: object, **kwargs: object) -> tuple[types.Implementation, mock.AsyncMock]: |
681 | | - group._session_exit_stacks[mock_session] = mock.AsyncMock() |
| 727 | + group._session_exit_stacks[mock_session] = session_stack |
682 | 728 | return mock_server_info, mock_session |
683 | 729 |
|
684 | 730 | group = ClientSessionGroup(exit_stack=mock_exit_stack) |
685 | 731 | with mock.patch.object(group, "_establish_session", side_effect=fake_establish): |
686 | 732 | await group.connect_to_server(StdioServerParameters(command="test")) |
687 | 733 |
|
| 734 | + assert mock_session in group._session_exit_stacks |
| 735 | + assert group.tools == {} |
| 736 | + assert group.sessions == [mock_session] |
| 737 | + |
| 738 | + await group.disconnect_from_server(mock_session) |
| 739 | + session_stack.aclose.assert_awaited_once() |
688 | 740 | assert mock_session not in group._session_exit_stacks |
| 741 | + assert group.sessions == [] |
| 742 | + |
| 743 | + |
| 744 | +@pytest.mark.anyio |
| 745 | +async def test_group_owned_session_refreshed_to_empty_and_disconnected_cleans_up( |
| 746 | + mock_exit_stack: contextlib.AsyncExitStack, |
| 747 | +) -> None: |
| 748 | + """Verify group-owned session refreshed to empty catalog retains exit stack and closes it on disconnect.""" |
| 749 | + mock_server_info = mock.Mock(spec=types.Implementation) |
| 750 | + mock_server_info.name = "TestServer" |
| 751 | + mock_session = mock.AsyncMock(spec=mcp.ClientSession) |
| 752 | + initial_tool = types.Tool(name="initial_tool", input_schema={}) |
| 753 | + mock_session.list_tools.return_value = types.ListToolsResult(tools=[initial_tool]) |
| 754 | + mock_session.list_resources.return_value = types.ListResourcesResult(resources=[]) |
| 755 | + mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) |
| 756 | + |
| 757 | + session_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack) |
| 758 | + |
| 759 | + async def fake_establish(*args: object, **kwargs: object) -> tuple[types.Implementation, mock.AsyncMock]: |
| 760 | + group._session_exit_stacks[mock_session] = session_stack |
| 761 | + return mock_server_info, mock_session |
| 762 | + |
| 763 | + group = ClientSessionGroup(exit_stack=mock_exit_stack) |
| 764 | + with mock.patch.object(group, "_establish_session", side_effect=fake_establish): |
| 765 | + await group.connect_to_server(StdioServerParameters(command="test")) |
| 766 | + |
| 767 | + assert mock_session in group._session_exit_stacks |
| 768 | + assert group.tools == {"initial_tool": initial_tool} |
| 769 | + assert group.sessions == [mock_session] |
| 770 | + |
| 771 | + # Refresh to empty catalog |
| 772 | + mock_session.list_tools.return_value = types.ListToolsResult(tools=[]) |
| 773 | + await group.connect_with_session(mock_server_info, mock_session) |
| 774 | + |
689 | 775 | assert group.tools == {} |
| 776 | + assert group.sessions == [mock_session] |
| 777 | + # Exit stack must NOT be deleted without closing |
| 778 | + assert mock_session in group._session_exit_stacks |
| 779 | + session_stack.aclose.assert_not_called() |
| 780 | + |
| 781 | + # Disconnect must close the exit stack |
| 782 | + await group.disconnect_from_server(mock_session) |
| 783 | + session_stack.aclose.assert_awaited_once() |
| 784 | + assert mock_session not in group._session_exit_stacks |
| 785 | + assert group.sessions == [] |
| 786 | + |
| 787 | + |
| 788 | +@pytest.mark.anyio |
| 789 | +async def test_connect_with_session_failed_refresh_preserves_existing_components() -> None: |
| 790 | + """Verify that MCPError when fetching components during refresh preserves existing registrations.""" |
| 791 | + server_info = types.Implementation(name="TestServer", version="1.0") |
| 792 | + session = mock.AsyncMock(spec=mcp.ClientSession) |
| 793 | + |
| 794 | + initial_prompt = types.Prompt(name="orig_prompt") |
| 795 | + initial_resource = types.Resource(uri="file:///orig_resource", name="orig_res") |
| 796 | + initial_tool = types.Tool(name="orig_tool", input_schema={}) |
| 797 | + |
| 798 | + session.list_prompts.return_value = types.ListPromptsResult(prompts=[initial_prompt]) |
| 799 | + session.list_resources.return_value = types.ListResourcesResult(resources=[initial_resource]) |
| 800 | + session.list_tools.return_value = types.ListToolsResult(tools=[initial_tool]) |
| 801 | + |
| 802 | + group = ClientSessionGroup() |
| 803 | + await group.connect_with_session(server_info, session) |
| 804 | + |
| 805 | + assert group.prompts == {"orig_prompt": initial_prompt} |
| 806 | + assert group.resources == {"orig_res": initial_resource} |
| 807 | + assert group.tools == {"orig_tool": initial_tool} |
| 808 | + assert group._tool_to_session == {"orig_tool": session} |
| 809 | + assert group.sessions == [session] |
| 810 | + |
| 811 | + # Simulate MCPError on all component listings during refresh |
| 812 | + session.list_prompts.side_effect = MCPError(code=types.INTERNAL_ERROR, message="prompts failed") |
| 813 | + session.list_resources.side_effect = MCPError(code=types.INTERNAL_ERROR, message="resources failed") |
| 814 | + session.list_tools.side_effect = MCPError(code=types.INTERNAL_ERROR, message="tools failed") |
| 815 | + |
| 816 | + await group.connect_with_session(server_info, session) |
| 817 | + |
| 818 | + # Invariants: existing components are preserved, not evicted |
| 819 | + assert group.prompts == {"orig_prompt": initial_prompt} |
| 820 | + assert group.resources == {"orig_res": initial_resource} |
| 821 | + assert group.tools == {"orig_tool": initial_tool} |
| 822 | + assert group._tool_to_session == {"orig_tool": session} |
| 823 | + assert group._sessions[session].prompts == {"orig_prompt"} |
| 824 | + assert group._sessions[session].resources == {"orig_res"} |
| 825 | + assert group._sessions[session].tools == {"orig_tool"} |
| 826 | + assert group.sessions == [session] |
| 827 | + |
| 828 | + # Preserved tools remain callable and routed |
| 829 | + session.call_tool.return_value = types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) |
| 830 | + call_res = await group.call_tool("orig_tool", {}) |
| 831 | + assert isinstance(call_res.content[0], types.TextContent) |
| 832 | + assert call_res.content[0].text == "ok" |
| 833 | + |
| 834 | + # Disconnect cleans up preserved components completely |
| 835 | + await group.disconnect_from_server(session) |
| 836 | + assert group.prompts == {} |
| 837 | + assert group.resources == {} |
| 838 | + assert group.tools == {} |
| 839 | + assert group._tool_to_session == {} |
| 840 | + assert group.sessions == [] |
| 841 | + |
| 842 | + |
| 843 | +@pytest.mark.anyio |
| 844 | +async def test_connect_with_session_initial_fetch_mcperror_handled() -> None: |
| 845 | + """Verify that MCPError on initial component queries is handled gracefully without error.""" |
| 846 | + server_info = types.Implementation(name="ErrorServer", version="1.0") |
| 847 | + session = mock.AsyncMock(spec=mcp.ClientSession) |
| 848 | + |
| 849 | + session.list_prompts.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no prompts") |
| 850 | + session.list_resources.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no resources") |
| 851 | + session.list_tools.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no tools") |
| 852 | + |
| 853 | + group = ClientSessionGroup() |
| 854 | + await group.connect_with_session(server_info, session) |
| 855 | + |
| 856 | + assert group.prompts == {} |
| 857 | + assert group.resources == {} |
| 858 | + assert group.tools == {} |
| 859 | + assert group._tool_to_session == {} |
| 860 | + assert group.sessions == [session] |
| 861 | + |
| 862 | + await group.disconnect_from_server(session) |
| 863 | + assert group.sessions == [] |
0 commit comments