""" 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 from .tools.workflow_tools import WORKFLOW_TOOLS, WorkflowTools from .tools.debug_tools import DEBUG_TOOLS, DebugTools from .integrations.storybook import STORYBOOK_TOOLS from .integrations.translations import TRANSLATION_TOOLS from .plugin_registry import PluginRegistry # 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") # Initialize Plugin Registry plugin_registry = PluginRegistry() plugin_registry.load_plugins() # 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) - Workflow orchestration tools - Debug tools - Storybook integration tools - Dynamically loaded plugins - User's enabled integrations (Figma, Jira, Confluence, etc.) """ # Start with base project tools tools = PROJECT_TOOLS.copy() # Add workflow orchestration tools tools.extend(WORKFLOW_TOOLS) # Add debug tools tools.extend(DEBUG_TOOLS) # Add Storybook integration tools tools.extend(STORYBOOK_TOOLS) # Add Translation tools tools.extend(TRANSLATION_TOOLS) # Add plugin tools tools.extend(plugin_registry.get_all_tools()) # 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), plugin_count=len(plugin_registry.plugins)) 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() # Check if this is a workflow tool workflow_tool_names = [tool.name for tool in WORKFLOW_TOOLS] debug_tool_names = [tool.name for tool in DEBUG_TOOLS] storybook_tool_names = [tool.name for tool in STORYBOOK_TOOLS] translation_tool_names = [tool.name for tool in TRANSLATION_TOOLS] # Execute tool try: if name in workflow_tool_names: # Handle workflow orchestration tools from .audit import AuditLog audit_log = AuditLog() workflow_tools = WorkflowTools(audit_log) result = await workflow_tools.handle_tool_call(name, arguments) elif name in debug_tool_names: # Handle debug tools debug_tools = DebugTools() result = await debug_tools.execute_tool(name, arguments) elif name in storybook_tool_names: # Handle Storybook tools from .integrations.storybook import StorybookTools storybook_tools = StorybookTools() result = await storybook_tools.execute_tool(name, arguments) elif name in translation_tool_names: # Handle Translation tools from .integrations.translations import TranslationTools translation_tools = TranslationTools() result = await translation_tools.execute_tool(name, arguments) elif name in plugin_registry.handlers: # Handle plugin tools result = await plugin_registry.execute_tool(name, arguments) # Plugin tools return MCP content objects directly, not dicts if isinstance(result, list): return result else: # Handle regular project tools 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() )