refactor(middleware): rename profile to proxy_profile and raise ToolError on failures

Avoids collisions with backend tool arguments by using a namespaced
proxy_profile parameter. Errors now raise ToolError instead of returning
ToolResult for proper MCP error propagation. Deep-copies tool parameters
to prevent mutating shared upstream dicts. Extracts profile override
middleware setup into a dedicated helper.

Author: benjstoll

mcp_proxy_for_aws/middleware/profile_switcher.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Middleware that enables per-call AWS profile overrides via a ``profile`` argument.
+"""Middleware that enables per-call AWS profile overrides via a ``proxy_profile`` argument.
 
-Pass ``profile`` as an extra argument on any tool call to route that single request
+Pass ``proxy_profile`` as an extra argument on any tool call to route that single request
 through a dedicated transport signed with the specified profile's credentials. The
 argument is stripped before forwarding to the backend.
 
@@ -23,12 +23,14 @@
 """
 
 import asyncio
+import copy
 import httpx
 import logging
 import mcp.types as mt
 from collections.abc import Sequence
 from fastmcp import Client
 from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
+from fastmcp.exceptions import ToolError
 from fastmcp.tools.tool import Tool, ToolResult
 from mcp_proxy_for_aws.utils import create_transport_with_sigv4
 from typing import Any, cast
@@ -39,12 +41,12 @@
 
 
 class ProfileOverrideMiddleware(Middleware):
-    """Middleware that intercepts ``profile`` on any tool call for per-request AWS identity switching.
+    """Middleware that intercepts ``proxy_profile`` on any tool call for per-request AWS identity switching.
 
-    When a tool call includes a ``profile`` argument, the middleware:
+    When a tool call includes a ``proxy_profile`` argument, the middleware:
 
     1. Validates the profile against the allowed list
-    2. Strips ``profile`` from the arguments
+    2. Strips ``proxy_profile`` from the arguments
     3. Forwards the call through a dedicated per-profile MCP client
 
     Each profile gets its own transport and session to the backend so that
@@ -79,22 +81,24 @@ async def on_list_tools(
         context: MiddlewareContext[mt.ListToolsRequest],
         call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]],
     ) -> Sequence[Tool]:
-        """Inject ``profile`` into every tool's schema."""
+        """Inject ``proxy_profile`` into every tool's schema."""
         tools = await call_next(context)
 
         for tool in tools:
-            params = tool.parameters
-            if not isinstance(params, dict):
+            if not isinstance(tool.parameters, dict):
                 continue
+            # Deep-copy to avoid mutating upstream cached/shared dicts
+            params = copy.deepcopy(tool.parameters)
             if 'properties' not in params:
                 params['properties'] = {}
-            params['properties']['profile'] = {
+            params['properties']['proxy_profile'] = {
                 'type': 'string',
                 'description': (
                     'AWS CLI profile to sign this request with. Omit to use the default profile.'
                 ),
                 'enum': sorted(self._allowed_profiles),
             }
+            tool.parameters = params
 
         return list(tools)
 
@@ -106,10 +110,10 @@ async def on_call_tool(
         context: MiddlewareContext[mt.CallToolRequestParams],
         call_next: CallNext[mt.CallToolRequestParams, ToolResult],
     ) -> ToolResult:
-        """Intercept ``profile`` and route through a dedicated per-profile client."""
+        """Intercept ``proxy_profile`` and route through a dedicated per-profile client."""
         arguments = context.message.arguments
-        if isinstance(arguments, dict) and 'profile' in arguments:
-            profile = arguments['profile']
+        if isinstance(arguments, dict) and 'proxy_profile' in arguments:
+            profile = arguments['proxy_profile']
             return await self._call_with_profile(profile, context, call_next)
 
         return await call_next(context)
@@ -157,14 +161,14 @@ async def _call_with_profile(
         """Forward a tool call through a dedicated per-profile connection."""
         if profile not in self._allowed_profiles:
             allowed = ', '.join(sorted(self._allowed_profiles))
-            return ToolResult(
-                content=f'Error: profile {profile!r} is not in the allowed list. '
+            raise ToolError(
+                f'Profile {profile!r} is not in the allowed list. '
                 f'Allowed profiles: {allowed}'
             )
 
-        # Strip profile before forwarding to the backend
+        # Strip proxy_profile before forwarding to the backend
         arguments: dict[str, Any] = dict(cast(dict[str, Any], context.message.arguments))
-        arguments.pop('profile', None)
+        arguments.pop('proxy_profile', None)
 
         logger.info(
             'Per-call profile override: routing through dedicated connection for %s', profile
@@ -174,8 +178,8 @@ async def _call_with_profile(
             client = await self._get_profile_client(profile)
         except Exception:
             logger.exception('Failed to create connection for profile %s', profile)
-            return ToolResult(
-                content=f'Error: failed to create connection for profile {profile!r}. '
+            raise ToolError(
+                f'Failed to create connection for profile {profile!r}. '
                 'Check that the profile is configured and credentials are valid.'
             )
 
@@ -188,7 +192,7 @@ async def _call_with_profile(
             )
         except Exception:
             logger.exception('Error calling tool via profile %s', profile)
-            return ToolResult(
-                content=f'Error: tool call failed using profile {profile!r}. '
+            raise ToolError(
+                f'Tool call failed using profile {profile!r}. '
                 'The request could not be completed with the specified profile.'
             )

mcp_proxy_for_aws/server.py

@@ -105,17 +105,9 @@ async def run_proxy(args) -> None:
         add_logging_middleware(proxy, args.log_level)
         add_tool_filtering_middleware(proxy, args.read_only)
 
-        allowed_profiles = getattr(args, 'allow_switch_profile', None)
-        if isinstance(allowed_profiles, list) and allowed_profiles:
-            profile_middleware = ProfileOverrideMiddleware(
-                allowed_profiles=allowed_profiles,
-                service=service,
-                region=region,
-                metadata=metadata,
-                timeout=timeout,
-                endpoint=args.endpoint,
-            )
-            proxy.add_middleware(profile_middleware)
+        profile_middleware = add_profile_override_middleware(
+            proxy, args, service, region, metadata, timeout
+        )
 
         if args.retries:
             add_retry_middleware(proxy, args.retries)
@@ -140,6 +132,43 @@ def add_tool_error_middleware(mcp: FastMCP, tool_timeout: float) -> None:
     mcp.add_middleware(ToolErrorMiddleware(tool_call_timeout=tool_timeout))
 
 
+def add_profile_override_middleware(
+    mcp: FastMCP,
+    args,
+    service: str,
+    region: str,
+    metadata: dict,
+    timeout: httpx.Timeout,
+) -> ProfileOverrideMiddleware | None:
+    """Add profile override middleware to target MCP server.
+
+    Args:
+        mcp: The FastMCP instance to add profile override to
+        args: The parsed CLI arguments
+        service: The AWS service name
+        region: The AWS region
+        metadata: The metadata dictionary
+        timeout: The httpx timeout configuration
+
+    Returns:
+        The ProfileOverrideMiddleware instance if added, None otherwise
+    """
+    allowed_profiles = getattr(args, 'allow_switch_profile', None)
+    if not isinstance(allowed_profiles, list) or not allowed_profiles:
+        return None
+    logger.info('Adding profile override middleware')
+    middleware = ProfileOverrideMiddleware(
+        allowed_profiles=allowed_profiles,
+        service=service,
+        region=region,
+        metadata=metadata,
+        timeout=timeout,
+        endpoint=args.endpoint,
+    )
+    mcp.add_middleware(middleware)
+    return middleware
+
+
 def add_tool_filtering_middleware(mcp: FastMCP, read_only: bool = False) -> None:
     """Add tool filtering middleware to target MCP server.
 

tests/unit/test_profile_switcher.py

@@ -17,6 +17,7 @@
 import asyncio
 import httpx
 import pytest
+from fastmcp.exceptions import ToolError
 from fastmcp.server.middleware import MiddlewareContext
 from mcp_proxy_for_aws.middleware.profile_switcher import ProfileOverrideMiddleware
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
@@ -51,8 +52,8 @@ class TestOnListTools:
     """Tests for the on_list_tools method."""
 
     @pytest.mark.asyncio
-    async def test_injects_profile_property_into_tool_schemas(self, middleware, mock_context):
-        """Every proxied tool gets a profile property in its schema."""
+    async def test_injects_proxy_profile_property_into_tool_schemas(self, middleware, mock_context):
+        """Every proxied tool gets a proxy_profile property in its schema."""
         tool = Mock()
         tool.name = 'some_tool'
         tool.parameters = {'type': 'object', 'properties': {'arg': {'type': 'string'}}}
@@ -62,7 +63,7 @@ async def test_injects_profile_property_into_tool_schemas(self, middleware, mock
 
         assert len(result) == 1
         assert result[0].name == 'some_tool'
-        profile_schema = result[0].parameters['properties']['profile']
+        profile_schema = result[0].parameters['properties']['proxy_profile']
         assert profile_schema['type'] == 'string'
         assert 'AWS CLI profile' in profile_schema['description']
         assert profile_schema['enum'] == sorted(ALLOWED_PROFILES)
@@ -101,15 +102,15 @@ async def test_adds_properties_key_when_missing(self, middleware, mock_context):
         result = await middleware.on_list_tools(mock_context, call_next)
 
         assert 'properties' in result[0].parameters
-        assert 'profile' in result[0].parameters['properties']
+        assert 'proxy_profile' in result[0].parameters['properties']
 
 
 class TestOnCallTool:
     """Tests for the on_call_tool method."""
 
     @pytest.mark.asyncio
-    async def test_passes_through_calls_without_profile(self, middleware, mock_context):
-        """Tool calls without profile are forwarded unchanged."""
+    async def test_passes_through_calls_without_proxy_profile(self, middleware, mock_context):
+        """Tool calls without proxy_profile are forwarded unchanged."""
         mock_context.message = Mock()
         mock_context.message.name = 'some_tool'
         mock_context.message.arguments = {'arg': 'value'}
@@ -141,20 +142,20 @@ class TestPerCallProfileOverride:
 
     @pytest.mark.asyncio
     async def test_profile_override_disallowed(self, middleware, mock_context):
-        """Profile with a disallowed profile returns an error."""
+        """Disallowed profile raises ToolError."""
         mock_context.message = Mock()
         mock_context.message.name = 'some_tool'
-        mock_context.message.arguments = {'arg': 'value', 'profile': 'evil-profile'}
+        mock_context.message.arguments = {'arg': 'value', 'proxy_profile': 'evil-profile'}
         call_next = AsyncMock()
 
-        result = await middleware.on_call_tool(mock_context, call_next)
+        with pytest.raises(ToolError, match='not in the allowed list'):
+            await middleware.on_call_tool(mock_context, call_next)
 
-        assert 'not in the allowed list' in result.content[0].text
         call_next.assert_not_called()
 
     @pytest.mark.asyncio
-    async def test_profile_override_strips_profile_arg(self, middleware, mock_context):
-        """Profile is stripped before forwarding to the backend."""
+    async def test_profile_override_strips_proxy_profile_arg(self, middleware, mock_context):
+        """proxy_profile is stripped before forwarding to the backend."""
         mock_client = AsyncMock()
         mock_call_result = MagicMock()
         mock_call_result.content = 'result'
@@ -164,7 +165,7 @@ async def test_profile_override_strips_profile_arg(self, middleware, mock_contex
 
         mock_context.message = Mock()
         mock_context.message.name = 'some_tool'
-        mock_context.message.arguments = {'arg': 'value', 'profile': 'dev-profile'}
+        mock_context.message.arguments = {'arg': 'value', 'proxy_profile': 'dev-profile'}
         call_next = AsyncMock()
 
         with patch.object(middleware, '_get_profile_client', return_value=mock_client):
@@ -175,36 +176,36 @@ async def test_profile_override_strips_profile_arg(self, middleware, mock_contex
 
     @pytest.mark.asyncio
     async def test_profile_override_connection_failure(self, middleware, mock_context):
-        """Connection failure returns a sanitized error."""
+        """Connection failure raises ToolError with sanitized message."""
         mock_context.message = Mock()
         mock_context.message.name = 'some_tool'
-        mock_context.message.arguments = {'arg': 'value', 'profile': 'dev-profile'}
+        mock_context.message.arguments = {'arg': 'value', 'proxy_profile': 'dev-profile'}
         call_next = AsyncMock()
 
         with patch.object(
             middleware, '_get_profile_client', side_effect=Exception('connection refused')
         ):
-            result = await middleware.on_call_tool(mock_context, call_next)
+            with pytest.raises(ToolError, match='Failed to create connection') as exc_info:
+                await middleware.on_call_tool(mock_context, call_next)
 
-        assert 'failed to create connection' in result.content[0].text
-        assert 'connection refused' not in result.content[0].text
+        assert 'connection refused' not in str(exc_info.value)
 
     @pytest.mark.asyncio
     async def test_profile_override_tool_call_failure(self, middleware, mock_context):
-        """Tool call failure returns a sanitized error."""
+        """Tool call failure raises ToolError with sanitized message."""
         mock_client = AsyncMock()
         mock_client.call_tool.side_effect = Exception('backend error')
 
         mock_context.message = Mock()
         mock_context.message.name = 'some_tool'
-        mock_context.message.arguments = {'arg': 'value', 'profile': 'dev-profile'}
+        mock_context.message.arguments = {'arg': 'value', 'proxy_profile': 'dev-profile'}
         call_next = AsyncMock()
 
         with patch.object(middleware, '_get_profile_client', return_value=mock_client):
-            result = await middleware.on_call_tool(mock_context, call_next)
+            with pytest.raises(ToolError, match='Tool call failed') as exc_info:
+                await middleware.on_call_tool(mock_context, call_next)
 
-        assert 'tool call failed' in result.content[0].text
-        assert 'backend error' not in result.content[0].text
+        assert 'backend error' not in str(exc_info.value)
 
 
 class TestGetProfileClient: