feat: replace forwarding region with metadata forwarding

Kyon Caldera · Kyon Caldera · commit 40c4c7419b4d · 2025-11-05T13:29:43.000+01:00
diff --git a/README.md b/README.md
@@ -47,6 +47,7 @@ docker build -t mcp-proxy-for-aws .
 | `--service`	         | AWS service name for SigV4 signing	                                                                                                                                                                                                     | Inferred from endpoint if not provided	                                     |No	|
 | `--profile`	         | AWS profile for AWS credentials to use	                                                                                                                                                                                                 | Uses `AWS_PROFILE` environment variable if not set                          |No	|
 | `--region`	          | AWS region to use	                                                                                                                                                                                                                      | Uses `AWS_REGION` environment variable if not set, defaults to `us-east-1`	 |No	|
+| `--metadata`	        | Metadata to inject into MCP requests as key=value pairs (e.g., `--metadata KEY1=value1 KEY2=value2`)                                                                                                                                    | `AWS_REGION` is automatically injected based on `--region` if not provided    |No	|
 | `--read-only`	       | Disable tools which may require write permissions (tools which DO NOT require write permissions are annotated with [`readOnlyHint=true`](https://modelcontextprotocol.io/specification/2025-06-18/schema#toolannotations-readonlyhint)) | `False`	                                                                    |No	|
 | `--retries`          | Configures number of retries done when calling upstream services, setting this to 0 disables retries.                                                                                                                                   | 0                                                                           |No |
 | `--log-level`	       | Set the logging level (`DEBUG/INFO/WARNING/ERROR/CRITICAL`)	                                                                                                                                                                            | `INFO`	                                                                     |No	|
diff --git a/mcp_proxy_for_aws/cli.py b/mcp_proxy_for_aws/cli.py
@@ -18,6 +18,38 @@
 import os
 from mcp_proxy_for_aws import __version__
 from mcp_proxy_for_aws.utils import within_range
+from typing import Dict, Optional, Sequence
+
+
+class KeyValueAction(argparse.Action):
+    """Custom argparse action to parse key=value pairs into a dictionary."""
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: str | Sequence[str],
+        option_string: Optional[str] = None,
+    ) -> None:
+        """Parse key=value pairs into a dictionary.
+
+        Args:
+            parser: The argument parser
+            namespace: The namespace object to update
+            values: The values to parse (list of key=value strings)
+            option_string: The option string that triggered this action
+        """
+        metadata: Dict[str, str] = {}
+        # Ensure values is a sequence
+        if isinstance(values, str):
+            values = [values]
+
+        for item in values:
+            if '=' not in item:
+                parser.error(f'Metadata must be in key=value format, got: {item}')
+            key, value = item.split('=', 1)
+            metadata[key] = value
+        setattr(namespace, self.dest, metadata)
 
 
 def parse_args():
@@ -65,9 +97,11 @@ def parse_args():
     )
 
     parser.add_argument(
-        '--forwarding-region',
-        help='AWS region to forward to server (uses --region if not provided)',
+        '--metadata',
+        nargs='*',
+        action=KeyValueAction,
         default=None,
+        help='Metadata to inject into MCP requests as key=value pairs (e.g., --metadata AWS_REGION=us-west-2 KEY=VALUE)',
     )
 
     parser.add_argument(
diff --git a/mcp_proxy_for_aws/server.py b/mcp_proxy_for_aws/server.py
@@ -52,9 +52,13 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
 
     # Validate and determine region
     region = determine_aws_region(args.endpoint, args.region)
-    forwarding_region = args.forwarding_region or region
     logger.debug('Using region: %s', region)
 
+    # Build metadata dictionary - start with AWS_REGION, then merge user metadata
+    metadata = {'AWS_REGION': region}
+    if args.metadata:
+        metadata.update(args.metadata)
+
     # Get profile
     profile = args.profile
 
@@ -63,7 +67,7 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
         'Using service: %s, region: %s, forwarding region: %s, profile: %s',
         service,
         region,
-        forwarding_region,
+        metadata.get('AWS_REGION'),
         profile,
     )
     logger.info('Running in MCP mode')
@@ -77,7 +81,7 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
 
     # Create transport with SigV4 authentication
     transport = create_transport_with_sigv4(
-        args.endpoint, service, region, forwarding_region, timeout, profile
+        args.endpoint, service, region, metadata, timeout, profile
     )
 
     # Create proxy with the transport
diff --git a/mcp_proxy_for_aws/utils.py b/mcp_proxy_for_aws/utils.py
@@ -21,7 +21,7 @@
 import re
 from fastmcp.client.transports import StreamableHttpTransport
 from mcp_proxy_for_aws.sigv4_helper import create_sigv4_client
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 from urllib.parse import urlparse
 
 
@@ -32,7 +32,7 @@ def create_transport_with_sigv4(
     url: str,
     service: str,
     region: str,
-    forwarding_region: str,
+    metadata: Dict[str, Any],
     custom_timeout: httpx.Timeout,
     profile: Optional[str] = None,
 ) -> StreamableHttpTransport:
@@ -42,7 +42,7 @@ def create_transport_with_sigv4(
         url: The endpoint URL
         service: AWS service name for SigV4 signing
         region: AWS region to use
-        forwarding_region: AWS region to forward to server
+        metadata: Metadata dictionary to inject into MCP requests
         custom_timeout: httpx.Timeout used to connect to the endpoint
         profile: AWS profile to use (optional)
 
@@ -62,7 +62,7 @@ def client_factory(
             region=region,
             headers=headers,
             timeout=custom_timeout,
-            metadata={'AWS_REGION': forwarding_region},
+            metadata=metadata,
             auth=auth,
         )
 
diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py
@@ -140,3 +140,39 @@ def test_parse_args_negative_timeout(self):
         """Test parsing fails with negative timeout (within_range validation)."""
         with pytest.raises(SystemExit):
             parse_args()
+
+    @patch(
+        'sys.argv',
+        ['mcp-proxy-for-aws', 'https://example.com', '--metadata', 'KEY1=value1', 'KEY2=value2'],
+    )
+    def test_parse_metadata_argument(self):
+        """Test parsing metadata key=value pairs."""
+        args = parse_args()
+        assert args.metadata == {'KEY1': 'value1', 'KEY2': 'value2'}
+
+    @patch(
+        'sys.argv',
+        ['mcp-proxy-for-aws', 'https://example.com', '--metadata', 'AWS_REGION=us-west-2'],
+    )
+    def test_parse_metadata_single_pair(self):
+        """Test parsing single metadata key=value pair."""
+        args = parse_args()
+        assert args.metadata == {'AWS_REGION': 'us-west-2'}
+
+    @patch(
+        'sys.argv',
+        ['mcp-proxy-for-aws', 'https://example.com', '--metadata', 'KEY=value with spaces'],
+    )
+    def test_parse_metadata_with_spaces_in_value(self):
+        """Test parsing metadata with spaces in value."""
+        args = parse_args()
+        assert args.metadata == {'KEY': 'value with spaces'}
+
+    @patch('sys.argv', ['mcp-proxy-for-aws', 'https://example.com', '--metadata', 'INVALID'])
+    def test_parse_metadata_invalid_format(self):
+        """Test that invalid metadata format raises an error."""
+        import argparse
+
+        with pytest.raises((SystemExit, argparse.ArgumentTypeError)):
+            # argparse may call sys.exit or raise ArgumentTypeError
+            parse_args()
diff --git a/tests/unit/test_server.py b/tests/unit/test_server.py
@@ -56,7 +56,7 @@ async def test_setup_mcp_mode(
         mock_args.profile = None
         mock_args.read_only = True
         mock_args.retries = 1
-        mock_args.forwarding_region = None
+        mock_args.metadata = None
         # Add timeout parameters
         mock_args.timeout = 180.0
         mock_args.connect_timeout = 60.0
@@ -87,7 +87,7 @@ async def test_setup_mcp_mode(
         assert call_args[0][0] == 'https://test.example.com'
         assert call_args[0][1] == 'test-service'
         assert call_args[0][2] == 'us-east-1'
-        assert call_args[0][3] == 'us-east-1'  # forwarding_region (defaults to region)
+        assert call_args[0][3] == {'AWS_REGION': 'us-east-1'}  # metadata
         # call_args[0][4] is the Timeout object
         assert call_args[0][5] is None  # profile
         mock_as_proxy.assert_called_once_with(mock_transport)
@@ -118,7 +118,7 @@ async def test_setup_mcp_mode_no_retries(
         mock_args.profile = 'test-profile'
         mock_args.read_only = False
         mock_args.retries = 0  # No retries
-        mock_args.forwarding_region = 'eu-west-1'
+        mock_args.metadata = {'AWS_REGION': 'eu-west-1', 'CUSTOM_KEY': 'custom_value'}
         # Add timeout parameters
         mock_args.timeout = 180.0
         mock_args.connect_timeout = 60.0
@@ -149,13 +149,116 @@ async def test_setup_mcp_mode_no_retries(
         assert call_args[0][0] == 'https://test.example.com'
         assert call_args[0][1] == 'test-service'
         assert call_args[0][2] == 'us-east-1'
-        assert call_args[0][3] == 'eu-west-1'  # forwarding_region
+        assert call_args[0][3] == {
+            'AWS_REGION': 'eu-west-1',
+            'CUSTOM_KEY': 'custom_value',
+        }  # metadata
         # call_args[0][4] is the Timeout object
         assert call_args[0][5] == 'test-profile'  # profile
         mock_as_proxy.assert_called_once_with(mock_transport)
         mock_add_filtering.assert_called_once_with(mock_proxy, False)
         mock_proxy.run_async.assert_called_once()
 
+    @patch('mcp_proxy_for_aws.server.create_transport_with_sigv4')
+    @patch('mcp_proxy_for_aws.server.FastMCP.as_proxy')
+    @patch('mcp_proxy_for_aws.server.determine_aws_region')
+    @patch('mcp_proxy_for_aws.server.determine_service_name')
+    @patch('mcp_proxy_for_aws.server.add_tool_filtering_middleware')
+    async def test_setup_mcp_mode_no_metadata_injects_aws_region(
+        self,
+        mock_add_filtering,
+        mock_determine_service,
+        mock_determine_region,
+        mock_as_proxy,
+        mock_create_transport,
+    ):
+        """Test that AWS_REGION is automatically injected when no metadata is provided."""
+        # Arrange
+        local_mcp = Mock(spec=FastMCP)
+        mock_args = Mock()
+        mock_args.endpoint = 'https://test.example.com'
+        mock_args.service = 'test-service'
+        mock_args.region = 'ap-southeast-1'
+        mock_args.profile = None
+        mock_args.read_only = False
+        mock_args.retries = 0
+        mock_args.metadata = None  # No metadata provided
+        mock_args.timeout = 180.0
+        mock_args.connect_timeout = 60.0
+        mock_args.read_timeout = 120.0
+        mock_args.write_timeout = 180.0
+        mock_args.log_level = 'INFO'
+
+        mock_determine_service.return_value = 'test-service'
+        mock_determine_region.return_value = 'ap-southeast-1'
+
+        mock_transport = Mock()
+        mock_create_transport.return_value = mock_transport
+        mock_proxy = Mock()
+        mock_proxy.run_async = AsyncMock()
+        mock_as_proxy.return_value = mock_proxy
+
+        # Act
+        await setup_mcp_mode(local_mcp, mock_args)
+
+        # Assert - verify AWS_REGION was automatically injected
+        assert mock_create_transport.call_count == 1
+        call_args = mock_create_transport.call_args
+        metadata = call_args[0][3]
+        assert metadata == {'AWS_REGION': 'ap-southeast-1'}
+
+    @patch('mcp_proxy_for_aws.server.create_transport_with_sigv4')
+    @patch('mcp_proxy_for_aws.server.FastMCP.as_proxy')
+    @patch('mcp_proxy_for_aws.server.determine_aws_region')
+    @patch('mcp_proxy_for_aws.server.determine_service_name')
+    @patch('mcp_proxy_for_aws.server.add_tool_filtering_middleware')
+    async def test_setup_mcp_mode_metadata_without_aws_region_injects_it(
+        self,
+        mock_add_filtering,
+        mock_determine_service,
+        mock_determine_region,
+        mock_as_proxy,
+        mock_create_transport,
+    ):
+        """Test that AWS_REGION is injected even when other metadata is provided."""
+        # Arrange
+        local_mcp = Mock(spec=FastMCP)
+        mock_args = Mock()
+        mock_args.endpoint = 'https://test.example.com'
+        mock_args.service = 'test-service'
+        mock_args.region = 'us-west-1'
+        mock_args.profile = None
+        mock_args.read_only = False
+        mock_args.retries = 0
+        mock_args.metadata = {'CUSTOM_KEY': 'custom_value', 'ANOTHER_KEY': 'another_value'}
+        mock_args.timeout = 180.0
+        mock_args.connect_timeout = 60.0
+        mock_args.read_timeout = 120.0
+        mock_args.write_timeout = 180.0
+        mock_args.log_level = 'INFO'
+
+        mock_determine_service.return_value = 'test-service'
+        mock_determine_region.return_value = 'us-west-1'
+
+        mock_transport = Mock()
+        mock_create_transport.return_value = mock_transport
+        mock_proxy = Mock()
+        mock_proxy.run_async = AsyncMock()
+        mock_as_proxy.return_value = mock_proxy
+
+        # Act
+        await setup_mcp_mode(local_mcp, mock_args)
+
+        # Assert - verify AWS_REGION was injected along with custom metadata
+        assert mock_create_transport.call_count == 1
+        call_args = mock_create_transport.call_args
+        metadata = call_args[0][3]
+        assert metadata == {
+            'AWS_REGION': 'us-west-1',
+            'CUSTOM_KEY': 'custom_value',
+            'ANOTHER_KEY': 'another_value',
+        }
+
     def test_add_tool_filtering_middleware(self):
         """Test that tool filtering middleware is added correctly."""
         # Arrange
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -39,11 +39,11 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client):
         service = 'test-service'
         profile = 'test-profile'
         region = 'us-east-1'
-        forwarding_region = 'us-west-2'
+        metadata = {'AWS_REGION': 'us-west-2', 'CUSTOM_KEY': 'custom_value'}
         custom_timeout = Timeout(30.0)
 
         result = create_transport_with_sigv4(
-            url, service, region, forwarding_region, custom_timeout, profile
+            url, service, region, metadata, custom_timeout, profile
         )
 
         # Verify result is StreamableHttpTransport
@@ -64,7 +64,7 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client):
                 headers={'test': 'header'},
                 timeout=custom_timeout,
                 auth=None,
-                metadata={'AWS_REGION': forwarding_region},
+                metadata=metadata,
             )
         else:
             # If we can't access the factory directly, just verify the transport was created
@@ -78,12 +78,10 @@ def test_create_transport_with_sigv4_no_profile(self, mock_create_sigv4_client):
         url = 'https://test-service.us-west-2.api.aws/mcp'
         service = 'test-service'
         region = 'test-region'
-        forwarding_region = 'test-forwarding-region'
+        metadata = {'AWS_REGION': 'test-forwarding-region'}
         custom_timeout = Timeout(60.0)
 
-        result = create_transport_with_sigv4(
-            url, service, region, forwarding_region, custom_timeout
-        )
+        result = create_transport_with_sigv4(url, service, region, metadata, custom_timeout)
 
         # Test that the httpx_client_factory calls create_sigv4_client correctly
         # We need to access the factory through the transport's internal structure
@@ -98,7 +96,7 @@ def test_create_transport_with_sigv4_no_profile(self, mock_create_sigv4_client):
                 headers=None,
                 timeout=custom_timeout,
                 auth=None,
-                metadata={'AWS_REGION': forwarding_region},
+                metadata=metadata,
             )
         else:
             # If we can't access the factory directly, just verify the transport was created
