fix: disconnect client in reverse order

Author: wzxxing

mcp_proxy_for_aws/proxy.py

@@ -130,7 +130,7 @@ class AWSMCPProxyClientFactory:
     def __init__(self, transport: ClientTransport) -> None:
         """Initialize a client factory with transport."""
         self._transport = transport
-        self._client = AWSMCPProxyClient(transport)
+        self._client: AWSMCPProxyClient | None = None
         self._clients: list[AWSMCPProxyClient] = []
         self._initialize_request: InitializeRequest | None = None
 
@@ -140,8 +140,9 @@ def set_init_params(self, initialize_request: InitializeRequest):
 
     async def get_client(self) -> Client:
         """Get client."""
-        if not self._client.is_connected():
+        if self._client is None or not self._client.is_connected():
             self._client = AWSMCPProxyClient(self._transport)
+            self._clients.append(self._client)
 
         return self._client
 

tests/unit/test_proxy.py

@@ -150,7 +150,8 @@ def test_client_factory_initialization():
     factory = AWSMCPProxyClientFactory(mock_transport)
 
     assert factory._transport == mock_transport
-    assert isinstance(factory._client, AWSMCPProxyClient)
+    assert factory._client is None
+    assert factory._clients == []
     assert factory._initialize_request is None
 
 
@@ -171,10 +172,12 @@ async def test_client_factory_get_client_when_connected():
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    factory._client.is_connected = Mock(return_value=True)
+    mock_client = Mock(spec=AWSMCPProxyClient)
+    mock_client.is_connected = Mock(return_value=True)
+    factory._client = mock_client
 
     client = await factory.get_client()
-    assert client == factory._client
+    assert client == mock_client
 
 
 @pytest.mark.asyncio
@@ -183,12 +186,14 @@ async def test_client_factory_get_client_when_disconnected():
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    old_client = factory._client
-    factory._client.is_connected = Mock(return_value=False)
+    mock_old_client = Mock(spec=AWSMCPProxyClient)
+    mock_old_client.is_connected = Mock(return_value=False)
+    factory._client = mock_old_client
 
     client = await factory.get_client()
-    assert client != old_client
+    assert client != mock_old_client
     assert isinstance(client, AWSMCPProxyClient)
+    assert client in factory._clients
 
 
 @pytest.mark.asyncio
@@ -197,10 +202,12 @@ async def test_client_factory_callable_interface():
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    factory._client.is_connected = Mock(return_value=True)
+    mock_client = Mock(spec=AWSMCPProxyClient)
+    mock_client.is_connected = Mock(return_value=True)
+    factory._client = mock_client
 
     client = await factory()
-    assert client == factory._client
+    assert client == mock_client
 
 
 @pytest.mark.asyncio
@@ -222,20 +229,49 @@ async def test_client_factory_disconnect_all():
     mock_client2._disconnect.assert_called_once_with(force=True)
 
 
+@pytest.mark.asyncio
+async def test_client_factory_disconnect_all_reverse_order():
+    """Test disconnect_all disconnects clients in reverse order."""
+    mock_transport = Mock(spec=ClientTransport)
+    factory = AWSMCPProxyClientFactory(mock_transport)
+
+    disconnect_order = []
+
+    mock_client1 = Mock()
+    mock_client1._disconnect = AsyncMock(side_effect=lambda **kwargs: disconnect_order.append(1))
+    mock_client2 = Mock()
+    mock_client2._disconnect = AsyncMock(side_effect=lambda **kwargs: disconnect_order.append(2))
+    mock_client3 = Mock()
+    mock_client3._disconnect = AsyncMock(side_effect=lambda **kwargs: disconnect_order.append(3))
+
+    factory._clients = [mock_client1, mock_client2, mock_client3]
+
+    await factory.disconnect_all()
+
+    assert disconnect_order == [3, 2, 1]
+
+
 @pytest.mark.asyncio
 async def test_client_factory_disconnect_all_handles_exceptions():
-    """Test disconnect_all handles exceptions gracefully."""
+    """Test disconnect_all handles exceptions gracefully and continues in reverse order."""
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
+    disconnect_order = []
+
     mock_client1 = Mock()
-    mock_client1._disconnect = AsyncMock(side_effect=Exception('Disconnect failed'))
+    mock_client1._disconnect = AsyncMock(side_effect=lambda **kwargs: disconnect_order.append(1))
     mock_client2 = Mock()
-    mock_client2._disconnect = AsyncMock()
+    mock_client2._disconnect = AsyncMock(side_effect=Exception('Disconnect failed'))
+    mock_client3 = Mock()
+    mock_client3._disconnect = AsyncMock(side_effect=lambda **kwargs: disconnect_order.append(3))
 
-    factory._clients = [mock_client1, mock_client2]
+    factory._clients = [mock_client1, mock_client2, mock_client3]
 
     await factory.disconnect_all()
 
+    # Verify client3 and client1 were disconnected despite client2 failing
+    assert disconnect_order == [3, 1]
     mock_client1._disconnect.assert_called_once_with(force=True)
     mock_client2._disconnect.assert_called_once_with(force=True)
+    mock_client3._disconnect.assert_called_once_with(force=True)