fix: connect remote mcp client immediately in the initialize middleware

Author: wzxxing

mcp_proxy_for_aws/middleware/initialize_middleware.py

@@ -25,6 +25,10 @@ async def on_initialize(
         try:
             logger.debug('Received initialize request %s.', context.message)
             self._client_factory.set_init_params(context.message)
+            client = await self._client_factory.get_client()
+            # connect the http client, fail and don't succeed the stdio connect
+            # if remote client cannot be connected
+            await client._connect()
             return await call_next(context)
         except Exception:
             logger.exception('Initialize failed in middleware.')

tests/unit/test_initialize_middleware.py

@@ -0,0 +1,52 @@
+import mcp.types as mt
+import pytest
+from mcp_proxy_for_aws.middleware.initialize_middleware import InitializeMiddleware
+from unittest.mock import AsyncMock, Mock
+
+
+@pytest.mark.asyncio
+async def test_on_initialize_connects_client():
+    """Test that on_initialize calls client._connect()."""
+    mock_client = Mock()
+    mock_client._connect = AsyncMock()
+
+    mock_factory = Mock()
+    mock_factory.set_init_params = Mock()
+    mock_factory.get_client = AsyncMock(return_value=mock_client)
+
+    middleware = InitializeMiddleware(mock_factory)
+
+    mock_context = Mock()
+    mock_context.message = Mock(spec=mt.InitializeRequest)
+
+    mock_call_next = AsyncMock()
+
+    await middleware.on_initialize(mock_context, mock_call_next)
+
+    mock_factory.set_init_params.assert_called_once_with(mock_context.message)
+    mock_factory.get_client.assert_called_once()
+    mock_client._connect.assert_called_once()
+    mock_call_next.assert_called_once_with(mock_context)
+
+
+@pytest.mark.asyncio
+async def test_on_initialize_fails_if_connect_fails():
+    """Test that on_initialize raises exception if _connect() fails."""
+    mock_client = Mock()
+    mock_client._connect = AsyncMock(side_effect=Exception('Connection failed'))
+
+    mock_factory = Mock()
+    mock_factory.set_init_params = Mock()
+    mock_factory.get_client = AsyncMock(return_value=mock_client)
+
+    middleware = InitializeMiddleware(mock_factory)
+
+    mock_context = Mock()
+    mock_context.message = Mock(spec=mt.InitializeRequest)
+
+    mock_call_next = AsyncMock()
+
+    with pytest.raises(Exception, match='Connection failed'):
+        await middleware.on_initialize(mock_context, mock_call_next)
+
+    mock_call_next.assert_not_called()