fix: do not call initialize for q dev cli / kiro cli

Author: wzxxing

mcp_proxy_for_aws/middleware/initialize_middleware.py

@@ -28,7 +28,17 @@ async def on_initialize(
             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()
+            client_name = context.message.params.clientInfo.name.lower()
+            if 'kiro cli' not in client_name and 'q dev cli' not in client_name:
+                # q cli / kiro cli uses the rust SDK which does not handle json rpc error
+                # properly during initialization.
+                # https://github.com/modelcontextprotocol/rust-sdk/pull/569
+                # if calling _connect below raise mcp error, the q cli will skip the message
+                # and continue wait for a json rpc response message which will never come.
+                # Luckily, q cli calls list tool immediately after being connected to a mcp server
+                # the list_tool call will require the client to be connected again, so the mcp error
+                # will be displayed in the q cli logs.
+                await client._connect()
             return await call_next(context)
         except Exception:
             logger.exception('Initialize failed in middleware.')

tests/unit/test_initialize_middleware.py

@@ -4,6 +4,18 @@
 from unittest.mock import AsyncMock, Mock
 
 
+def create_initialize_request(client_name: str) -> mt.InitializeRequest:
+    """Create a real InitializeRequest object."""
+    return mt.InitializeRequest(
+        method='initialize',
+        params=mt.InitializeRequestParams(
+            protocolVersion='2024-11-05',
+            capabilities=mt.ClientCapabilities(),
+            clientInfo=mt.Implementation(name=client_name, version='1.0'),
+        ),
+    )
+
+
 @pytest.mark.asyncio
 async def test_on_initialize_connects_client():
     """Test that on_initialize calls client._connect()."""
@@ -17,7 +29,7 @@ async def test_on_initialize_connects_client():
     middleware = InitializeMiddleware(mock_factory)
 
     mock_context = Mock()
-    mock_context.message = Mock(spec=mt.InitializeRequest)
+    mock_context.message = create_initialize_request('test-client')
 
     mock_call_next = AsyncMock()
 
@@ -42,11 +54,45 @@ async def test_on_initialize_fails_if_connect_fails():
     middleware = InitializeMiddleware(mock_factory)
 
     mock_context = Mock()
-    mock_context.message = Mock(spec=mt.InitializeRequest)
+    mock_context.message = create_initialize_request('test-client')
 
     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()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'client_name',
+    [
+        'Kiro CLI',
+        'kiro cli',
+        'KIRO CLI',
+        'Amazon Q Dev CLI',
+        'amazon q dev cli',
+        'Q DEV CLI',
+    ],
+)
+async def test_on_initialize_skips_connect_for_special_clients(client_name):
+    """Test that on_initialize skips _connect() for Kiro CLI and Q Dev CLI."""
+    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 = create_initialize_request(client_name)
+
+    mock_call_next = AsyncMock()
+
+    await middleware.on_initialize(mock_context, mock_call_next)
+
+    mock_client._connect.assert_not_called()
+    mock_call_next.assert_called_once_with(mock_context)