Merge branch 'main' into aws-mcp-integration-tests

Author: detti456

CHANGELOG.md

@@ -5,6 +5,22 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## v1.1.4 (2025-12-04)
+
+### Fix
+
+- do not call initialize for q dev cli / kiro cli
+- patch fastmcp lowlevel session method
+- connect remote mcp client immediately in the initialize middleware
+
+## v1.1.3 (2025-12-04)
+
+### Fix
+
+- avoid infinite recursion (#111)
+- init client and show errors in all clients (#108)
+- set the fastmcp log level to be the same as the proxy (#101)
+
 ## v1.1.2 (2025-11-27)
 
 ### Fix

README.md

@@ -315,3 +315,6 @@ Licensed under the Apache License, Version 2.0 (the "License").
 ## Disclaimer
 
 LLMs are non-deterministic and they make mistakes, we advise you to always thoroughly test and follow the best practices of your organization before using these tools on customer facing accounts. Users of this package are solely responsible for implementing proper security controls and MUST use AWS Identity and Access Management (IAM) to manage access to AWS resources. You are responsible for configuring appropriate IAM policies, roles, and permissions, and any security vulnerabilities resulting from improper IAM configuration are your sole responsibility. By using this package, you acknowledge that you have read and understood this disclaimer and agree to use the package at your own risk.
+
+<!-- mcp-name: io.github.aws/mcp-proxy-for-aws -->
+<!-- mcp-name: io.github.aws/aws-mcp -->

mcp_proxy_for_aws/__init__.py

@@ -16,6 +16,8 @@
 
 from importlib.metadata import version as _metadata_version
 
+import mcp_proxy_for_aws.fastmcp_patch as _fastmcp_patch
+
 
 __all__ = ['__version__']
 __version__ = _metadata_version('mcp-proxy-for-aws')

mcp_proxy_for_aws/fastmcp_patch.py

@@ -0,0 +1,34 @@
+import fastmcp.server.low_level as low_level_module
+import mcp.types
+from functools import wraps
+from mcp import McpError
+from mcp.server.stdio import stdio_server as stdio_server
+from mcp.shared.session import RequestResponder
+
+
+original_receive_request = low_level_module.MiddlewareServerSession._received_request
+
+
+@wraps(original_receive_request)
+async def _received_request(
+    self,
+    responder: RequestResponder[mcp.types.ClientRequest, mcp.types.ServerResult],
+):
+    """Monkey patch fastmcp so that the initialize error from the middleware can be send back to the client.
+
+    https://github.com/jlowin/fastmcp/pull/2531
+    """
+    if isinstance(responder.request.root, mcp.types.InitializeRequest):
+        try:
+            return await original_receive_request(self, responder)
+        except McpError as e:
+            if not responder._completed:
+                with responder:
+                    return await responder.respond(e.error)
+
+            raise e
+    else:
+        return await original_receive_request(self, responder)
+
+
+low_level_module.MiddlewareServerSession._received_request = _received_request

mcp_proxy_for_aws/middleware/initialize_middleware.py

@@ -25,6 +25,20 @@ 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
+            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.')

mcp_proxy_for_aws/proxy.py

@@ -78,12 +78,13 @@ def __init__(
 class AWSMCPProxyClient(_ProxyClient):
     """Proxy client that handles HTTP errors when connection fails."""
 
-    def __init__(self, transport: ClientTransport, **kwargs):
+    def __init__(self, transport: ClientTransport, max_connect_retry=3, **kwargs):
         """Constructor of AutoRefreshProxyCilent."""
         super().__init__(transport, **kwargs)
+        self._max_connect_retry = max_connect_retry
 
     @override
-    async def _connect(self):
+    async def _connect(self, retry=0):
         """Enter as normal && initialize only once."""
         logger.debug('Connecting %s', self)
         try:
@@ -96,27 +97,36 @@ async def _connect(self):
             try:
                 body = await response.aread()
                 jsonrpc_msg = JSONRPCMessage.model_validate_json(body).root
-            except Exception:
-                logger.debug('HTTP error is not a valid MCP message.')
+            except Exception as e:
+                logger.debug('HTTP error is not a valid MCP message.', exc_info=e)
                 raise http_error
 
             if isinstance(jsonrpc_msg, JSONRPCError):
-                logger.debug('Converting HTTP error to MCP error %s', http_error)
+                logger.debug('Converting HTTP error to MCP error', exc_info=http_error)
                 # raising McpError so that the sdk can handle the exception properly
                 raise McpError(error=jsonrpc_msg.error) from http_error
             else:
                 raise http_error
-        except RuntimeError:
+        except RuntimeError as e:
+            if isinstance(e.__cause__, McpError):
+                raise e.__cause__
+
+            if retry > self._max_connect_retry:
+                raise e
+
             try:
-                logger.warning('encountered runtime error, try force disconnect.')
+                logger.warning('encountered runtime error, try force disconnect.', exc_info=e)
                 await self._disconnect(force=True)
-            except Exception:
+            except httpx.TimeoutException:
                 # _disconnect awaits on the session_task,
                 # which raises the timeout error that caused the client session to be terminated.
                 # the error is ignored as long as the counter is force set to 0.
                 # TODO: investigate how timeout error is handled by fastmcp and httpx
-                logger.exception('encountered another error, ignoring.')
-            return await self._connect()
+                logger.exception(
+                    'Session was terminated due to timeout error, ignore and reconnect'
+                )
+
+            return await self._connect(retry + 1)
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """The MCP Proxy for AWS project is a proxy from stdio to http (sigv4).

pyproject.toml

@@ -10,7 +10,7 @@ members = [
 name = "mcp-proxy-for-aws"
 
 # NOTE: "Patch"=9223372036854775807 bumps next release to zero.
-version = "1.1.2"
+version = "1.1.4"
 
 description = "MCP Proxy for AWS"
 readme = "README.md"

tests/unit/test_fastmcp_patch.py

@@ -0,0 +1,106 @@
+import mcp.types as mt
+import pytest
+from mcp import McpError
+from mcp.shared.session import RequestResponder
+from unittest.mock import AsyncMock, Mock, patch
+
+
+@pytest.mark.asyncio
+async def test_patched_received_request_initialize_success():
+    """Test that patched _received_request calls original for successful initialize."""
+    # Import after patching is applied
+    import fastmcp.server.low_level as low_level_module
+    from mcp_proxy_for_aws import fastmcp_patch
+
+    mock_self = Mock()
+    mock_self.fastmcp = Mock()
+
+    mock_request = Mock()
+    mock_request.root = Mock(spec=mt.InitializeRequest)
+
+    mock_responder = Mock(spec=RequestResponder)
+    mock_responder.request = mock_request
+
+    with patch.object(
+        fastmcp_patch, 'original_receive_request', new_callable=AsyncMock
+    ) as mock_original:
+        await low_level_module.MiddlewareServerSession._received_request(mock_self, mock_responder)
+        mock_original.assert_called_once_with(mock_self, mock_responder)
+
+
+@pytest.mark.asyncio
+async def test_patched_received_request_initialize_mcp_error_not_completed():
+    """Test that patched _received_request handles McpError when responder not completed."""
+    import fastmcp.server.low_level as low_level_module
+    from mcp_proxy_for_aws import fastmcp_patch
+
+    mock_self = Mock()
+    mock_self.fastmcp = Mock()
+
+    mock_request = Mock()
+    mock_request.root = Mock(spec=mt.InitializeRequest)
+
+    mock_responder = Mock(spec=RequestResponder)
+    mock_responder.request = mock_request
+    mock_responder._completed = False
+    mock_responder.__enter__ = Mock(return_value=mock_responder)
+    mock_responder.__exit__ = Mock(return_value=False)
+    mock_responder.respond = AsyncMock()
+
+    error = mt.ErrorData(code=1, message='test error')
+    mcp_error = McpError(error=error)
+
+    with patch.object(
+        fastmcp_patch, 'original_receive_request', new_callable=AsyncMock, side_effect=mcp_error
+    ):
+        await low_level_module.MiddlewareServerSession._received_request(mock_self, mock_responder)
+        mock_responder.respond.assert_called_once_with(error)
+
+
+@pytest.mark.asyncio
+async def test_patched_received_request_initialize_mcp_error_completed():
+    """Test that patched _received_request re-raises McpError when responder completed."""
+    import fastmcp.server.low_level as low_level_module
+    from mcp_proxy_for_aws import fastmcp_patch
+
+    mock_self = Mock()
+    mock_self.fastmcp = Mock()
+
+    mock_request = Mock()
+    mock_request.root = Mock(spec=mt.InitializeRequest)
+
+    mock_responder = Mock(spec=RequestResponder)
+    mock_responder.request = mock_request
+    mock_responder._completed = True
+
+    error = mt.ErrorData(code=1, message='test error')
+    mcp_error = McpError(error=error)
+
+    with patch.object(
+        fastmcp_patch, 'original_receive_request', new_callable=AsyncMock, side_effect=mcp_error
+    ):
+        with pytest.raises(McpError):
+            await low_level_module.MiddlewareServerSession._received_request(
+                mock_self, mock_responder
+            )
+
+
+@pytest.mark.asyncio
+async def test_patched_received_request_non_initialize():
+    """Test that patched _received_request calls original for non-initialize requests."""
+    import fastmcp.server.low_level as low_level_module
+    from mcp_proxy_for_aws import fastmcp_patch
+
+    mock_self = Mock()
+
+    mock_request = Mock()
+    mock_request.root = Mock(spec=mt.CallToolRequest)
+
+    mock_responder = Mock(spec=RequestResponder)
+    mock_responder.request = mock_request
+
+    with patch.object(
+        fastmcp_patch, 'original_receive_request', new_callable=AsyncMock
+    ) as mock_original:
+        await low_level_module.MiddlewareServerSession._received_request(mock_self, mock_responder)
+        mock_original.assert_called_once_with(mock_self, mock_responder)

tests/unit/test_initialize_middleware.py

@@ -0,0 +1,98 @@
+import mcp.types as mt
+import pytest
+from mcp_proxy_for_aws.middleware.initialize_middleware import InitializeMiddleware
+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()."""
+    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('test-client')
+
+    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 = 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)