fix: avoid infinite recursion

can happen when the network is down and connect keeps failing

Author: wzxxing

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).

tests/unit/test_proxy.py

@@ -227,3 +227,72 @@ async def test_client_factory_disconnect_all_handles_exceptions():
     await factory.disconnect()
 
     mock_client._disconnect.assert_called_once_with(force=True)
+
+
+@pytest.mark.asyncio
+async def test_proxy_client_connect_runtime_error_with_mcp_error():
+    """Test connection handles RuntimeError wrapping McpError."""
+    mock_transport = Mock(spec=ClientTransport)
+    client = AWSMCPProxyClient(mock_transport)
+
+    error_data = ErrorData(code=-32600, message='Invalid Request')
+    mcp_error = McpError(error=error_data)
+    runtime_error = RuntimeError('Connection failed')
+    runtime_error.__cause__ = mcp_error
+
+    with patch('mcp_proxy_for_aws.proxy._ProxyClient._connect', side_effect=runtime_error):
+        with pytest.raises(McpError) as exc_info:
+            await client._connect()
+        assert exc_info.value.error.code == -32600
+
+
+@pytest.mark.asyncio
+async def test_proxy_client_connect_runtime_error_max_retries():
+    """Test connection stops retrying after max_connect_retry."""
+    mock_transport = Mock(spec=ClientTransport)
+    client = AWSMCPProxyClient(mock_transport, max_connect_retry=2)
+
+    runtime_error = RuntimeError('Connection failed')
+
+    with patch('mcp_proxy_for_aws.proxy._ProxyClient._connect', side_effect=runtime_error):
+        with patch.object(client, '_disconnect', new_callable=AsyncMock) as mock_disconnect:
+            with pytest.raises(RuntimeError):
+                await client._connect()
+            assert mock_disconnect.call_count == 3
+
+
+@pytest.mark.asyncio
+async def test_proxy_client_connect_runtime_error_with_timeout():
+    """Test connection handles TimeoutException during disconnect."""
+    mock_transport = Mock(spec=ClientTransport)
+    client = AWSMCPProxyClient(mock_transport, max_connect_retry=1)
+
+    runtime_error = RuntimeError('Connection failed')
+    call_count = 0
+
+    async def mock_connect_side_effect(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count <= 2:
+            raise runtime_error
+        return 'connected'
+
+    with patch(
+        'mcp_proxy_for_aws.proxy._ProxyClient._connect', side_effect=mock_connect_side_effect
+    ):
+        with patch.object(
+            client,
+            '_disconnect',
+            new_callable=AsyncMock,
+            side_effect=httpx.TimeoutException('timeout'),
+        ):
+            result = await client._connect()
+            assert result == 'connected'
+
+
+@pytest.mark.asyncio
+async def test_proxy_client_max_connect_retry_default():
+    """Test default max_connect_retry is 3."""
+    mock_transport = Mock(spec=ClientTransport)
+    client = AWSMCPProxyClient(mock_transport)
+    assert client._max_connect_retry == 3