Adding support for configuring a KMS key when creating a Gateway

documentation/docs/user-guide/gateway/quickstart.md

@@ -78,6 +78,8 @@ def setup_gateway():
         authorizer_config=cognito_response["authorizer_config"],
         # enable semantic search
         enable_semantic_search=True,
+        # optional KMS key ARN for encrypting data stored on Gateway
+        kms_key_arn=None,
     )
     print(f"✓ Gateway created: {gateway['gatewayUrl']}\n")
 

src/bedrock_agentcore_starter_toolkit/cli/gateway/commands.py

@@ -19,6 +19,7 @@ def create_mcp_gateway(
     role_arn: Optional[str] = None,
     authorizer_config: Optional[str] = None,
     enable_semantic_search: Optional[bool] = typer.Option(True, "--enable_semantic_search", "-sem"),
+    kms_key_arn: Optional[str] = None,
 ) -> None:
     """Creates an MCP Gateway.
 
@@ -27,13 +28,14 @@ def create_mcp_gateway(
     :param role_arn: optional - the role arn to use (creates one if none provided).
     :param authorizer_config: optional - the serialized authorizer config (will create one if none provided).
     :param enable_semantic_search: optional - whether to enable search tool (defaults to True).
+    :param kms_key_arn: optional - the KMS key ARN to use for encryption.
     :return:
     """
     client = GatewayClient(region_name=region)
     json_authorizer_config = ""
     if authorizer_config:
         json_authorizer_config = json.loads(authorizer_config)
-    gateway = client.create_mcp_gateway(name, role_arn, json_authorizer_config, enable_semantic_search)
+    gateway = client.create_mcp_gateway(name, role_arn, json_authorizer_config, enable_semantic_search, kms_key_arn)
     console.print(gateway)
 
 

src/bedrock_agentcore_starter_toolkit/operations/gateway/client.py

@@ -58,6 +58,7 @@ def create_mcp_gateway(
         role_arn=None,
         authorizer_config=None,
         enable_semantic_search=True,
+        kms_key_arn=None,
     ) -> dict:
         """Creates an MCP Gateway.
 
@@ -88,6 +89,8 @@ def create_mcp_gateway(
         }
         if enable_semantic_search:
             create_request["protocolConfiguration"] = {"mcp": {"searchType": "SEMANTIC"}}
+        if kms_key_arn:
+            create_request["kmsKeyArn"] = kms_key_arn
         self.logger.info("Creating Gateway")
         self.logger.debug("Creating gateway with params: %s", json.dumps(create_request, indent=2))
         gateway = self.client.create_gateway(**create_request)

tests/cli/gateway/test_commands.py

@@ -59,6 +59,7 @@ def test_create_mcp_gateway_command_basic(self):
                 "arn:aws:iam::123456789012:role/TestGatewayRole",
                 "",  # empty authorizer config
                 True,  # enable_semantic_search default
+                None,  # kms_key_arn default
             )
 
     def test_create_mcp_gateway_with_defaults(self):
@@ -84,6 +85,56 @@ def test_create_mcp_gateway_with_defaults(self):
                 None,  # role_arn default
                 "",  # empty authorizer config
                 True,  # enable_semantic_search default
+                None,  # kms_key_arn default
+            )
+
+    def test_create_mcp_gateway_with_kms_key_arn(self):
+        """Test create_mcp_gateway command with KMS key ARN."""
+        with patch("bedrock_agentcore_starter_toolkit.cli.gateway.commands.GatewayClient") as mock_gateway_client:
+            mock_client_instance = Mock()
+            mock_gateway_client.return_value = mock_client_instance
+
+            mock_gateway_response = {
+                "gatewayId": "test-gateway-kms",
+                "gatewayArn": "arn:aws:bedrock-agentcore:us-west-2:123456789012:gateway/test-gateway-kms",
+                "gatewayUrl": "https://test-gateway-kms.us-west-2.amazonaws.com",
+                "status": "CREATING",
+                "name": "TestGatewayWithKMS",
+                "roleArn": "arn:aws:iam::123456789012:role/TestGatewayRole",
+                "kmsKeyArn": "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012",
+            }
+            mock_client_instance.create_mcp_gateway.return_value = mock_gateway_response
+
+            # Test the command with KMS key ARN
+            kms_key_arn = "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012"
+            result = self.runner.invoke(
+                gateway_app,
+                [
+                    "create-mcp-gateway",
+                    "--region",
+                    "us-west-2",
+                    "--name",
+                    "TestGatewayWithKMS",
+                    "--role-arn",
+                    "arn:aws:iam::123456789012:role/TestGatewayRole",
+                    "--kms-key-arn",
+                    kms_key_arn,
+                ],
+            )
+
+            # Verify the command executed successfully
+            assert result.exit_code == 0
+
+            # Verify GatewayClient was initialized with correct region
+            mock_gateway_client.assert_called_once_with(region_name="us-west-2")
+
+            # Verify create_mcp_gateway was called with correct parameters including KMS key ARN
+            mock_client_instance.create_mcp_gateway.assert_called_once_with(
+                "TestGatewayWithKMS",
+                "arn:aws:iam::123456789012:role/TestGatewayRole",
+                "",  # empty authorizer config
+                True,  # enable_semantic_search default
+                kms_key_arn,  # kms_key_arn provided
             )
 
     def test_create_mcp_gateway_target_command_basic(self):

tests/operations/gateway/test_gateway_client.py

@@ -402,3 +402,123 @@ def test_wait_for_ready_error(self, gateway_client):
             )
 
         assert "failed" in str(excinfo.value).lower()
+
+    @patch("bedrock_agentcore_starter_toolkit.operations.gateway.GatewayClient.create_oauth_authorizer_with_cognito")
+    def test_create_mcp_gateway_with_kms_key_arn(self, mock_create_oauth_authorizer_with_cognito, gateway_client):
+        """Test creating MCP gateway with KMS key ARN"""
+        # Mock responses
+        mock_bedrock = Mock()
+        gateway_client.client = mock_bedrock
+
+        mock_create_oauth_authorizer_with_cognito.return_value = {
+            "authorizer_config": {
+                "customJWTAuthorizer": {"allowedClients": ["allowedClient"], "discoveryUrl": "aRandomUrl"}
+            },
+            "client_info": {
+                "client_id": "client",
+                "client_secret": "clientSecret",
+                "user_pool_id": "poolId",
+                "token_endpoint": "tokenEndpoint",
+                "scope": "my-gateway/invoke",
+                "domain_prefix": "some-prefix",
+            },
+        }
+
+        # Mock gateway creation
+        mock_bedrock.create_gateway.return_value = {
+            "gatewayId": "TEST123",
+            "gatewayArn": "arn:aws:bedrock_agentcore:us-west-2:123:gateway/TEST123",
+            "gatewayUrl": "https://TEST456.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp",
+            "status": "READY",
+            "roleArn": "roleArn",
+        }
+
+        # Mock get operations for status checking
+        mock_bedrock.get_gateway.return_value = {
+            "gatewayId": "TEST456",
+            "gatewayArn": "arn:aws:bedrock-agentcore:us-west-2:gateway/TEST456",
+            "gatewayUrl": "https://TEST456.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp",
+            "status": "READY",
+        }
+
+        # Test with KMS key ARN
+        kms_key_arn = "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012"
+        gateway_client.create_mcp_gateway(
+            name="test-gateway-with-kms",
+            role_arn="arn:aws:iam::123:role/TestRole",
+            kms_key_arn=kms_key_arn,
+        )
+
+        # Verify the create_gateway call includes the KMS key ARN
+        mock_bedrock.create_gateway.assert_called_once()
+        call_args = mock_bedrock.create_gateway.call_args[1]
+
+        assert "kmsKeyArn" in call_args
+        assert call_args["kmsKeyArn"] == kms_key_arn
+
+        # Verify other expected parameters are still present
+        assert call_args["name"] == "test-gateway-with-kms"
+        assert call_args["roleArn"] == "arn:aws:iam::123:role/TestRole"
+        assert call_args["protocolType"] == "MCP"
+        assert call_args["authorizerType"] == "CUSTOM_JWT"
+        assert "authorizerConfiguration" in call_args
+        assert call_args["exceptionLevel"] == "DEBUG"
+
+    @patch("bedrock_agentcore_starter_toolkit.operations.gateway.GatewayClient.create_oauth_authorizer_with_cognito")
+    def test_create_mcp_gateway_without_kms_key_arn(self, mock_create_oauth_authorizer_with_cognito, gateway_client):
+        """Test creating MCP gateway without KMS key ARN (default behavior)"""
+        # Mock responses
+        mock_bedrock = Mock()
+        gateway_client.client = mock_bedrock
+
+        mock_create_oauth_authorizer_with_cognito.return_value = {
+            "authorizer_config": {
+                "customJWTAuthorizer": {"allowedClients": ["allowedClient"], "discoveryUrl": "aRandomUrl"}
+            },
+            "client_info": {
+                "client_id": "client",
+                "client_secret": "clientSecret",
+                "user_pool_id": "poolId",
+                "token_endpoint": "tokenEndpoint",
+                "scope": "my-gateway/invoke",
+                "domain_prefix": "some-prefix",
+            },
+        }
+
+        # Mock gateway creation
+        mock_bedrock.create_gateway.return_value = {
+            "gatewayId": "TEST123",
+            "gatewayArn": "arn:aws:bedrock_agentcore:us-west-2:123:gateway/TEST123",
+            "gatewayUrl": "https://TEST456.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp",
+            "status": "READY",
+            "roleArn": "roleArn",
+        }
+
+        # Mock get operations for status checking
+        mock_bedrock.get_gateway.return_value = {
+            "gatewayId": "TEST456",
+            "gatewayArn": "arn:aws:bedrock-agentcore:us-west-2:gateway/TEST456",
+            "gatewayUrl": "https://TEST456.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp",
+            "status": "READY",
+        }
+
+        # Test without KMS key ARN (default behavior)
+        gateway_client.create_mcp_gateway(
+            name="test-gateway-no-kms",
+            role_arn="arn:aws:iam::123:role/TestRole",
+        )
+
+        # Verify the create_gateway call does NOT include the KMS key ARN
+        mock_bedrock.create_gateway.assert_called_once()
+        call_args = mock_bedrock.create_gateway.call_args[1]
+
+        assert "kmsKeyArn" not in call_args
+
+        # Verify other expected parameters are still present
+        assert call_args["name"] == "test-gateway-no-kms"
+        assert call_args["roleArn"] == "arn:aws:iam::123:role/TestRole"
+        assert call_args["protocolType"] == "MCP"
+        assert call_args["authorizerType"] == "CUSTOM_JWT"
+        assert "authorizerConfiguration" in call_args
+        assert call_args["exceptionLevel"] == "DEBUG"
+