Building multi-tenant agents on AWS

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "dynamodb:DeleteItem",
        "dynamodb:GetItem",
        "dynamodb:PutItem",
        "dynamodb:Query",
        "dynamodb:UpdateItem"
      ],
      "Resource": "arn:aws:dynamodb:us-east-1:905418320997:table/pooled-tenant-table",
      "Condition": {
        "ForAllValues:StringEquals": {
          "dynamodb:LeadingKeys": [
            "${aws:PrincipalTag/tenantId}"
          ]
        }
      }
    }
  ]
}

"""
MCP Server for tenant-specific Amazon DynamoDB access using FastMCP and ABAC.
"""

import json
import os
import sys
import boto3
import uuid
import re
import logging
from datetime import datetime
from typing import Dict, List, Any, Optional, TypeVar, Generic
from pydantic import BaseModel
from fastmcp import FastMCP

# Import configuration (assumed to be defined in config.py)
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import MCP_SERVER_NAME, MCP_SERVER_DESCRIPTION, ROLE_ARN_TEMPLATE, DYNAMODB_TABLE_NAME

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

T = TypeVar('T')

class MCPResponse(BaseModel, Generic[T]):
    result: Optional[T] = None
    error: Optional[str] = None
    status_code: int = 200

    model_config = {"arbitrary_types_allowed": True}

    def dict(self) -> Dict[str, Any]:
        response = {}
        if self.result is not None:
            response["result"] = self.result
        if self.error is not None:
            response["error"] = self.error
        return response

class TenantItem(BaseModel):
    tenant_id: str
    item_id: str
    product: str
    value: Any
    description: Optional[str] = None
    created_at: str

class TenantCredentialManager:
    def __init__(self, role_arn_template: str):
        self.role_arn_template = role_arn_template
        self.sts_client = boto3.client('sts')

    def get_tenant_credentials(self, tenant_id: str) -> Dict[str, str]:
        if not re.match(r'^[a-zA-Z0-9-]{1,36}$', tenant_id):
            raise ValueError(f"Invalid tenant ID format: {tenant_id}")
        try:
            response = self.sts_client.assume_role(
                RoleArn=self.role_arn_template,
                RoleSessionName=f"tenant-{tenant_id}-session",
                DurationSeconds=3600,
                Tags=[{'Key': 'tenantId', 'Value': tenant_id}]
            )
            credentials = response['Credentials']
            return {
                'aws_access_key_id': credentials['AccessKeyId'],
                'aws_secret_access_key': credentials['SecretAccessKey'],
                'aws_session_token': credentials['SessionToken']
            }
        except Exception as e:
            raise Exception(f"Failed to assume role for tenant {tenant_id}: {str(e)}")

class TenantDynamoDBManager:
    def __init__(self, table_name: str):
        self.table_name = table_name

    def get_dynamodb_client(self, credentials: Dict[str, str]):
        return boto3.client(
            'dynamodb',
            aws_access_key_id=credentials['aws_access_key_id'],
            aws_secret_access_key=credentials['aws_secret_access_key'],
            aws_session_token=credentials['aws_session_token']
        )

    def get_tenant_items(self, dynamodb_client, tenant_id: str) -> List[Dict]:
        try:
            response = dynamodb_client.query(
                TableName=self.table_name,
                KeyConditionExpression="tenant_id = :tid",
                ExpressionAttributeValues={":tid": {"S": tenant_id}}
            )
            items = []
            for item in response.get('Items', []):
                parsed_item = {k: v[list(v.keys())[0]] for k, v in item.items()}
                items.append(parsed_item)
            return items
        except Exception as e:
            logger.error(f"Failed to query items for tenant {tenant_id}: {str(e)}")
            raise

    def put_tenant_item(self, dynamodb_client, tenant_id: str, item_data: Dict) -> Dict:
        try:
            item_with_tenant = item_data.copy()
            item_with_tenant['tenant_id'] = tenant_id
            item_with_tenant['item_id'] = item_with_tenant.get('item_id', str(uuid.uuid4()))
            item_with_tenant['created_at'] = datetime.now().isoformat()
            dynamodb_item = {
                k: {"S": str(v)} if isinstance(v, (str, dict, list)) else
                   {"N": str(v)} if isinstance(v, (int, float)) else
                   {"BOOL": v} if isinstance(v, bool) else
                   {"S": str(v)}
                for k, v in item_with_tenant.items()
            }
            dynamodb_client.put_item(TableName=self.table_name, Item=dynamodb_item)
            return item_with_tenant
        except Exception as e:
            logger.error(f"Failed to put item for tenant {tenant_id}: {str(e)}")
            raise

app = FastMCP(
    title=MCP_SERVER_NAME,
    description=MCP_SERVER_DESCRIPTION
)

credential_manager = TenantCredentialManager(ROLE_ARN_TEMPLATE)
dynamodb_manager = TenantDynamoDBManager(DYNAMODB_TABLE_NAME)

@app.tool(
    name="get_tenant_items",
    description="Get items for a specific tenant from the pooled DynamoDB table"
)
def get_tenant_items(tenant_id: str) -> MCPResponse[List[TenantItem]]:
    try:
        credentials = credential_manager.get_tenant_credentials(tenant_id)
        dynamodb_client = dynamodb_manager.get_dynamodb_client(credentials)
        items = dynamodb_manager.get_tenant_items(dynamodb_client, tenant_id)
        tenant_items = [
            TenantItem(
                tenant_id=item.get("tenant_id"),
                item_id=item.get("item_id"),
                product=item.get("product"),
                value=item.get("value"),
                description=item.get("description"),
                created_at=item.get("created_at")
            ) for item in items
        ]
        return MCPResponse(result=tenant_items)
    except ValueError as e:
        return MCPResponse(error=str(e), status_code=400)
    except Exception as e:
        return MCPResponse(error=str(e), status_code=500)

@app.tool(
    name="put_tenant_item",
    description="Add an item for a specific tenant to the pooled DynamoDB table"
)
def put_tenant_item(tenant_id: str, product: str, value: int) -> MCPResponse[TenantItem]:
    try:
        credentials = credential_manager.get_tenant_credentials(tenant_id)
        dynamodb_client = dynamodb_manager.get_dynamodb_client(credentials)
        item_data_dict = {"product": product, "value": value}
        created_item = dynamodb_manager.put_tenant_item(dynamodb_client, tenant_id, item_data_dict)
        tenant_item = TenantItem(**created_item)
        return MCPResponse(result=tenant_item)
    except ValueError as e:
        return MCPResponse(error=str(e), status_code=400)
    except Exception as e:
        return MCPResponse(error=str(e), status_code=500)

if __name__ == "__main__":
    app.run('streamable-http', host="0.0.0.0", port=8000, path='/mcp')

import os
import json
import re
import uuid
import logging
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field
from mcp.client.streamable_http import streamablehttp_client
from strands.agent.conversation_manager import ConversationManager, SlidingWindowConversationManager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import configuration
from config import MCP_SERVER_URL, AGENT_NAME, AGENT_DESCRIPTION, BEDROCK_REGION, BEDROCK_MODEL_ID

# Import Strands Agents SDK
from strands import Agent
from strands.tools.mcp import MCPClient

async def main(input_text: str = None):
    """
    Main function to create and run the agent.
    
    Args:
        input_text: Optional input text to process. If None, runs in interactive mode.
    """
    logger.info(f"Starting Bedrock Agent for MCP Server at {MCP_SERVER_URL}")
    
    # Create agent
    mcp_client = MCPClient(lambda: streamablehttp_client(MCP_SERVER_URL))

    with mcp_client:
            
        # Create agent
        agent = Agent(
            model=BEDROCK_MODEL_ID,
            system_prompt="You help users to create products and get product details.
            ",
            tools=mcp_client.list_tools_sync(),
            conversation_manager=CustomConversationManager()
        )
        
        tenant_id = uuid.uuid4()
        
        input_text_with_context = f'{input_text} for tenant id {tenant_id}'
        response = agent(input_text_with_context)
        logger.info(response)

# Entry point
if __name__ == "__main__":
    try:
        import asyncio
        asyncio.run(main())
    except Exception as e:
        logger.error(f"Failed to run agent: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())

class CustomConversationManager(SlidingWindowConversationManager): 

     def __scrub_tenant_id(self, messages, uuid_pattern=r'[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}'):
        def scrub_value(value):
            if isinstance(value, str):
                return re.sub(uuid_pattern, '<tenant_id_scrubbed>', value)
            elif isinstance(value, list):
                return [scrub_value(v) for v in value]
            elif isinstance(value, dict):
                return {k: scrub_value(v) for k, v in value.items()}
            else:
                return value

        for i, msg in enumerate(messages):
            messages[i] = scrub_value(msg)    
     
     def apply_management(self, messages):
         self.__scrub_tenant_id(messages=messages)
         return super().apply_management(messages)