"""
Lambda Handler for MCP Gateway
This handler can be deployed as a Lambda function to handle MCP requests
"""
import json
import logging
import os
import sys
from typing import Dict, Any
# Import directly from gateway module file (for Lambda deployment where files are siblings)
# Add current directory to path to ensure we can import sibling modules
_current_dir = os.path.dirname(os.path.abspath(__file__))
if _current_dir not in sys.path:
sys.path.insert(0, _current_dir)
# Import gateway module directly
import importlib.util
gateway_path = os.path.join(_current_dir, 'gateway.py')
gateway_spec = importlib.util.spec_from_file_location("gateway", gateway_path)
gateway = importlib.util.module_from_spec(gateway_spec)
gateway_spec.loader.exec_module(gateway)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# Initialize gateway
aws_region = os.getenv('AWS_REGION', 'us-east-1')
environment = os.getenv('ENVIRONMENT', 'dev')
mcp_gateway = gateway.MCPGateway(region=aws_region)
# Auto-register tools on module load
def auto_register_tools():
"""Automatically register Lambda tools based on environment"""
try:
import boto3
# Get account ID and region
sts_client = boto3.client('sts', region_name=aws_region)
account_id = sts_client.get_caller_identity()['Account']
auto_register_tools_with_account(account_id)
except Exception as e:
logger.warning(f"Could not auto-register tools: {e}")
def auto_register_tools_with_account(account_id: str):
"""Register tools with a known account ID"""
try:
# Register Weather Tool
weather_arn = f"arn:aws:lambda:{aws_region}:{account_id}:function:{environment}-weather-tool"
mcp_gateway.register_lambda_tool(
name="weather_lookup",
lambda_arn=weather_arn,
description="Get weather information for a location",
input_schema={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "City name or location"
},
"units": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "Temperature units",
"default": "celsius"
}
},
"required": ["location"]
}
)
# Register Calculator Tool
calculator_arn = f"arn:aws:lambda:{aws_region}:{account_id}:function:{environment}-calculator-tool"
mcp_gateway.register_lambda_tool(
name="calculator",
lambda_arn=calculator_arn,
description="Perform mathematical calculations",
input_schema={
"type": "object",
"properties": {
"operation": {
"type": "string",
"enum": ["add", "subtract", "multiply", "divide"],
"description": "Mathematical operation"
},
"a": {
"type": "number",
"description": "First number"
},
"b": {
"type": "number",
"description": "Second number"
}
},
"required": ["operation", "a", "b"]
}
)
# Register Data Lookup Tool
data_lookup_arn = f"arn:aws:lambda:{aws_region}:{account_id}:function:{environment}-data-lookup-tool"
mcp_gateway.register_lambda_tool(
name="data_lookup",
lambda_arn=data_lookup_arn,
description="Lookup data by query type and value",
input_schema={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Type of query (e.g., 'user_id', 'order_id')"
},
"value": {
"type": "string",
"description": "Value to lookup"
}
},
"required": ["query", "value"]
}
)
# Register REST API tools
# Get REST API URL from environment or construct from account/region
rest_api_id = os.getenv('REST_API_SERVER_API_ID', '')
if rest_api_id:
rest_api_base_url = f"https://{rest_api_id}.execute-api.{aws_region}.amazonaws.com/{environment}"
else:
# Try to construct from account ID (will be set after first deployment)
rest_api_base_url = f"https://<rest-api-id>.execute-api.{aws_region}.amazonaws.com/{environment}"
logger.warning("REST_API_SERVER_API_ID not set, REST tools will not be registered")
if rest_api_id:
# Register Stock Price Tool
mcp_gateway.register_rest_tool(
name="stock_price",
endpoint=f"{rest_api_base_url}/stock/price",
description="Get stock price information",
input_schema={
"type": "object",
"properties": {
"symbol": {
"type": "string",
"description": "Stock symbol (e.g., AAPL, GOOGL)"
}
},
"required": ["symbol"]
},
method="POST",
headers={"Content-Type": "application/json"}
)
# Register Current Time Tool
mcp_gateway.register_rest_tool(
name="current_time",
endpoint=f"{rest_api_base_url}/time/current",
description="Get current UTC time",
input_schema={
"type": "object",
"properties": {}
},
method="GET"
)
# Register Text Summarize Tool
mcp_gateway.register_rest_tool(
name="summarize_text",
endpoint=f"{rest_api_base_url}/text/summarize",
description="Summarize text content",
input_schema={
"type": "object",
"properties": {
"text": {
"type": "string",
"description": "Text to summarize"
}
},
"required": ["text"]
},
method="POST",
headers={"Content-Type": "application/json"}
)
logger.info(f"Auto-registered {len(mcp_gateway.tool_registry.list_tools())} tools")
except Exception as e:
logger.error(f"Error registering tools: {e}", exc_info=True)
# Register tools on module load
auto_register_tools()
def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
"""
AWS Lambda handler for MCP Gateway
Args:
event: Lambda event (contains MCP request)
context: Lambda context
Returns:
Lambda response with MCP protocol response
"""
# Ensure tools are registered (in case auto-registration failed)
if len(mcp_gateway.tool_registry.list_tools()) == 0 and context:
try:
# Try to get account ID from context
function_arn = context.invoked_function_arn
# Extract account ID from ARN: arn:aws:lambda:region:account-id:function:name
account_id = function_arn.split(':')[4]
auto_register_tools_with_account(account_id)
except Exception as e:
logger.warning(f"Could not register tools from context: {e}")
try:
# Extract MCP request from event
# For API Gateway integration, request is in event body
if 'body' in event:
request_body = json.loads(event['body'])
else:
request_body = event
# Extract session ID from headers if present
headers = event.get('headers', {})
session_id = headers.get('Mcp-Session-Id') or headers.get('mcp-session-id')
# Handle request (async in sync context - gateway handles this)
import asyncio
response = asyncio.run(mcp_gateway.handle_request(request_body, session_id))
return {
'statusCode': 200,
'headers': {
'Content-Type': 'application/json',
'Access-Control-Allow-Origin': '*'
},
'body': json.dumps(response)
}
except Exception as e:
logger.error(f"Error in lambda_handler: {e}", exc_info=True)
return {
'statusCode': 500,
'headers': {
'Content-Type': 'application/json'
},
'body': json.dumps({
'error': str(e)
})
}
