dss/demo/tools/dss_mcp/server.py

"""
DSS MCP Server

SSE-based Model Context Protocol server for Claude.
Provides project-isolated context and tools with user-scoped integrations.
"""

import asyncio
import json
import logging
import structlog
from typing import Optional, Dict, Any
from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette.sse import EventSourceResponse
from mcp.server import Server
from mcp import types

from .config import mcp_config, validate_config
from .context.project_context import get_context_manager
from .tools.project_tools import PROJECT_TOOLS, ProjectTools

# Configure logging
logging.basicConfig(
    level=mcp_config.LOG_LEVEL,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = structlog.get_logger()

# FastAPI app for SSE endpoints
app = FastAPI(
    title="DSS MCP Server",
    description="Model Context Protocol server for Design System Swarm",
    version="0.8.0"
)

# CORS configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # TODO: Configure based on environment
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# MCP Server instance
mcp_server = Server("dss-mcp")

# Store active sessions
_active_sessions: Dict[str, Dict[str, Any]] = {}


def get_session_key(project_id: str, user_id: Optional[int] = None) -> str:
    """Generate session key for caching"""
    return f"{project_id}:{user_id or 'anonymous'}"


@app.on_event("startup")
async def startup():
    """Startup tasks"""
    logger.info("Starting DSS MCP Server")

    # Validate configuration
    warnings = validate_config()
    if warnings:
        for warning in warnings:
            logger.warning(warning)

    logger.info(
        "DSS MCP Server started",
        host=mcp_config.HOST,
        port=mcp_config.PORT
    )


@app.on_event("shutdown")
async def shutdown():
    """Cleanup on shutdown"""
    logger.info("Shutting down DSS MCP Server")


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    context_manager = get_context_manager()
    return {
        "status": "healthy",
        "server": "dss-mcp",
        "version": "0.8.0",
        "cache_size": len(context_manager._cache),
        "active_sessions": len(_active_sessions)
    }


@app.get("/sse")
async def sse_endpoint(
    project_id: str = Query(..., description="Project ID for context isolation"),
    user_id: Optional[int] = Query(None, description="User ID for user-scoped integrations")
):
    """
    Server-Sent Events endpoint for MCP communication.

    This endpoint maintains a persistent connection with the client
    and streams MCP protocol messages.
    """
    session_key = get_session_key(project_id, user_id)

    logger.info(
        "SSE connection established",
        project_id=project_id,
        user_id=user_id,
        session_key=session_key
    )

    # Load project context
    context_manager = get_context_manager()
    try:
        project_context = await context_manager.get_context(project_id, user_id)
        if not project_context:
            raise HTTPException(status_code=404, detail=f"Project not found: {project_id}")
    except Exception as e:
        logger.error("Failed to load project context", error=str(e))
        raise HTTPException(status_code=500, detail=f"Failed to load project: {str(e)}")

    # Create project tools instance
    project_tools = ProjectTools(user_id)

    # Track session
    _active_sessions[session_key] = {
        "project_id": project_id,
        "user_id": user_id,
        "connected_at": asyncio.get_event_loop().time(),
        "project_tools": project_tools
    }

    async def event_generator():
        """Generate SSE events for MCP communication"""
        try:
            # Send initial connection confirmation
            yield {
                "event": "connected",
                "data": json.dumps({
                    "project_id": project_id,
                    "project_name": project_context.name,
                    "available_tools": len(PROJECT_TOOLS),
                    "integrations_enabled": list(project_context.integrations.keys())
                })
            }

            # Keep connection alive
            while True:
                await asyncio.sleep(30)  # Heartbeat every 30 seconds
                yield {
                    "event": "heartbeat",
                    "data": json.dumps({"timestamp": asyncio.get_event_loop().time()})
                }

        except asyncio.CancelledError:
            logger.info("SSE connection closed", session_key=session_key)
        finally:
            # Cleanup session
            if session_key in _active_sessions:
                del _active_sessions[session_key]

    return EventSourceResponse(event_generator())


# MCP Protocol Handlers
@mcp_server.list_tools()
async def list_tools() -> list[types.Tool]:
    """
    List all available tools.

    Tools are dynamically determined based on:
    - Base DSS project tools (always available)
    - User's enabled integrations (Figma, Jira, Confluence, etc.)
    """
    # Start with base project tools
    tools = PROJECT_TOOLS.copy()

    # TODO: Add integration-specific tools based on user's enabled integrations
    # This will be implemented in Phase 3

    logger.debug("Listed tools", tool_count=len(tools))
    return tools


@mcp_server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
    """
    Execute a tool by name.

    Args:
        name: Tool name
        arguments: Tool arguments (must include project_id)

    Returns:
        Tool execution results
    """
    logger.info("Tool called", tool_name=name, arguments=arguments)

    project_id = arguments.get("project_id")
    if not project_id:
        return [
            types.TextContent(
                type="text",
                text=json.dumps({"error": "project_id is required"})
            )
        ]

    # Find active session for this project
    # For now, use first matching session (can be enhanced with session management)
    session_key = None
    project_tools = None

    for key, session in _active_sessions.items():
        if session["project_id"] == project_id:
            session_key = key
            project_tools = session["project_tools"]
            break

    if not project_tools:
        # Create temporary tools instance
        project_tools = ProjectTools()

    # Execute tool
    try:
        result = await project_tools.execute_tool(name, arguments)

        return [
            types.TextContent(
                type="text",
                text=json.dumps(result, indent=2)
            )
        ]
    except Exception as e:
        logger.error("Tool execution failed", tool_name=name, error=str(e))
        return [
            types.TextContent(
                type="text",
                text=json.dumps({"error": str(e)})
            )
        ]


@mcp_server.list_resources()
async def list_resources() -> list[types.Resource]:
    """
    List available resources.

    Resources provide static or dynamic content that Claude can access.
    Examples: project documentation, component specs, design system guidelines.
    """
    # TODO: Implement resources based on project context
    # For now, return empty list
    return []


@mcp_server.read_resource()
async def read_resource(uri: str) -> str:
    """
    Read a specific resource by URI.

    Args:
        uri: Resource URI (e.g., "dss://project-id/components/Button")

    Returns:
        Resource content
    """
    # TODO: Implement resource reading
    # For now, return not implemented
    return json.dumps({"error": "Resource reading not yet implemented"})


@mcp_server.list_prompts()
async def list_prompts() -> list[types.Prompt]:
    """
    List available prompt templates.

    Prompts provide pre-configured conversation starters for Claude.
    """
    # TODO: Add DSS-specific prompt templates
    # Examples: "Analyze component consistency", "Review token usage", etc.
    return []


@mcp_server.get_prompt()
async def get_prompt(name: str, arguments: dict) -> types.GetPromptResult:
    """
    Get a specific prompt template.

    Args:
        name: Prompt name
        arguments: Prompt arguments

    Returns:
        Prompt content
    """
    # TODO: Implement prompt templates
    return types.GetPromptResult(
        description="Prompt not found",
        messages=[]
    )


# API endpoint to call MCP tools directly (for testing/debugging)
@app.post("/api/tools/{tool_name}")
async def call_tool_api(tool_name: str, arguments: Dict[str, Any]):
    """
    Direct API endpoint to call MCP tools.

    Useful for testing tools without MCP client.
    """
    project_tools = ProjectTools()
    result = await project_tools.execute_tool(tool_name, arguments)
    return result


# API endpoint to list active sessions
@app.get("/api/sessions")
async def list_sessions():
    """List all active SSE sessions"""
    return {
        "active_sessions": len(_active_sessions),
        "sessions": [
            {
                "project_id": session["project_id"],
                "user_id": session["user_id"],
                "connected_at": session["connected_at"]
            }
            for session in _active_sessions.values()
        ]
    }


# API endpoint to clear context cache
@app.post("/api/cache/clear")
async def clear_cache(project_id: Optional[str] = None):
    """Clear context cache for a project or all projects"""
    context_manager = get_context_manager()
    context_manager.clear_cache(project_id)

    return {
        "status": "cache_cleared",
        "project_id": project_id or "all"
    }


if __name__ == "__main__":
    import uvicorn

    logger.info(
        "Starting DSS MCP Server",
        host=mcp_config.HOST,
        port=mcp_config.PORT
    )

    uvicorn.run(
        "server:app",
        host=mcp_config.HOST,
        port=mcp_config.PORT,
        reload=True,
        log_level=mcp_config.LOG_LEVEL.lower()
    )