main_remote.py

"""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())