feat: Add RateLimiting middleware for client calls to proxy and exponential backoff middleware for network mcp. (#19)

Author: arangatang

aws_mcp_proxy/mcp_proxy_manager.py

@@ -15,6 +15,8 @@
 """MCP Proxy Manager for handling proxy content integration."""
 
 import logging
+from fastmcp.server.middleware.error_handling import RetryMiddleware
+from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
 from fastmcp.server.server import FastMCP
 
 
@@ -30,6 +32,7 @@ def __init__(self, target_mcp: FastMCP, read_only: bool = False):
             target_mcp: The target MCP server to add content to
             read_only: If true, disable tools that require write permissions OR that do not have `readOnlyHint` set.
         """
+        self._add_rate_limiting_middleware(target_mcp)
         self.target_mcp = target_mcp
         self.read_only = read_only
 
@@ -42,6 +45,8 @@ async def add_proxy_content(self, proxy: FastMCP) -> None:
         Raises:
             Exception: If tools cannot be retrieved or added
         """
+        self._add_retry_middleware(proxy)
+
         try:
             await self._add_tools(proxy)
             await self._add_resources(proxy)
@@ -115,3 +120,26 @@ async def _add_prompts(self, proxy: FastMCP) -> None:
             self.logger.debug("Proxy doesn't have prompts method")
         except Exception as e:
             self.logger.warning(f'Failed to get prompts from proxy: {e}')
+
+    def _add_retry_middleware(self, mcp: FastMCP) -> None:
+        """Add retry with exponential backoff middleware to target MCP server.
+
+        Args:
+            mcp: The FastMCP instance to add exponential backoff to
+        """
+        self.logger.info('Adding retry middleware')
+        mcp.add_middleware(RetryMiddleware())
+
+    def _add_rate_limiting_middleware(self, mcp):
+        """Add retry with exponential backoff middleware to target MCP server.
+
+        Args:
+            mcp: The FastMCP instance to add rate limiting to
+        """
+        self.logger.info('Adding rate limiting middleware')
+        mcp.add_middleware(
+            RateLimitingMiddleware(
+                max_requests_per_second=5,
+                burst_capacity=10,
+            )
+        )

aws_mcp_proxy/server.py

@@ -39,7 +39,7 @@
 logger = logging.getLogger(__name__)
 
 
-async def setup_mcp_mode(mcp: FastMCP, args) -> None:
+async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
     """Set up the server in MCP mode."""
     logger.info('Setting up server in MCP mode')
 
@@ -63,7 +63,8 @@ async def setup_mcp_mode(mcp: FastMCP, args) -> None:
     proxy = FastMCP.as_proxy(transport)
 
     # Use McpProxyManager to add proxy content
-    proxy_manager = McpProxyManager(mcp, args.read_only)
+
+    proxy_manager = McpProxyManager(local_mcp, args.read_only)
     await proxy_manager.add_proxy_content(proxy)