teslamate-mcp

"""TeslaMate MCP Server - HTTP Transport (Remote)""" import contextlib import json import logging from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any, Dict, List, Optional import click import mcp.types as types import uvicorn from mcp.server.lowlevel import Server from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from psycopg_pool import AsyncConnectionPool from starlette.applications import Starlette from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse from starlette.routing import Mount from starlette.types import Receive, Scope, Send from src.config import Config from src.database import DatabaseManager, create_async_pool from src.tools import TOOL_DEFINITIONS, get_tool_by_name from src.validators import validate_sql_query logger = logging.getLogger(__name__) @dataclass class AppContext: """Application context with database pool and manager""" db_pool: AsyncConnectionPool db_manager: DatabaseManager db_schema: List[Dict[str, str]] # Global app context app_context: AppContext | None = None class BearerAuthMiddleware(BaseHTTPMiddleware): """Bearer token authentication middleware for the MCP server""" def __init__(self, app, auth_token: Optional[str] = None): super().__init__(app) self.auth_token = auth_token async def dispatch(self, request, call_next): # Skip auth if no token is configured if not self.auth_token: return await call_next(request) # Skip auth for non-MCP endpoints if not request.url.path.startswith("/mcp"): return await call_next(request) # Check Authorization header auth_header = request.headers.get("Authorization") if not auth_header or not auth_header.startswith("Bearer "): return JSONResponse( status_code=401, content={"error": "Authorization required"}, headers={"WWW-Authenticate": "Bearer"}, ) # Validate token try: provided_token = auth_header.split(" ", 1)[1] if provided_token != self.auth_token: raise ValueError("Invalid token") except (IndexError, ValueError): return JSONResponse( status_code=401, content={"error": "Invalid token"}, headers={"WWW-Authenticate": "Bearer"}, ) # Continue with the request return await call_next(request) @click.command() @click.option("--port", default=8888, help="Port to listen on for HTTP") @click.option( "--host", default="0.0.0.0", help="Host to listen on", ) @click.option( "--json-response", is_flag=True, default=False, help="Enable JSON responses instead of SSE streams", ) @click.option( "--auth-token", default=None, help="Bearer authentication token (optional)", envvar="AUTH_TOKEN", ) def main( port: int, host: str, json_response: bool, auth_token: str | None, ) -> int: global app_context # Load configuration config = Config.from_env() # Create MCP server app = Server("teslamate") # Tool handler functions async def execute_predefined_tool(tool_name: str) -> List[Dict[str, Any]]: """Execute a predefined tool by name""" if not app_context: raise RuntimeError("Application context not initialized") tool = get_tool_by_name(tool_name) return await app_context.db_manager.execute_query_async( tool.sql_file, app_context.db_pool ) async def get_database_schema() -> List[Dict[str, str]]: """Return the database schema information""" if not app_context: raise RuntimeError("Application context not initialized") return app_context.db_schema async def run_sql(query: str) -> List[Dict[str, Any]]: """Execute a custom SQL query with validation""" if not app_context: raise RuntimeError("Application context not initialized") # Validate the SQL query is_valid, error_msg = validate_sql_query(query) if not is_valid: raise ValueError(error_msg) return await app_context.db_manager.execute_custom_query_async( query, app_context.db_pool ) @app.call_tool() async def call_tool(name: str, arguments: dict) -> list[types.ContentBlock]: """Route tool calls to appropriate handlers""" # Handle custom SQL query tool if name == "run_sql": query = arguments.get("query") if not query: raise ValueError("Missing required argument 'query' for run_sql") result = await run_sql(query) # Handle database schema tool elif name == "get_database_schema": result = await get_database_schema() # Handle predefined tools else: result = await execute_predefined_tool(name) # Convert result to MCP content blocks return [ types.TextContent( type="text", text=json.dumps(result, indent=2, default=str), ) ] @app.list_tools() async def list_tools() -> list[types.Tool]: """List all available TeslaMate tools""" tools = [] # Add all predefined tools for tool_def in TOOL_DEFINITIONS: tools.append( types.Tool( name=tool_def.name, description=tool_def.description, inputSchema={"type": "object", "properties": {}}, ) ) # Add database schema tool tools.append( types.Tool( name="get_database_schema", description="Get the TeslaMate database schema information including all tables and columns with their data types. Use this to understand the database structure before writing SQL queries.", inputSchema={"type": "object", "properties": {}}, ) ) # Add custom SQL query tool tools.append( types.Tool( name="run_sql", description="Execute a custom SELECT SQL query on the TeslaMate database. Only SELECT queries are allowed. Use get_database_schema first to understand the available tables and columns.", inputSchema={ "type": "object", "properties": { "query": { "type": "string", "description": "The SELECT SQL query to execute. Must be a single SELECT statement.", } }, "required": ["query"], }, ) ) return tools # Create the session manager with true stateless mode session_manager = StreamableHTTPSessionManager( app=app, event_store=None, json_response=json_response, stateless=True, ) async def handle_streamable_http( scope: Scope, receive: Receive, send: Send ) -> None: await session_manager.handle_request(scope, receive, send) @contextlib.asynccontextmanager async def lifespan(app: Starlette) -> AsyncIterator[None]: """Context manager for application lifecycle""" global app_context # Create database connection pool db_pool = create_async_pool(config.database_url) try: # Initialize the pool await db_pool.open() logger.info("Database connection pool initialized") # Load database schema db_schema = DatabaseManager.load_db_schema() # Initialize app context db_manager = DatabaseManager(config) app_context = AppContext( db_pool=db_pool, db_manager=db_manager, db_schema=db_schema ) # Start session manager async with session_manager.run(): logger.info("Application started with StreamableHTTP session manager!") yield finally: logger.info("Application shutting down...") if app_context: await app_context.db_pool.close() logger.info("Database connection pool closed") app_context = None # Create an ASGI application using the transport starlette_app = Starlette( debug=True, routes=[ Mount("/mcp", app=handle_streamable_http), ], lifespan=lifespan, ) # Add bearer auth middleware if token is provided if auth_token: starlette_app.add_middleware(BearerAuthMiddleware, auth_token=auth_token) logger.info("Bearer token authentication enabled") # Run with uvicorn logger.info(f"Starting TeslaMate MCP server on {host}:{port}") uvicorn.run(starlette_app, host=host, port=port) return 0 if __name__ == "__main__": import sys sys.exit(main())