Adding as a Streamable HTTP MCP Server

src/index.ts

@@ -1,6 +1,9 @@
 #!/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 { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import {
   CallToolRequestSchema,
   ErrorCode,
@@ -8,12 +11,11 @@ import {
   ListToolsRequestSchema,
   McpError,
   ReadResourceRequestSchema,
+  InitializeRequestSchema,
 } from '@modelcontextprotocol/sdk/types.js';
 
 // Import tools and utilities
 import { createDiagram, validateCreateDiagramInput, getSupportedDiagramTypes, getDiagramTypeDescription } from './tools/create-diagram.js';
-import { createFileConfig, findDiagramFiles, getDiagramFilesWithMetadata } from './utils/file-manager.js';
-import { createVSCodeConfig, openDiagramInVSCode, setupVSCodeEnvironment } from './utils/vscode-integration.js';
 import { DiagramType, DiagramFormat } from './types/diagram-types.js';
 
 // Functional types and configurations
@@ -21,6 +23,7 @@ type ServerConfig = Readonly<{
   name: string;
   version: string;
   workspaceRoot: string;
+  httpPort: number;
   capabilities: Readonly<{
     resources: Record<string, unknown>;
     tools: Record<string, unknown>;
@@ -32,22 +35,19 @@ type ResourceHandler = (uri: string) => Promise<any>;
 
 type ToolHandlers = Readonly<{
   create_diagram: ToolHandler;
-  list_diagrams: ToolHandler;
-  open_diagram_in_vscode: ToolHandler;
-  setup_vscode_environment: ToolHandler;
   get_diagram_types: ToolHandler;
 }>;
 
 type ResourceHandlers = Readonly<{
-  'diagrams://workspace/list': ResourceHandler;
   'diagrams://types/supported': ResourceHandler;
 }>;
 
 // Pure function to create server configuration
-const createServerConfig = (workspaceRoot?: string): ServerConfig => ({
+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: {},
@@ -58,7 +58,7 @@ const createServerConfig = (workspaceRoot?: string): ServerConfig => ({
 const createToolDefinitions = () => [
   {
     name: 'create_diagram',
-    description: 'Create a new diagram of specified type (BPMN, UML, ER, Network, Architecture, etc.)',
+    description: 'Create a new diagram of specified type (BPMN, UML, ER, Network, Architecture, etc.) with AI-powered generation',
     inputSchema: {
       type: 'object',
       properties: {
@@ -79,26 +79,37 @@ const createToolDefinitions = () => [
         },
         description: {
           type: 'string',
-          description: 'Description of the diagram'
+          description: 'Natural language description of the diagram to generate using AI'
         },
         outputPath: {
           type: 'string',
           description: 'Output directory path (relative to workspace)'
         },
-        // BPMN specific parameters
+        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'
+          description: 'Name of the BPMN process (legacy)'
         },
         tasks: {
           type: 'array',
           items: { type: 'string' },
-          description: 'List of tasks for BPMN process'
+          description: 'List of tasks for BPMN process (legacy)'
         },
         gatewayType: {
           type: 'string',
           enum: ['exclusive', 'parallel'],
-          description: 'Type of gateway for BPMN process'
+          description: 'Type of gateway for BPMN process (legacy)'
         },
         branches: {
           type: 'array',
@@ -106,90 +117,42 @@ const createToolDefinitions = () => [
             type: 'array',
             items: { type: 'string' }
           },
-          description: 'Branches for BPMN gateway'
+          description: 'Branches for BPMN gateway (legacy)'
         },
         beforeGateway: {
           type: 'array',
           items: { type: 'string' },
-          description: 'Tasks before gateway'
+          description: 'Tasks before gateway (legacy)'
         },
         afterGateway: {
           type: 'array',
           items: { type: 'string' },
-          description: 'Tasks after gateway'
+          description: 'Tasks after gateway (legacy)'
         },
-        // UML specific parameters
         classes: {
           type: 'array',
           items: { type: 'string' },
-          description: 'List of classes for UML class diagram'
+          description: 'List of classes for UML class diagram (legacy)'
         },
-        // ER specific parameters
         entities: {
           type: 'array',
           items: { type: 'string' },
-          description: 'List of entities for ER diagram'
+          description: 'List of entities for ER diagram (legacy)'
         },
-        // Network/Architecture specific parameters
         components: {
           type: 'array',
           items: { type: 'string' },
-          description: 'List of components for network or architecture diagram'
+          description: 'List of components for network or architecture diagram (legacy)'
         },
-        // Flowchart specific parameters
         processes: {
           type: 'array',
           items: { type: 'string' },
-          description: 'List of processes for flowchart'
+          description: 'List of processes for flowchart (legacy)'
         }
       },
       required: ['name', 'type']
     }
   },
-  {
-    name: 'list_diagrams',
-    description: 'List all diagram files in the workspace',
-    inputSchema: {
-      type: 'object',
-      properties: {
-        workspaceRoot: {
-          type: 'string',
-          description: 'Workspace root directory (optional)'
-        }
-      }
-    }
-  },
-  {
-    name: 'open_diagram_in_vscode',
-    description: 'Open a diagram file in VSCode with draw.io extension',
-    inputSchema: {
-      type: 'object',
-      properties: {
-        filePath: {
-          type: 'string',
-          description: 'Path to the diagram file'
-        },
-        workspaceRoot: {
-          type: 'string',
-          description: 'Workspace root directory (optional)'
-        }
-      },
-      required: ['filePath']
-    }
-  },
-  {
-    name: 'setup_vscode_environment',
-    description: 'Setup VSCode environment for draw.io (install extension if needed)',
-    inputSchema: {
-      type: 'object',
-      properties: {
-        workspaceRoot: {
-          type: 'string',
-          description: 'Workspace root directory (optional)'
-        }
-      }
-    }
-  },
   {
     name: 'get_diagram_types',
     description: 'Get list of supported diagram types with descriptions',
@@ -202,12 +165,6 @@ const createToolDefinitions = () => [
 
 // Pure function to create resource definitions
 const createResourceDefinitions = () => [
-  {
-    uri: 'diagrams://workspace/list',
-    name: 'Workspace Diagrams',
-    mimeType: 'application/json',
-    description: 'List of all diagram files in the workspace'
-  },
   {
     uri: 'diagrams://types/supported',
     name: 'Supported Diagram Types',
@@ -241,86 +198,6 @@ const createToolHandlers = (config: ServerConfig): ToolHandlers => ({
     };
   },
 
-  list_diagrams: async (args: any) => {
-    const workspaceRoot = args?.workspaceRoot || config.workspaceRoot;
-    const fileConfig = createFileConfig(workspaceRoot);
-    const diagrams = await getDiagramFilesWithMetadata(fileConfig)();
-
-    return {
-      content: [
-        {
-          type: 'text',
-          text: JSON.stringify({
-            success: true,
-            count: diagrams.length,
-            diagrams: diagrams.map(d => ({
-              path: d.relativePath,
-              format: d.format,
-              size: d.stats.size,
-              modified: d.stats.mtime
-            }))
-          }, null, 2)
-        }
-      ]
-    };
-  },
-
-  open_diagram_in_vscode: async (args: any) => {
-    if (!args?.filePath) {
-      throw new McpError(
-        ErrorCode.InvalidParams,
-        'filePath is required'
-      );
-    }
-
-    const workspaceRoot = args.workspaceRoot || config.workspaceRoot;
-    const vscodeConfig = createVSCodeConfig(workspaceRoot);
-
-    try {
-      await openDiagramInVSCode(vscodeConfig)(args.filePath);
-      
-      return {
-        content: [
-          {
-            type: 'text',
-            text: JSON.stringify({
-              success: true,
-              message: `Successfully opened diagram: ${args.filePath}`
-            }, null, 2)
-          }
-        ]
-      };
-    } catch (error) {
-      return {
-        content: [
-          {
-            type: 'text',
-            text: JSON.stringify({
-              success: false,
-              message: `Failed to open diagram: ${error}`
-            }, null, 2)
-          }
-        ]
-      };
-    }
-  },
-
-  setup_vscode_environment: async (args: any) => {
-    const workspaceRoot = args?.workspaceRoot || config.workspaceRoot;
-    const vscodeConfig = createVSCodeConfig(workspaceRoot);
-
-    const result = await setupVSCodeEnvironment(vscodeConfig)();
-
-    return {
-      content: [
-        {
-          type: 'text',
-          text: JSON.stringify(result, null, 2)
-        }
-      ]
-    };
-  },
-
   get_diagram_types: async () => {
     const types = getSupportedDiagramTypes();
     const typesWithDescriptions = types.map(type => ({
@@ -344,21 +221,6 @@ const createToolHandlers = (config: ServerConfig): ToolHandlers => ({
 
 // Pure function to create resource handlers
 const createResourceHandlers = (config: ServerConfig): ResourceHandlers => ({
-  'diagrams://workspace/list': async () => {
-    const fileConfig = createFileConfig(config.workspaceRoot);
-    const diagrams = await getDiagramFilesWithMetadata(fileConfig)();
-    
-    return {
-      contents: [
-        {
-          uri: 'diagrams://workspace/list',
-          mimeType: 'application/json',
-          text: JSON.stringify(diagrams, null, 2)
-        }
-      ]
-    };
-  },
-
   'diagrams://types/supported': async () => {
     const types = getSupportedDiagramTypes();
     const typesWithDescriptions = types.map(type => ({
@@ -449,16 +311,10 @@ const setupResourceRequestHandlers = (server: Server, resourceHandlers: Resource
 // Pure function to setup error handling
 const setupErrorHandling = (server: Server) => {
   server.onerror = (error) => console.error('[MCP Error]', error);
-  
-  process.on('SIGINT', async () => {
-    await server.close();
-    process.exit(0);
-  });
-
   return server;
 };
 
-// Pure function to create and configure server
+// Pure function to create and configure MCP server
 const createMCPServer = (config: ServerConfig): Server => {
   const server = new Server(
     {
@@ -481,19 +337,313 @@ const createMCPServer = (config: ServerConfig): Server => {
   return serverWithErrorHandling;
 };
 
-// Pure function to run the server
-const runServer = async (server: Server): Promise<void> => {
-  const transport = new StdioServerTransport();
-  await server.connect(transport);
-  console.error('Draw.io MCP server running on stdio');
+// 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;
 };
 
-// Main execution - functional composition
+// 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> => {
-  const config = createServerConfig();
-  const server = createMCPServer(config);
-  await runServer(server);
+  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(console.error);
+main().catch((error) => {
+  console.error('❌ Fatal error:', error);
+  process.exit(1);
+});