fix: do not use multiple clients

Author: wzxxing

mcp_proxy_for_aws/proxy.py

@@ -140,8 +140,7 @@ class AWSMCPProxyClientFactory:
     def __init__(self, transport: ClientTransport) -> None:
         """Initialize a client factory with transport."""
         self._transport = transport
-        self._client = AWSMCPProxyClient(transport)
-        self._clients: list[AWSMCPProxyClient] = []
+        self._client: AWSMCPProxyClient | None = None
         self._initialize_request: InitializeRequest | None = None
 
     def set_init_params(self, initialize_request: InitializeRequest):
@@ -150,7 +149,7 @@ 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:
             self._client = AWSMCPProxyClient(self._transport)
 
         return self._client
@@ -159,10 +158,10 @@ async def __call__(self) -> Client:
         """Implement the callable factory interface."""
         return await self.get_client()
 
-    async def disconnect_all(self):
+    async def disconnect(self):
         """Disconnect all the clients (no throw)."""
-        for client in reversed(self._clients):
-            try:
-                await client._disconnect(force=True)
-            except Exception:
-                logger.exception('Failed to disconnect client.')
+        try:
+            if self._client:
+                await self._client._disconnect(force=True)
+        except Exception:
+            logger.exception('Failed to disconnect client.')

mcp_proxy_for_aws/server.py

@@ -105,7 +105,7 @@ async def run_proxy(args) -> None:
         logger.error('Cannot start proxy server: %s', e)
         raise e
     finally:
-        await client_factory.disconnect_all()
+        await client_factory.disconnect()
 
 
 def add_tool_filtering_middleware(mcp: FastMCP, read_only: bool = False) -> None:

tests/unit/test_proxy.py

@@ -150,7 +150,7 @@ 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._initialize_request is None
 
 
@@ -171,10 +171,11 @@ 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)
+    factory._client = mock_client
 
     client = await factory.get_client()
-    assert client == factory._client
+    assert client == mock_client
 
 
 @pytest.mark.asyncio
@@ -183,12 +184,9 @@ 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)
-
     client = await factory.get_client()
-    assert client != old_client
     assert isinstance(client, AWSMCPProxyClient)
+    assert factory._client == client
 
 
 @pytest.mark.asyncio
@@ -197,45 +195,35 @@ async def test_client_factory_callable_interface():
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    factory._client.is_connected = Mock(return_value=True)
-
     client = await factory()
-    assert client == factory._client
+    assert isinstance(client, AWSMCPProxyClient)
 
 
 @pytest.mark.asyncio
 async def test_client_factory_disconnect_all():
-    """Test disconnect_all disconnects all clients."""
+    """Test disconnect disconnects the client."""
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    mock_client1 = Mock()
-    mock_client1._disconnect = AsyncMock()
-    mock_client2 = Mock()
-    mock_client2._disconnect = AsyncMock()
+    mock_client = Mock()
+    mock_client._disconnect = AsyncMock()
+    factory._client = mock_client
 
-    factory._clients = [mock_client1, mock_client2]
+    await factory.disconnect()
 
-    await factory.disconnect_all()
-
-    mock_client1._disconnect.assert_called_once_with(force=True)
-    mock_client2._disconnect.assert_called_once_with(force=True)
+    mock_client._disconnect.assert_called_once_with(force=True)
 
 
 @pytest.mark.asyncio
 async def test_client_factory_disconnect_all_handles_exceptions():
-    """Test disconnect_all handles exceptions gracefully."""
+    """Test disconnect handles exceptions gracefully."""
     mock_transport = Mock(spec=ClientTransport)
     factory = AWSMCPProxyClientFactory(mock_transport)
 
-    mock_client1 = Mock()
-    mock_client1._disconnect = AsyncMock(side_effect=Exception('Disconnect failed'))
-    mock_client2 = Mock()
-    mock_client2._disconnect = AsyncMock()
-
-    factory._clients = [mock_client1, mock_client2]
+    mock_client = Mock()
+    mock_client._disconnect = AsyncMock(side_effect=Exception('Disconnect failed'))
+    factory._client = mock_client
 
-    await factory.disconnect_all()
+    await factory.disconnect()
 
-    mock_client1._disconnect.assert_called_once_with(force=True)
-    mock_client2._disconnect.assert_called_once_with(force=True)
+    mock_client._disconnect.assert_called_once_with(force=True)

tests/unit/test_server.py

@@ -73,7 +73,7 @@ async def test_setup_mcp_mode(
         mock_create_transport.return_value = mock_transport
 
         mock_client_factory = Mock()
-        mock_client_factory.disconnect_all = AsyncMock()
+        mock_client_factory.disconnect = AsyncMock()
         mock_client_factory_class.return_value = mock_client_factory
 
         mock_proxy = Mock()
@@ -145,7 +145,7 @@ async def test_setup_mcp_mode_no_retries(
         mock_create_transport.return_value = mock_transport
 
         mock_client_factory = Mock()
-        mock_client_factory.disconnect_all = AsyncMock()
+        mock_client_factory.disconnect = AsyncMock()
         mock_client_factory_class.return_value = mock_client_factory
 
         mock_proxy = Mock()
@@ -216,7 +216,7 @@ async def test_setup_mcp_mode_no_metadata_injects_aws_region(
         mock_create_transport.return_value = mock_transport
 
         mock_client_factory = Mock()
-        mock_client_factory.disconnect_all = AsyncMock()
+        mock_client_factory.disconnect = AsyncMock()
         mock_client_factory_class.return_value = mock_client_factory
 
         mock_proxy = Mock()
@@ -271,7 +271,7 @@ async def test_setup_mcp_mode_metadata_without_aws_region_injects_it(
         mock_create_transport.return_value = mock_transport
 
         mock_client_factory = Mock()
-        mock_client_factory.disconnect_all = AsyncMock()
+        mock_client_factory.disconnect = AsyncMock()
         mock_client_factory_class.return_value = mock_client_factory
 
         mock_proxy = Mock()