drawio-mcp-server/src/index.ts

#!/usr/bin/env node
import express from 'express';
import cors from 'cors';
import { randomUUID } from 'node:crypto';
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListResourcesRequestSchema,
  ListToolsRequestSchema,
  McpError,
  ReadResourceRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';

// Import tools and utilities
import { createDiagram, validateCreateDiagramInput, getSupportedDiagramTypes, getDiagramTypeDescription } from './tools/create-diagram.js';
import { DiagramType, DiagramFormat } from './types/diagram-types.js';

// Functional types and configurations
type ServerConfig = Readonly<{
  name: string;
  version: string;
  workspaceRoot: string;
  httpPort: number;
  capabilities: Readonly<{
    resources: Record<string, unknown>;
    tools: Record<string, unknown>;
  }>;
}>;

type ToolHandler = (args: any) => Promise<any>;
type ResourceHandler = (uri: string) => Promise<any>;

type ToolHandlers = Readonly<{
  create_diagram: ToolHandler;
  get_diagram_types: ToolHandler;
}>;

type ResourceHandlers = Readonly<{
  'diagrams://types/supported': ResourceHandler;
}>;

// Pure function to create server configuration
const createServerConfig = (workspaceRoot?: string, httpPort?: number): ServerConfig => ({
  name: 'drawio-mcp-server',
  version: '0.1.0',
  workspaceRoot: workspaceRoot || process.env.WORKSPACE_ROOT || process.cwd(),
  httpPort: httpPort || parseInt(process.env.HTTP_PORT || '3000', 10),
  capabilities: {
    resources: {},
    tools: {},
  },
});

// Pure function to create tool definitions
const createToolDefinitions = () => [
  {
    name: 'create_diagram',
    description: 'Create a new diagram of specified type (BPMN, UML, ER, Network, Architecture, etc.) with AI-powered generation',
    inputSchema: {
      type: 'object',
      properties: {
        name: {
          type: 'string',
          description: 'Name of the diagram file (without extension)',
        },
        type: {
          type: 'string',
          enum: Object.values(DiagramType),
          description: 'Type of diagram to create',
        },
        format: {
          type: 'string',
          enum: Object.values(DiagramFormat),
          default: 'drawio',
          description: 'File format for the diagram',
        },
        description: {
          type: 'string',
          description: 'Natural language description of the diagram to generate using AI',
        },
        outputPath: {
          type: 'string',
          description: 'Output directory path (relative to workspace)',
        },
        complexity: {
          type: 'string',
          enum: ['simple', 'detailed'],
          default: 'detailed',
          description: 'Complexity level of the generated diagram',
        },
        language: {
          type: 'string',
          default: 'es',
          description: 'Language for diagram labels and text',
        },
        // Legacy parameters for backward compatibility
        processName: {
          type: 'string',
          description: 'Name of the BPMN process (legacy)',
        },
        tasks: {
          type: 'array',
          items: { type: 'string' },
          description: 'List of tasks for BPMN process (legacy)',
        },
        gatewayType: {
          type: 'string',
          enum: ['exclusive', 'parallel'],
          description: 'Type of gateway for BPMN process (legacy)',
        },
        branches: {
          type: 'array',
          items: {
            type: 'array',
            items: { type: 'string' },
          },
          description: 'Branches for BPMN gateway (legacy)',
        },
        beforeGateway: {
          type: 'array',
          items: { type: 'string' },
          description: 'Tasks before gateway (legacy)',
        },
        afterGateway: {
          type: 'array',
          items: { type: 'string' },
          description: 'Tasks after gateway (legacy)',
        },
        classes: {
          type: 'array',
          items: { type: 'string' },
          description: 'List of classes for UML class diagram (legacy)',
        },
        entities: {
          type: 'array',
          items: { type: 'string' },
          description: 'List of entities for ER diagram (legacy)',
        },
        components: {
          type: 'array',
          items: { type: 'string' },
          description: 'List of components for network or architecture diagram (legacy)',
        },
        processes: {
          type: 'array',
          items: { type: 'string' },
          description: 'List of processes for flowchart (legacy)',
        },
      },
      required: ['name', 'type'],
    },
  },
  {
    name: 'get_diagram_types',
    description: 'Get list of supported diagram types with descriptions',
    inputSchema: {
      type: 'object',
      properties: {},
    },
  },
];

// Pure function to create resource definitions
const createResourceDefinitions = () => [
  {
    uri: 'diagrams://types/supported',
    name: 'Supported Diagram Types',
    mimeType: 'application/json',
    description: 'List of supported diagram types and their descriptions',
  },
];

// Pure function to create tool handlers
const createToolHandlers = (config: ServerConfig): ToolHandlers => ({
  create_diagram: async (args: any) => {
    if (!validateCreateDiagramInput(args)) {
      throw new McpError(
        ErrorCode.InvalidParams,
        'Invalid create_diagram arguments',
      );
    }

    const result = await createDiagram({
      ...args,
      workspaceRoot: args.workspaceRoot || config.workspaceRoot,
    });

    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify(result, null, 2),
        },
      ],
    };
  },

  get_diagram_types: async () => {
    const types = getSupportedDiagramTypes();
    const typesWithDescriptions = types.map(type => ({
      type,
      description: getDiagramTypeDescription(type),
    }));

    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify({
            success: true,
            supportedTypes: typesWithDescriptions,
          }, null, 2),
        },
      ],
    };
  },
});

// Pure function to create resource handlers
const createResourceHandlers = (_config: ServerConfig): ResourceHandlers => ({
  'diagrams://types/supported': async () => {
    const types = getSupportedDiagramTypes();
    const typesWithDescriptions = types.map(type => ({
      type,
      description: getDiagramTypeDescription(type),
    }));
    
    return {
      contents: [
        {
          uri: 'diagrams://types/supported',
          mimeType: 'application/json',
          text: JSON.stringify(typesWithDescriptions, null, 2),
        },
      ],
    };
  },
});

// Pure function to setup tool request handlers
const setupToolRequestHandlers = (server: Server, toolHandlers: ToolHandlers) => {
  // List available tools
  server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: createToolDefinitions(),
  }));

  // Handle tool calls
  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    try {
      const handler = toolHandlers[request.params.name as keyof ToolHandlers];
      if (!handler) {
        throw new McpError(
          ErrorCode.MethodNotFound,
          `Unknown tool: ${request.params.name}`,
        );
      }
      
      return await handler(request.params.arguments);
    } catch (error) {
      console.error(`Error in tool ${request.params.name}:`, error);
      return {
        content: [
          {
            type: 'text',
            text: `Error: ${error instanceof Error ? error.message : String(error)}`,
          },
        ],
        isError: true,
      };
    }
  });

  return server;
};

// Pure function to setup resource request handlers
const setupResourceRequestHandlers = (server: Server, resourceHandlers: ResourceHandlers) => {
  // List available resources
  server.setRequestHandler(ListResourcesRequestSchema, async () => ({
    resources: createResourceDefinitions(),
  }));

  // Handle resource requests
  server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
    const uri = request.params.uri;
    
    try {
      const handler = resourceHandlers[uri as keyof ResourceHandlers];
      if (!handler) {
        throw new McpError(
          ErrorCode.InvalidRequest,
          `Unknown resource URI: ${uri}`,
        );
      }
      
      return await handler(uri);
    } catch (error) {
      throw new McpError(
        ErrorCode.InternalError,
        `Failed to read resource: ${error}`,
      );
    }
  });

  return server;
};

// Pure function to setup error handling
const setupErrorHandling = (server: Server) => {
  server.onerror = (error) => console.error('[MCP Error]', error);
  return server;
};

// Pure function to create and configure MCP server
const createMCPServer = (config: ServerConfig): Server => {
  const server = new Server(
    {
      name: config.name,
      version: config.version,
    },
    {
      capabilities: config.capabilities,
    },
  );

  const toolHandlers = createToolHandlers(config);
  const resourceHandlers = createResourceHandlers(config);

  // Compose server setup using function composition
  const serverWithTools = setupToolRequestHandlers(server, toolHandlers);
  const serverWithResources = setupResourceRequestHandlers(serverWithTools, resourceHandlers);
  const serverWithErrorHandling = setupErrorHandling(serverWithResources);

  return serverWithErrorHandling;
};

// Pure function to create Express app with CORS
const createExpressApp = () => {
  const app = express();
  
  // Add CORS middleware with proper headers for MCP
  app.use(cors({
    origin: '*', // Configure appropriately for production
    exposedHeaders: ['Mcp-Session-Id'],
    allowedHeaders: ['Content-Type', 'mcp-session-id'],
  }));
  
  app.use(express.json());
  
  return app;
};

// Session management for stateful connections
type SessionTransport = {
  transport: StreamableHTTPServerTransport;
  server: Server;
};

// Map to store transports by session ID
const transports: { [sessionId: string]: SessionTransport } = {};

// Pure function to check if request is initialize request
const isInitializeRequest = (body: any): boolean => {
  return body && body.method === 'initialize';
};

// Pure function to create session transport
const createSessionTransport = (_config: ServerConfig): StreamableHTTPServerTransport => {
  return new StreamableHTTPServerTransport({
    sessionIdGenerator: () => randomUUID(),
    onsessioninitialized: (sessionId) => {
      // Session will be stored when server is connected
      console.log(`Session initialized: ${sessionId}`);
    },
    enableDnsRebindingProtection: true,
    allowedHosts: ['127.0.0.1', 'localhost'],
    allowedOrigins: ['*'], // Configure appropriately for production
  });
};

// Function to setup MCP routes on Express app
const setupMCPRoutes = (app: express.Application, config: ServerConfig) => {
  // Handle POST requests for client-to-server communication
  app.post('/mcp', async (req, res) => {
    try {
      // Check for existing session ID
      const sessionId = req.headers['mcp-session-id'] as string | undefined;
      let sessionTransport: SessionTransport;

      if (sessionId && transports[sessionId]) {
        // Reuse existing transport
        sessionTransport = transports[sessionId];
      } else if (!sessionId && isInitializeRequest(req.body)) {
        // New initialization request
        const transport = createSessionTransport(config);
        const server = createMCPServer(config);

        // Clean up transport when closed
        transport.onclose = () => {
          if (transport.sessionId) {
            delete transports[transport.sessionId];
            console.log(`Session closed: ${transport.sessionId}`);
          }
        };

        // Connect to the MCP server
        await server.connect(transport);

        sessionTransport = { transport, server };

        // Store the transport by session ID after connection
        if (transport.sessionId) {
          transports[transport.sessionId] = sessionTransport;
        }
      } else {
        // Invalid request
        res.status(400).json({
          jsonrpc: '2.0',
          error: {
            code: -32000,
            message: 'Bad Request: No valid session ID provided',
          },
          id: null,
        });
        return;
      }

      // Handle the request
      await sessionTransport.transport.handleRequest(req, res, req.body);
    } catch (error) {
      console.error('Error handling MCP POST request:', error);
      if (!res.headersSent) {
        res.status(500).json({
          jsonrpc: '2.0',
          error: {
            code: -32603,
            message: 'Internal server error',
          },
          id: null,
        });
      }
    }
  });

  // Reusable handler for GET and DELETE requests
  const handleSessionRequest = async (req: express.Request, res: express.Response) => {
    try {
      const sessionId = req.headers['mcp-session-id'] as string | undefined;
      if (!sessionId || !transports[sessionId]) {
        res.status(400).send('Invalid or missing session ID');
        return;
      }
      
      const sessionTransport = transports[sessionId];
      await sessionTransport.transport.handleRequest(req, res);
    } catch (error) {
      console.error('Error handling session request:', error);
      if (!res.headersSent) {
        res.status(500).send('Internal server error');
      }
    }
  };

  // Handle GET requests for server-to-client notifications via SSE
  app.get('/mcp', handleSessionRequest);

  // Handle DELETE requests for session termination
  app.delete('/mcp', handleSessionRequest);

  return app;
};

// Function to setup additional API routes
const setupAPIRoutes = (app: express.Application, config: ServerConfig) => {
  // Health check endpoint
  app.get('/health', (req, res) => {
    res.json({
      status: 'healthy',
      timestamp: new Date().toISOString(),
      version: config.version,
      services: {
        mcp: 'operational',
        diagramGeneration: 'operational',
        sessions: Object.keys(transports).length,
      },
    });
  });

  // API documentation endpoint
  app.get('/', (req, res) => {
    const documentation = `
<!DOCTYPE html>
<html>
<head>
    <title>Draw.io MCP Server</title>
    <style>
        body { font-family: Arial, sans-serif; margin: 40px; }
        .endpoint { background: #f5f5f5; padding: 15px; margin: 10px 0; border-radius: 5px; }
        .method { color: #fff; padding: 3px 8px; border-radius: 3px; font-weight: bold; }
        .post { background: #28a745; }
        .get { background: #007bff; }
        .delete { background: #dc3545; }
        code { background: #e9ecef; padding: 2px 4px; border-radius: 3px; }
        pre { background: #f8f9fa; padding: 15px; border-radius: 5px; overflow-x: auto; }
    </style>
</head>
<body>
    <h1>Draw.io MCP Server</h1>
    <p>Model Context Protocol server for AI-powered diagram generation with draw.io integration</p>
    
    <h2>MCP Endpoints</h2>
    <div class="endpoint">
        <h3><span class="method post">POST</span> /mcp</h3>
        <p>MCP client-to-server communication</p>
        <p>Headers: <code>Content-Type: application/json</code>, <code>mcp-session-id</code> (optional)</p>
    </div>

    <div class="endpoint">
        <h3><span class="method get">GET</span> /mcp</h3>
        <p>Server-to-client notifications via Server-Sent Events</p>
        <p>Headers: <code>mcp-session-id</code> (required)</p>
    </div>

    <div class="endpoint">
        <h3><span class="method delete">DELETE</span> /mcp</h3>
        <p>Session termination</p>
        <p>Headers: <code>mcp-session-id</code> (required)</p>
    </div>

    <h2>API Endpoints</h2>
    <div class="endpoint">
        <h3><span class="method get">GET</span> /health</h3>
        <p>Health check and server status</p>
    </div>

    <h2>Available Tools</h2>
    <ul>
        <li><strong>create_diagram:</strong> Create AI-powered diagrams (BPMN, UML, ER, Architecture, etc.)</li>
        <li><strong>list_diagrams:</strong> List all diagram files in workspace</li>
        <li><strong>open_diagram_in_vscode:</strong> Open diagrams in VSCode with draw.io extension</li>
        <li><strong>setup_vscode_environment:</strong> Setup VSCode environment for draw.io</li>
        <li><strong>get_diagram_types:</strong> Get supported diagram types and descriptions</li>
    </ul>

    <h2>Resources</h2>
    <ul>
        <li><strong>diagrams://workspace/list:</strong> Workspace diagram files</li>
        <li><strong>diagrams://types/supported:</strong> Supported diagram types</li>
    </ul>

    <h2>Example Usage</h2>
    <p>Connect using any MCP-compatible client to start generating diagrams with AI!</p>
    
    <h3>Active Sessions: ${Object.keys(transports).length}</h3>
</body>
</html>`;

    res.send(documentation);
  });

  return app;
};

// Main function to start the server
const main = async (): Promise<void> => {
  try {
    const config = createServerConfig();
    
    console.log('🚀 Starting Draw.io MCP Server...');
    console.log(`📊 Version: ${config.version}`);
    console.log(`📁 Workspace: ${config.workspaceRoot}`);
    console.log(`🌐 HTTP Port: ${config.httpPort}`);
    console.log('');

    // Create Express app
    const app = createExpressApp();
    
    // Setup routes
    setupMCPRoutes(app, config);
    setupAPIRoutes(app, config);

    // Start the server
    const server = app.listen(config.httpPort, (error?: Error) => {
      if (error) {
        console.error('❌ Failed to start server:', error);
        process.exit(1);
      }
      
      console.log('✅ Draw.io MCP Server started successfully!');
      console.log(`🔗 MCP Endpoint: http://localhost:${config.httpPort}/mcp`);
      console.log(`📖 Documentation: http://localhost:${config.httpPort}/`);
      console.log(`🏥 Health Check: http://localhost:${config.httpPort}/health`);
      console.log('');
      console.log('Ready to accept MCP connections and generate diagrams! 🎨');
      console.log('Press Ctrl+C to stop the server');
    });

    // Setup graceful shutdown
    const gracefulShutdown = async (signal: string) => {
      console.log(`\n🛑 Received ${signal}, shutting down gracefully...`);
      
      try {
        // Close all active sessions
        for (const [sessionId, sessionTransport] of Object.entries(transports)) {
          console.log(`Closing session: ${sessionId}`);
          await sessionTransport.server.close();
          sessionTransport.transport.close();
        }
        
        // Close HTTP server
        server.close(() => {
          console.log('✅ Server stopped successfully');
          process.exit(0);
        });
      } catch (error) {
        console.error('❌ Error during shutdown:', error);
        process.exit(1);
      }
    };

    // Handle shutdown signals
    process.on('SIGINT', () => gracefulShutdown('SIGINT'));
    process.on('SIGTERM', () => gracefulShutdown('SIGTERM'));

    // Handle uncaught exceptions
    process.on('uncaughtException', (error) => {
      console.error('❌ Uncaught Exception:', error);
      process.exit(1);
    });

    process.on('unhandledRejection', (reason, promise) => {
      console.error('❌ Unhandled Rejection at:', promise, 'reason:', reason);
      process.exit(1);
    });

  } catch (error) {
    console.error('❌ Failed to start server:', error);
    process.exit(1);
  }
};

// Start the server
main().catch((error) => {
  console.error('❌ Fatal error:', error);
  process.exit(1);
});