Manage MCP Servers

Makefile

@@ -378,4 +378,4 @@ test-coverage:
           --cov-report term-missing \
           --cov-report html:build/coverage \
           --cov-report xml:build/coverage/coverage.xml \
-          --cov-fail-under 42
+          --cov-fail-under 50

lambda/mcp_server/__init__.py

@@ -0,0 +1,13 @@
+#   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.

lambda/mcp_server/lambda_functions.py

@@ -0,0 +1,170 @@
+#   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.
+
+"""Lambda functions for managing MCP Servers in AWS DynamoDB."""
+import json
+import logging
+import os
+from decimal import Decimal
+from typing import Any, Dict, Optional
+
+import boto3
+from boto3.dynamodb.conditions import Attr, Key
+from utilities.common_functions import api_wrapper, get_item, get_username, is_admin, retry_config
+
+from .models import McpServerModel
+
+logger = logging.getLogger(__name__)
+
+# Initialize the DynamoDB resource and the table using environment variables
+dynamodb = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION"], config=retry_config)
+table = dynamodb.Table(os.environ["MCP_SERVERS_TABLE_NAME"])
+
+
+def _get_mcp_servers(
+    user_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Helper function to retrieve mcp servers from DynamoDB."""
+    filter_expression = None
+
+    # Filter by user_id if provided
+    if user_id:
+        condition = Attr("owner").eq(user_id) | Attr("owner").eq("lisa:public")
+        filter_expression = condition if filter_expression is None else filter_expression & condition
+
+    scan_arguments = {
+        "TableName": os.environ["MCP_SERVERS_TABLE_NAME"],
+        "IndexName": os.environ["MCP_SERVERS_BY_OWNER_INDEX_NAME"],
+    }
+
+    # Set FilterExpression if applicable
+    if filter_expression:
+        scan_arguments["FilterExpression"] = filter_expression
+
+    # Scan the DynamoDB table to retrieve items
+    items = []
+    while True:
+        response = table.scan(**scan_arguments)
+        items.extend(response.get("Items", []))
+        if "LastEvaluatedKey" in response:
+            scan_arguments["ExclusiveStartKey"] = response["LastEvaluatedKey"]
+        else:
+            break
+
+    return {"Items": items}
+
+
+@api_wrapper
+def get(event: dict, context: dict) -> Any:
+    """Retrieve a specific mcp server from DynamoDB."""
+    user_id = get_username(event)
+    mcp_server_id = get_mcp_server_id(event)
+
+    # Query for the mcp server
+    response = table.query(KeyConditionExpression=Key("id").eq(mcp_server_id), Limit=1, ScanIndexForward=False)
+    item = get_item(response)
+
+    if item is None:
+        raise ValueError(f"MCP Server {mcp_server_id} not found.")
+
+    # Check if the user is authorized to get the mcp server
+    is_owner = item["owner"] == user_id or item["owner"] == "lisa:public"
+    if is_owner or is_admin(event):
+        # add extra attribute so the frontend doesn't have to determine this
+        if is_owner:
+            item["isOwner"] = True
+        return item
+
+    raise ValueError(f"Not authorized to get {mcp_server_id}.")
+
+
+@api_wrapper
+def list(event: dict, context: dict) -> Dict[str, Any]:
+    """List mcp servers for a user from DynamoDB."""
+    user_id = get_username(event)
+
+    if is_admin(event):
+        logger.info(f"Listing all mcp servers for user {user_id} (is_admin)")
+        return _get_mcp_servers()
+    else:
+        logger.info(f"Listing mcp servers for user {user_id}")
+    return _get_mcp_servers(user_id=user_id)
+
+
+@api_wrapper
+def create(event: dict, context: dict) -> Any:
+    """Create a new mcp server in DynamoDB."""
+    user_id = get_username(event)
+    body = json.loads(event["body"], parse_float=Decimal)
+    body["owner"] = user_id if body.get("owner", None) is None else body["owner"]  # Set the owner of the mcp server
+    mcp_server_model = McpServerModel(**body)
+
+    # Insert the new mcp server item into the DynamoDB table
+    table.put_item(Item=mcp_server_model.model_dump(exclude_none=True))
+    return mcp_server_model.model_dump()
+
+
+@api_wrapper
+def update(event: dict, context: dict) -> Any:
+    """Update an existing mcp server in DynamoDB."""
+    user_id = get_username(event)
+    mcp_server_id = get_mcp_server_id(event)
+    body = json.loads(event["body"], parse_float=Decimal)
+    mcp_server_model = McpServerModel(**body)
+
+    if mcp_server_id != mcp_server_model.id:
+        raise ValueError(f"URL id {mcp_server_id} doesn't match body id {mcp_server_model.id}")
+
+    # Query for the latest mcp server revision
+    response = table.query(KeyConditionExpression=Key("id").eq(mcp_server_id), Limit=1, ScanIndexForward=False)
+    item = get_item(response)
+
+    if item is None:
+        raise ValueError(f"MCP Server {mcp_server_model} not found.")
+
+    # Check if the user is authorized to update the mcp server
+    if is_admin(event) or item["owner"] == user_id:
+        # Update the mcp server
+        logger.info(f"new model: {mcp_server_model.model_dump(exclude_none=True)}")
+        table.put_item(Item=mcp_server_model.model_dump(exclude_none=True))
+        return mcp_server_model.model_dump()
+
+    raise ValueError(f"Not authorized to update {mcp_server_id}.")
+
+
+@api_wrapper
+def delete(event: dict, context: dict) -> Dict[str, str]:
+    """Logically delete a mcp server from DynamoDB."""
+    user_id = get_username(event)
+    mcp_server_id = get_mcp_server_id(event)
+
+    # Query for the mcp server
+    response = table.query(KeyConditionExpression=Key("id").eq(mcp_server_id), Limit=1, ScanIndexForward=False)
+    item = get_item(response)
+
+    if item is None:
+        raise ValueError(f"MCP Server {mcp_server_id} not found.")
+
+    # Check if the user is authorized to delete the mcp server
+    if is_admin(event) or item["owner"] == user_id:
+        logger.info(f"Deleting mcp server {mcp_server_id} for user {user_id}")
+        table.delete_item(Key={"id": mcp_server_id, "owner": item.get("owner")})
+        return {"status": "ok"}
+
+    raise ValueError(f"Not authorized to delete {mcp_server_id}.")
+
+
+def get_mcp_server_id(event: dict) -> str:
+    """Extract the mcp server id from the event's path parameters."""
+    return str(event["pathParameters"]["serverId"])

lambda/mcp_server/models.py

@@ -0,0 +1,47 @@
+#   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 uuid
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class McpServerModel(BaseModel):
+    """
+    A Pydantic model representing a template for prompts.
+    Contains metadata and functionality to create new revisions.
+    """
+
+    # Unique identifier for the mcp server
+    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
+
+    # Timestamp of when the mcp server was created
+    created: Optional[str] = Field(default_factory=lambda: datetime.now().isoformat())
+
+    # Owner of the MCP user
+    owner: str
+
+    # URL of the MCP server
+    url: str
+
+    # Name of the MCP server
+    name: str
+
+    # Custom headers for the MCP client
+    customHeaders: Optional[dict] = Field(default_factory=lambda: None)
+
+    # Custom client properties for the MCP client
+    clientConfig: Optional[dict] = Field(default_factory=lambda: None)

lambda/prompt_templates/lambda_functions.py

@@ -22,7 +22,7 @@
 
 import boto3
 from boto3.dynamodb.conditions import Attr, Key
-from utilities.common_functions import api_wrapper, get_groups, get_username, is_admin, retry_config
+from utilities.common_functions import api_wrapper, get_groups, get_item, get_username, is_admin, retry_config
 
 from .models import PromptTemplateModel
 
@@ -89,8 +89,7 @@ def get(event: dict, context: dict) -> Any:
 
     # Query for the latest prompt template revision
     response = table.query(KeyConditionExpression=Key("id").eq(prompt_template_id), Limit=1, ScanIndexForward=False)
-    items = response.get("Items", [])
-    item = items[0] if items else None
+    item = get_item(response)
 
     if item is None:
         raise ValueError(f"Prompt template {prompt_template_id} not found.")
@@ -159,8 +158,7 @@ def update(event: dict, context: dict) -> Any:
 
     # Query for the latest prompt template revision
     response = table.query(KeyConditionExpression=Key("id").eq(prompt_template_id), Limit=1, ScanIndexForward=False)
-    items = response.get("Items", [])
-    item = items[0] if items else None
+    item = get_item(response)
 
     if item is None:
         raise ValueError(f"Prompt template {prompt_template_model} not found.")
@@ -195,8 +193,7 @@ def delete(event: dict, context: dict) -> Dict[str, str]:
 
     # Query for the latest prompt template revision
     response = table.query(KeyConditionExpression=Key("id").eq(prompt_template_id), Limit=1, ScanIndexForward=False)
-    items = response.get("Items", [])
-    item = items[0] if items else None
+    item = get_item(response)
 
     if item is None:
         raise ValueError(f"Prompt template {prompt_template_id} not found.")

lambda/utilities/common_functions.py

@@ -471,3 +471,8 @@ def get_lambda_role_name() -> str:
     arn = _get_lambda_role_arn()
     parts = arn.split(":assumed-role/")[1].split("/")
     return parts[0]  # This is the role name
+
+
+def get_item(response: Any) -> Any:
+    items = response.get("Items", [])
+    return items[0] if items else None