Add region utils

Author: acmlau

aws_mcp_proxy/server.py

@@ -31,6 +31,7 @@
 from aws_mcp_proxy.utils import (
     create_transport_with_sigv4,
     determine_service_name,
+    determine_aws_region,
 )
 from fastmcp.server.server import FastMCP
 from typing import Any
@@ -46,9 +47,11 @@ async def setup_mcp_mode(mcp: FastMCP, args) -> None:
     # Validate and determine service
     service = determine_service_name(args.endpoint, args.service)
 
-    # Get profile and region
+    # Validate and determine region
+    region = determine_aws_region(args.endpoint, args.region)
+
+    # Get profile
     profile = args.profile
-    region = args.region
 
     # Log server configuration
     logger.info(
@@ -57,7 +60,7 @@ async def setup_mcp_mode(mcp: FastMCP, args) -> None:
     logger.info('Running in MCP mode')
 
     # Create transport with SigV4 authentication
-    transport = create_transport_with_sigv4(args.endpoint, service, profile)
+    transport = create_transport_with_sigv4(args.endpoint, service, profile, region)
 
     # Create proxy with the transport
     proxy = FastMCP.as_proxy(transport)
@@ -104,7 +107,6 @@ def parse_args():
     parser.add_argument(
         '--region',
         help='AWS region to use (uses AWS_REGION if not provided)',
-        default=os.getenv('AWS_REGION', 'us-east-1'),
     )
 
     parser.add_argument(

aws_mcp_proxy/sigv4_helper.py

@@ -190,7 +190,6 @@ def create_sigv4_client(
     profile: Optional[str] = None,
     region: Optional[str] = None,
     headers: Optional[Dict[str, str]] = None,
-    auth: Optional[httpx.Auth] = None,
     **kwargs: Any,
 ) -> httpx.AsyncClient:
     """Create an httpx.AsyncClient with SigV4 authentication.

aws_mcp_proxy/utils.py

@@ -23,7 +23,7 @@
 
 
 def create_transport_with_sigv4(
-    url: str, service: str, profile: Optional[str] = None
+    url: str, service: str, profile: Optional[str] = None, region: Optional[str] = None,
 ) -> StreamableHttpTransport:
     """Create a StreamableHttpTransport with SigV4 authentication.
 
@@ -42,7 +42,7 @@ def client_factory(
         auth: Optional[httpx.Auth] = None,
     ) -> httpx.AsyncClient:
         return create_sigv4_client(
-            service=service, profile=profile, headers=headers, timeout=timeout, auth=auth
+            service=service, profile=profile, region=region, headers=headers, timeout=timeout, auth=auth
         )
 
     return StreamableHttpTransport(
@@ -81,3 +81,34 @@ def determine_service_name(endpoint: str, service: Optional[str] = None) -> str:
             'Please provide the service name explicitly using --service argument.'
         )
     return determined_service
+
+def determine_aws_region(endpoint: str, region: Optional[str] = None) -> str:
+    """Validate and determine the AWS region.
+
+    Args:
+        endpoint: The endpoint URL
+        region: Optional region name
+
+    Returns:
+        Validated AWS region
+
+    Raises:
+        ValueError: If region cannot be determined
+    """
+    if region:
+        return region
+
+    # Parse AWS region from endpoint URL
+    parsed = urlparse(endpoint)
+    hostname = parsed.hostname or ''
+
+    # Extract region name (pattern: service.region.api.aws or service-name.region.api.aws)
+    region_match = re.search(r'\.([a-z0-9-]+)\.api\.aws', hostname)
+    determined_region = region_match.group(1) if region_match else None
+
+    if not determined_region:
+        raise ValueError(
+            f"Could not determine AWS region from endpoint '{endpoint}'. "
+            'Please provide the region explicitly using --region argument.'
+        )
+    return determined_region

tests/test_server.py

@@ -36,6 +36,7 @@ async def test_setup_mcp_mode(
         mock_args = Mock()
         mock_args.endpoint = 'https://test.example.com'
         mock_args.service = 'test-service'
+        mock_args.region = 'us-east-1'
         mock_args.profile = None
         mock_args.allow_write = False
 
@@ -71,6 +72,7 @@ async def test_setup_mcp_mode_with_tools(
         mock_args = Mock()
         mock_args.endpoint = 'https://test.example.com'
         mock_args.service = 'test-service'
+        mock_args.region = 'us-east-1'
         mock_args.profile = None
         mock_args.allow_write = False
 
@@ -106,6 +108,7 @@ async def test_setup_mcp_mode_tool_registration_error(
         mock_args = Mock()
         mock_args.endpoint = 'https://test.example.com'
         mock_args.service = 'test-service'
+        mock_args.region = 'us-east-1'
         mock_args.profile = None
         mock_args.allow_write = False
 

tests/test_sigv4_helper.py

@@ -226,6 +226,7 @@ class TestCreateSigv4Auth:
     """Test cases for the create_sigv4_auth function."""
 
     @patch('aws_mcp_proxy.sigv4_helper.create_aws_session')
+    # @patch.dict(os.environ, {"AWS_REGION": ""})
     def test_create_sigv4_auth_default(self, mock_create_session):
         """Test creating SigV4 auth with default parameters."""
         # Mock session and credentials
@@ -237,6 +238,9 @@ def test_create_sigv4_auth_default(self, mock_create_session):
         mock_session.get_credentials.return_value = mock_credentials
         mock_create_session.return_value = mock_session
 
+        # Remove AWS_REGION to check default 
+        del os.environ["AWS_REGION"]
+
         # Test auth creation
         result = create_sigv4_auth('test-service')
 

tests/test_utils.py

@@ -18,6 +18,7 @@
 from aws_mcp_proxy.utils import (
     create_transport_with_sigv4,
     determine_service_name,
+    determine_aws_region,
 )
 from fastmcp.client.transports import StreamableHttpTransport
 from unittest.mock import MagicMock, patch
@@ -35,8 +36,9 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client):
         url = 'https://eks-mcp.us-west-2.api.aws/mcp'
         service = 'eks-mcp'
         profile = 'test-profile'
+        region = 'us-east-1'
 
-        result = create_transport_with_sigv4(url, service, profile)
+        result = create_transport_with_sigv4(url, service, profile, region)
 
         # Verify result is StreamableHttpTransport
         assert isinstance(result, StreamableHttpTransport)
@@ -54,6 +56,7 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client):
             mock_create_sigv4_client.assert_called_once_with(
                 service=service,
                 profile=profile,
+                region=region,
                 headers={'test': 'header'},
                 timeout=Timeout(30.0),
                 auth=None,
@@ -77,7 +80,7 @@ def test_create_transport_with_sigv4_no_profile(self, mock_create_sigv4_client):
             factory(headers=None, timeout=None, auth=None)
 
             mock_create_sigv4_client.assert_called_once_with(
-                service=service, profile=None, headers=None, timeout=None, auth=None
+                service=service, profile=None, region=None, headers=None, timeout=None, auth=None
             )
         else:
             # If we can't access the factory directly, just verify the transport was created
@@ -144,3 +147,45 @@ def test_validate_service_name_invalid_url_failure(self):
         assert 'Could not determine AWS service name' in str(exc_info.value)
         assert endpoint in str(exc_info.value)
         assert '--service argument' in str(exc_info.value)
+
+
+class TestDetermineRegion:
+    """Test cases for determine_aws_region function."""
+
+    def test_determine_region_with_region(self):
+        """Test determination when region is provided."""
+        endpoint = 'https://mcp.us-east-1.api.aws/mcp'
+        region = 'custom-region'
+
+        result = determine_aws_region(endpoint, region)
+
+        assert result == region
+
+    def test_determine_region_without_region_success(self):
+        """Test determination when region is not provided but can be parsed."""
+        endpoint = 'https://mcp.us-east-1.api.aws/mcp'
+        expected_region = 'us-east-1'
+
+        result = determine_aws_region(endpoint)
+
+        assert result == expected_region
+
+    def test_determine_region_with_complex_service_name(self):
+        """Test parsing region from endpoint with complex service name."""
+        endpoint = 'https://eks-mcp-beta.us-west-2.api.aws/mcp'
+        expected_region = 'us-west-2'
+
+        result = determine_aws_region(endpoint)
+
+        assert result == expected_region
+
+    def test_determine_region_without_region_failure(self):
+        """Test determination when region cannot be determined."""
+        endpoint = 'https://service.example.com'
+
+        with pytest.raises(ValueError) as exc_info:
+            determine_aws_region(endpoint)
+
+        assert 'Could not determine AWS region' in str(exc_info.value)
+        assert endpoint in str(exc_info.value)
+        assert '--region argument' in str(exc_info.value)