Fix dynamic tooling (#28)

* fix(server.py): remove redundant proxy and bypass mc_proxy_manager

* feat(tools_filter.py): add tool_filter.py as filtering middleware

* feat(server.py): add middleware directly to proxy

Instead of using the mcp_proxy_manager, here we directly add middleware to the proxy. This effectively makes the mcp_proxy_manager redundant.

* style(server.py,-mcp_proxy_manager.py): cleanup remnants of mcp_proxy_manager

* fix(unit_tests): fix unit tests that were broken as a result of removing mcp_proxy_manager.py

* fix(server.py): remove rate limiting middleware

* fix(tool_filter.py): add type hint to call_next

* refactor(tool_filter.py): use class name for logger config

* perf(tool_filter.py): return early in tool filtering if read_only is false

* fix(logging_config.py): change log output to stderr

this fixes issues with logging, especially with printing to the terminal.

* style: fix failing stylecheck

* feat(server.py): backport retries as an argument to work without an mcp_proxy_manager

---------

Co-authored-by: Kyon Caldera <kyonc@amazon.com>

Author: kyoncal

aws_mcp_proxy/logging_config.py

@@ -34,7 +34,7 @@ def configure_logging(level: Optional[str] = None) -> None:
     date_format = '%Y-%m-%d %H:%M:%S'
 
     # Create console handler with formatting
-    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler = logging.StreamHandler(sys.stderr)
     console_handler.setFormatter(logging.Formatter(log_format, date_format))
 
     # Configure root logger

aws_mcp_proxy/mcp_proxy_manager.py

@@ -0,0 +1,57 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from collections.abc import Awaitable, Callable
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+from fastmcp.tools.tool import Tool
+
+
+class ToolFilteringMiddleware(Middleware):
+    """Middleware to filter tools based on read only flag."""
+
+    def __init__(self, read_only: bool, logger: logging.Logger | None = None):
+        """Initialize the middleware."""
+        self.read_only = read_only
+        self.logger = logger or logging.getLogger(__name__)
+
+    async def on_list_tools(
+        self,
+        context: MiddlewareContext,
+        call_next: Callable[[MiddlewareContext], Awaitable[list[Tool]]],
+    ):
+        """Filter tools based on read only flag."""
+        # Get list of FastMCP Components
+        tools = await call_next(context)
+        self.logger.info(f'Filtering tools for read only: {self.read_only}')
+
+        # If not read only, return the list of tools as is
+        if not self.read_only:
+            return tools
+
+        filtered_tools = []
+        for tool in tools:
+            # Check the tool annotations and disable if needed
+            annotations = tool.annotations
+
+            # Skip the tools with no readOnlyHint=True annotation
+            read_only_hint = getattr(annotations, 'readOnlyHint', False)
+            if not read_only_hint:
+                # Skip tools that don't have readOnlyHint=True
+                self.logger.info(f'Skipping tool {tool.name} needing write permissions')
+                continue
+
+            filtered_tools.append(tool)
+
+        return filtered_tools

aws_mcp_proxy/middleware/tool_filter.py

@@ -27,12 +27,13 @@
 import logging
 import os
 from aws_mcp_proxy.logging_config import configure_logging
-from aws_mcp_proxy.mcp_proxy_manager import McpProxyManager
+from aws_mcp_proxy.middleware.tool_filter import ToolFilteringMiddleware
 from aws_mcp_proxy.utils import (
     create_transport_with_sigv4,
     determine_aws_region,
     determine_service_name,
 )
+from fastmcp.server.middleware.error_handling import RetryMiddleware
 from fastmcp.server.server import FastMCP
 from typing import Any
 
@@ -62,10 +63,38 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
 
     # Create proxy with the transport
     proxy = FastMCP.as_proxy(transport)
+    add_tool_filtering_middleware(proxy, args.read_only)
 
-    # Use McpProxyManager to add proxy content
-    proxy_manager = McpProxyManager(local_mcp, args.read_only)
-    await proxy_manager.add_proxy_content(proxy, args.retries)
+    if args.retries:
+        add_retry_middleware(proxy, args.retries)
+
+    await proxy.run_async()
+
+
+def add_tool_filtering_middleware(mcp: FastMCP, read_only: bool = False) -> None:
+    """Add tool filtering middleware to target MCP server.
+
+    Args:
+        mcp: The FastMCP instance to add tool filtering to
+        read_only: Whether or not to filter out tools that require write permissions
+    """
+    logger.info('Adding tool filtering middleware')
+    mcp.add_middleware(
+        ToolFilteringMiddleware(
+            read_only=read_only,
+        )
+    )
+
+
+def add_retry_middleware(mcp: FastMCP, retries: int) -> None:
+    """Add retry with exponential backoff middleware to target MCP server.
+
+    Args:
+        mcp: The FastMCP instance to add exponential backoff to
+        retries: number of retries with which to configure the retry middleware
+    """
+    logger.info('Adding retry middleware')
+    mcp.add_middleware(RetryMiddleware(retries))
 
 
 def parse_args():
@@ -155,7 +184,6 @@ async def setup_and_run():
             await setup_mcp_mode(mcp, args)
 
             logger.info('Server setup complete, starting MCP server')
-            await mcp.run_async()
 
         except Exception as e:
             logger.error('Failed to start server: %s', e)