diff --git a/mcp_proxy_for_aws/proxy.py b/mcp_proxy_for_aws/proxy.py index b0b3bec..399b618 100644 --- a/mcp_proxy_for_aws/proxy.py +++ b/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). diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index eaba425..3a37e13 100644 --- a/tests/unit/test_proxy.py +++ b/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 diff --git a/uv.lock b/uv.lock index df83e59..45ee02a 100644 --- a/uv.lock +++ b/uv.lock @@ -2458,8 +2458,8 @@ dev = [ [package.metadata] requires-dist = [ - { name = "boto3", specifier = ">=1.34.0" }, - { name = "botocore", extras = ["crt"], specifier = ">=1.34.0" }, + { name = "boto3", specifier = ">=1.41.0" }, + { name = "botocore", extras = ["crt"], specifier = ">=1.41.0" }, { name = "fastmcp", specifier = ">=2.13.1" }, ]