spex-mcp
Version:
MCP server for Figma SpeX plugin and Cursor AI integration
834 lines (713 loc) • 29.6 kB
JavaScript
import { WebSocketServer } from 'ws';
import fs from 'fs';
import path from 'path';
import { createWriteStream } from 'fs';
import { promisify } from 'util';
import { pipeline } from 'stream';
import { createReadStream } from 'fs';
import { exec } from 'child_process';
import AdmZip from 'adm-zip';
const pipelineAsync = promisify(pipeline);
const execAsync = promisify(exec);
/**
* FigmaPluginManager handles all communication with Figma SpeX plugins via WebSocket
*/
export class FigmaPluginManager {
constructor(port = 8080, force = false) {
this.figmaConnections = new Set();
this.wss = null;
this.fileTransfers = new Map(); // Track ongoing file transfers
this.port = port;
this.force = force;
}
/**
* Kill any process using the specified port
*/
async killPortProcess() {
try {
console.error(`🔍 Checking for processes using port ${this.port}...`);
let command;
const platform = process.platform;
if (platform === 'win32') {
// Windows: Use netstat to find process using the port
command = `netstat -ano | findstr :${this.port}`;
} else {
// Linux/macOS: Use lsof to find process using the port
command = `lsof -i :${this.port}`;
}
const { stdout, stderr } = await execAsync(command);
if (!stdout.trim()) {
console.error(`✅ Port ${this.port} is available`);
return;
}
console.error(`⚠️ Found process(es) using port ${this.port}:`);
console.error(stdout);
// Extract PIDs and kill them
if (platform === 'win32') {
// Windows: Extract PID from netstat output
const lines = stdout.split('\n');
const pids = new Set();
for (const line of lines) {
const match = line.trim().match(/\s+(\d+)\s*$/);
if (match) {
pids.add(match[1]);
}
}
for (const pid of pids) {
if (pid && pid !== '0') {
try {
console.error(`💀 Killing process ${pid}...`);
await execAsync(`taskkill /F /PID ${pid}`);
console.error(`✅ Process ${pid} killed`);
} catch (error) {
console.error(`⚠️ Failed to kill process ${pid}:`, error.message);
}
}
}
} else {
// Linux/macOS: Extract PID from lsof output
const lines = stdout.split('\n');
const pids = new Set();
for (const line of lines) {
const parts = line.trim().split(/\s+/);
if (parts.length >= 2 && parts[1].match(/^\d+$/)) {
pids.add(parts[1]);
}
}
for (const pid of pids) {
if (pid && pid !== '0') {
try {
console.error(`💀 Killing process ${pid}...`);
await execAsync(`kill -9 ${pid}`);
console.error(`✅ Process ${pid} killed`);
} catch (error) {
console.error(`⚠️ Failed to kill process ${pid}:`, error.message);
}
}
}
}
// Wait a moment for processes to fully terminate
await new Promise(resolve => setTimeout(resolve, 1000));
} catch (error) {
// If the command fails, it usually means no process is using the port
if (error.code === 1 || error.message.includes('No such process')) {
console.error(`✅ Port ${this.port} is available`);
} else {
console.error(`⚠️ Error checking port ${this.port}:`, error.message);
}
}
}
/**
* Start WebSocket server for Figma SpeX plugin connections
*/
async setupWebSocketServer() {
// Kill existing processes on the port if --force is specified
if (this.force) {
await this.killPortProcess();
}
this.wss = new WebSocketServer({ port: this.port });
this.wss.on('connection', (ws) => {
console.error('Figma SpeX plugin connected');
this.figmaConnections.add(ws);
ws.on('message', async (data, isBinary) => {
try {
if (isBinary) {
// Handle binary data (file uploads)
await this.handleBinaryMessage(ws, data);
} else {
// Handle text data (JSON messages)
const textData = Buffer.isBuffer(data) ? data.toString('utf8') : data;
const message = JSON.parse(textData);
console.error('Received from Figma:', message);
this.handleFigmaMessage(ws, message);
}
} catch (error) {
console.error('Error parsing message from Figma:', error);
}
});
ws.on('close', () => {
console.error('Figma SpeX plugin disconnected');
this.figmaConnections.delete(ws);
// Clean up any ongoing file transfers for this connection
this.cleanupFileTransfers(ws);
});
ws.on('error', (error) => {
console.error('WebSocket error:', error);
this.figmaConnections.delete(ws);
this.cleanupFileTransfers(ws);
});
// Send welcome message
this.sendWelcomeMessage(ws);
});
console.error(`WebSocket server started on port ${this.port} for Figma SpeX plugins`);
}
/**
* Handle different types of messages from Figma SpeX plugins
*/
handleFigmaMessage(ws, message) {
switch (message.name || message.type) {
case 'hello-world':
this.handleHelloWorld(ws);
break;
case 'plugin-ready':
this.handlePluginReady(ws, message);
break;
case 'specs-data':
this.handleSpecsData(ws, message.data);
break;
case 'file-upload-start':
this.handleFileUploadStart(ws, message.data);
break;
case 'file-upload-end':
this.handleFileUploadEnd(ws, message.data);
break;
default:
console.error('Unknown message type:', message.type || message.name);
}
}
/**
* Handle binary message data (file chunks)
*/
async handleBinaryMessage(ws, binaryData) {
try {
const transferId = this.getActiveTransferForConnection(ws);
if (!transferId) {
console.error('Received binary data without active transfer');
return;
}
const transfer = this.fileTransfers.get(transferId);
if (!transfer) {
console.error('Invalid transfer ID');
return;
}
// Write binary chunk to file
if (!transfer.writeStream.write(binaryData)) {
// Handle backpressure
await new Promise(resolve => transfer.writeStream.once('drain', resolve));
}
transfer.receivedBytes += binaryData.length;
console.error(`Received chunk: ${binaryData.length} bytes (Total: ${transfer.receivedBytes}/${transfer.expectedSize})`);
// Mark transfer as ready for completion (let file-upload-end handle the processing)
if (transfer.receivedBytes >= transfer.expectedSize) {
transfer.isComplete = true;
console.error('File transfer completed, waiting for file-upload-end message');
}
} catch (error) {
console.error('Error handling binary message:', error);
this.sendErrorResponse(ws, 'Error processing binary data: ' + error.message);
}
}
/**
* Handle file upload start message
*/
handleFileUploadStart(ws, data) {
try {
const { filename, fileSize, transferId } = data;
if (!filename || !fileSize || !transferId) {
throw new Error('Missing required file upload parameters');
}
const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
const safeFilename = filename || `specs-${timestamp}.zip`;
const zipPath = path.join('/tmp', safeFilename);
// Ensure /tmp directory exists
if (!fs.existsSync('/tmp')) {
fs.mkdirSync('/tmp', { recursive: true });
}
// Create write stream
const writeStream = createWriteStream(zipPath);
// Store transfer info
this.fileTransfers.set(transferId, {
ws: ws,
filename: safeFilename,
zipPath: zipPath,
expectedSize: fileSize,
receivedBytes: 0,
writeStream: writeStream,
startTime: Date.now(),
isComplete: false
});
console.error(`Starting file upload: ${safeFilename} (${fileSize} bytes)`);
// Send acknowledgment
ws.send(JSON.stringify({
type: 'file-upload-ready',
transferId: transferId,
message: 'Ready to receive file data',
timestamp: new Date().toISOString()
}));
} catch (error) {
console.error('Error starting file upload:', error);
this.sendErrorResponse(ws, 'Error starting file upload: ' + error.message);
}
}
/**
* Handle file upload end message
*/
async handleFileUploadEnd(ws, data) {
try {
const { transferId } = data;
if (!transferId || !this.fileTransfers.has(transferId)) {
throw new Error('Invalid or missing transfer ID');
}
const transfer = this.fileTransfers.get(transferId);
// Check if file transfer is actually complete
if (!transfer.isComplete) {
throw new Error(`File transfer incomplete: ${transfer.receivedBytes}/${transfer.expectedSize} bytes received`);
}
await this.finishFileTransfer(ws, transferId);
} catch (error) {
console.error('Error ending file upload:', error);
this.sendErrorResponse(ws, 'Error ending file upload: ' + error.message);
}
}
/**
* Finish file transfer and process the uploaded file
*/
async finishFileTransfer(ws, transferId) {
const transfer = this.fileTransfers.get(transferId);
if (!transfer) {
throw new Error('Transfer not found');
}
try {
// Close the write stream
transfer.writeStream.end();
// Wait for stream to finish
await new Promise((resolve, reject) => {
transfer.writeStream.on('finish', resolve);
transfer.writeStream.on('error', reject);
});
console.error(`File upload completed: ${transfer.filename}`);
// Process the uploaded zip file
const result = await this.processUploadedZipFile(transfer);
// Send success response
ws.send(JSON.stringify({
type: 'specs-data-response',
success: true,
message: 'File uploaded and processed successfully',
data: result,
transferId: transferId,
timestamp: new Date().toISOString(),
}));
// Clean up transfer
this.fileTransfers.delete(transferId);
} catch (error) {
console.error('Error finishing file transfer:', error);
// Clean up on error
try {
if (fs.existsSync(transfer.zipPath)) {
fs.unlinkSync(transfer.zipPath);
}
} catch (cleanupError) {
console.error('Error during cleanup:', cleanupError);
}
this.fileTransfers.delete(transferId);
throw error;
}
}
/**
* Process uploaded file (supports zip files and other formats)
*/
async processUploadedZipFile(transfer) {
const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
const extractPath = path.join('/tmp', `extracted-${timestamp}`);
try {
console.error(`Processing uploaded file: ${transfer.filename} (${transfer.expectedSize} bytes)`);
// Verify file exists and get actual size
if (!fs.existsSync(transfer.zipPath)) {
throw new Error(`File not found: ${transfer.zipPath}`);
}
const stats = fs.statSync(transfer.zipPath);
// Validate that we received all expected bytes
if (stats.size !== transfer.expectedSize) {
throw new Error(`File size mismatch: expected ${transfer.expectedSize} bytes, got ${stats.size} bytes`);
}
// Read file and check for zip signature
const fileBuffer = fs.readFileSync(transfer.zipPath);
// Check for zip file signature (PK\003\004 or PK\005\006)
const zipSignature1 = Buffer.from([0x50, 0x4B, 0x03, 0x04]); // Local file header
const zipSignature2 = Buffer.from([0x50, 0x4B, 0x05, 0x06]); // End of central directory
const hasValidSignature = fileBuffer.indexOf(zipSignature1) === 0 || fileBuffer.indexOf(zipSignature2) >= 0;
if (hasValidSignature && stats.size >= 22) {
// Process as zip file
return await this.processZipFile(transfer, fileBuffer, extractPath, stats);
} else {
// Process as regular file (text, JSON, etc.)
return await this.processRegularFile(transfer, fileBuffer, stats);
}
} catch (error) {
// Clean up on error
try {
if (fs.existsSync(transfer.zipPath)) {
fs.unlinkSync(transfer.zipPath);
}
if (fs.existsSync(extractPath)) {
fs.rmSync(extractPath, { recursive: true, force: true });
}
} catch (cleanupError) {
console.error('Error during cleanup:', cleanupError);
}
throw new Error(`Failed to process file: ${error.message}`);
}
}
/**
* Process zip file extraction
*/
async processZipFile(transfer, fileBuffer, extractPath, stats) {
console.error(`Extracting zip file to: ${extractPath}`);
const zip = new AdmZip(transfer.zipPath);
// Create extraction directory
if (!fs.existsSync(extractPath)) {
fs.mkdirSync(extractPath, { recursive: true });
}
// Extract all files
zip.extractAllTo(extractPath, true);
// Get list of extracted files
const extractedFiles = this.getDirectoryContents(extractPath);
console.error(`Zip file extracted successfully. Files extracted: ${extractedFiles.length}`);
return {
type: 'zip',
originalFilename: transfer.filename,
filePath: transfer.zipPath,
extractPath: extractPath,
fileSize: stats.size,
uploadDuration: Date.now() - transfer.startTime,
extractedFiles: extractedFiles,
processedAt: new Date().toISOString()
};
}
/**
* Process regular file (text, JSON, etc.)
*/
async processRegularFile(transfer, fileBuffer, stats) {
console.error(`Processing as regular file`);
// Try to detect file type
let fileType = 'binary';
let content = null;
try {
// Try to parse as text
content = fileBuffer.toString('utf8');
fileType = 'text';
// Try to parse as JSON
JSON.parse(content);
fileType = 'json';
} catch (e) {
// Not JSON, keep as text or binary
if (fileType === 'text') {
// Check if it looks like valid text
const hasNullBytes = fileBuffer.indexOf(0) !== -1;
if (hasNullBytes) {
fileType = 'binary';
content = null;
}
}
}
console.error(`File processed as ${fileType} (${stats.size} bytes)`);
return {
type: 'file',
fileType: fileType,
originalFilename: transfer.filename,
filePath: transfer.zipPath,
fileSize: stats.size,
uploadDuration: Date.now() - transfer.startTime,
content: fileType === 'text' || fileType === 'json' ? content : null,
contentPreview: content ? content.substring(0, 200) + (content.length > 200 ? '...' : '') : 'Binary content',
processedAt: new Date().toISOString()
};
}
/**
* Get active transfer ID for a WebSocket connection
*/
getActiveTransferForConnection(ws) {
for (const [transferId, transfer] of this.fileTransfers.entries()) {
if (transfer.ws === ws) {
return transferId;
}
}
return null;
}
/**
* Clean up file transfers for a disconnected WebSocket
*/
cleanupFileTransfers(ws) {
const transfersToDelete = [];
for (const [transferId, transfer] of this.fileTransfers.entries()) {
if (transfer.ws === ws) {
// Close write stream if open
if (transfer.writeStream && !transfer.writeStream.destroyed) {
transfer.writeStream.destroy();
}
// Delete incomplete file
try {
if (fs.existsSync(transfer.zipPath)) {
fs.unlinkSync(transfer.zipPath);
}
} catch (error) {
console.error('Error cleaning up file:', error);
}
transfersToDelete.push(transferId);
}
}
transfersToDelete.forEach(id => this.fileTransfers.delete(id));
}
/**
* Send error response
*/
sendErrorResponse(ws, errorMessage) {
ws.send(JSON.stringify({
type: 'specs-data-response',
success: false,
message: 'Error processing specs data',
error: errorMessage,
timestamp: new Date().toISOString(),
}));
}
/**
* Handle plugin ready message
*/
handlePluginReady(ws, message) {
console.error('Figma plugin is ready:', message.timestamp);
ws.send(JSON.stringify({
type: 'plugin-ready-response',
message: 'Server received plugin ready signal',
serverTime: new Date().toISOString(),
clientTime: message.timestamp,
}));
}
/**
* Send hello world response to Figma plugin
*/
handleHelloWorld(ws) {
console.error('Hello world received from Figma plugin');
ws.send(JSON.stringify({
type: 'function-response',
message: 'Hello world received at ' + new Date().toISOString(),
timestamp: new Date().toISOString(),
}));
}
/**
* Handle specs data from Figma plugin - now supports both JSON and binary file uploads
*/
async handleSpecsData(ws, data) {
try {
console.error('Specs data received:', typeof data, data?.type || 'unknown type');
if (!data) {
throw new Error('No data received');
}
// Handle regular JSON specs data
console.error('Processing regular specs data:', data);
ws.send(JSON.stringify({
type: 'specs-data-response',
success: true,
message: 'Specs data received and processed successfully',
data: {
type: 'json',
receivedAt: new Date().toISOString(),
dataSize: JSON.stringify(data).length
},
timestamp: new Date().toISOString(),
}));
} catch (error) {
console.error('Error processing specs data:', error);
this.sendErrorResponse(ws, error.message);
}
}
/**
* Recursively get directory contents
*/
getDirectoryContents(dirPath, relativePath = '') {
const files = [];
try {
const items = fs.readdirSync(dirPath);
for (const item of items) {
const fullPath = path.join(dirPath, item);
const relativeFilePath = path.join(relativePath, item);
const stats = fs.statSync(fullPath);
if (stats.isDirectory()) {
files.push({
type: 'directory',
name: item,
path: relativeFilePath,
size: 0
});
// Recursively get contents of subdirectory
const subFiles = this.getDirectoryContents(fullPath, relativeFilePath);
files.push(...subFiles);
} else {
files.push({
type: 'file',
name: item,
path: relativeFilePath,
size: stats.size,
modified: stats.mtime.toISOString()
});
}
}
} catch (error) {
console.error(`Error reading directory ${dirPath}:`, error);
}
return files;
}
/**
* Send welcome message to newly connected plugin
*/
sendWelcomeMessage(ws) {
ws.send(JSON.stringify({
type: 'welcome',
message: 'Connected to MCP Figma Server',
timestamp: new Date().toISOString(),
}));
}
/**
* Broadcast message to all connected Figma plugins
*/
broadcastToFigmaPlugins(message) {
const promises = Array.from(this.figmaConnections).map(ws =>
this.sendMessageToFigma(ws, message)
);
return Promise.allSettled(promises);
}
/**
* Send message to specific Figma plugin with timeout
*/
sendMessageToFigma(ws, message) {
return new Promise((resolve, reject) => {
if (ws.readyState !== ws.OPEN) {
reject(new Error('WebSocket not open'));
return;
}
const timeout = setTimeout(() => {
reject(new Error('Timeout waiting for Figma response'));
}, 5000);
const responseHandler = (data) => {
try {
const response = JSON.parse(data);
if (response.type === message.type + '-response') {
clearTimeout(timeout);
ws.off('message', responseHandler);
resolve(response.data);
}
} catch (error) {
// Ignore parsing errors for other messages
}
};
ws.on('message', responseHandler);
ws.send(JSON.stringify(message));
});
}
/**
* Request a specific spec file or image from connected Figma plugins
* Enhanced to support image requests with additional parameters
*/
async requestSpecFile(requestType, fileName = null, requestData = null) {
if (this.figmaConnections.size === 0) {
return {
success: false,
error: 'No Figma plugins connected'
};
}
const requestId = `req_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
// Build message based on request type
let message = {
type: 'spec-request',
requestType: requestType,
requestId: requestId,
timestamp: new Date().toISOString(),
};
// Add fileName for file-based requests
if (fileName) {
message.fileName = fileName;
}
// Add requestData for complex requests (like image generation)
if (requestData) {
try {
// If requestData is a string, try to parse it as JSON
message.requestData = typeof requestData === 'string' ? JSON.parse(requestData) : requestData;
} catch (error) {
// If parsing fails, treat as raw string
message.requestData = requestData;
}
}
console.error(`Requesting: ${requestType}${fileName ? ` (${fileName})` : ''}${requestData ? ' with parameters' : ''}`);
// Send request to the first connected plugin (assuming single plugin for now)
const ws = Array.from(this.figmaConnections)[0];
if (!ws || ws.readyState !== ws.OPEN) {
return {
success: false,
error: 'No active WebSocket connection to Figma plugin'
};
}
return new Promise((resolve) => {
const timeout = setTimeout(() => {
resolve({
success: false,
error: 'Timeout waiting for response from Figma plugin'
});
}, 15000); // 15 second timeout for image generation
const responseHandler = (event) => {
try {
const response = JSON.parse(event.data);
if (response.type === 'spec-response' && response.requestId === requestId) {
clearTimeout(timeout);
ws.removeEventListener('message', responseHandler);
if (response.success) {
resolve({
success: true,
data: response.data
});
} else {
resolve({
success: false,
error: response.error || 'Unknown error from plugin'
});
}
}
} catch (error) {
// Ignore parsing errors for other messages
}
};
ws.addEventListener('message', responseHandler);
ws.send(JSON.stringify(message));
});
}
/**
* Close all connections and shutdown server
*/
shutdown() {
return new Promise((resolve) => {
if (this.wss) {
// Clean up all file transfers
for (const [transferId, transfer] of this.fileTransfers.entries()) {
if (transfer.writeStream && !transfer.writeStream.destroyed) {
transfer.writeStream.destroy();
}
try {
if (fs.existsSync(transfer.zipPath)) {
fs.unlinkSync(transfer.zipPath);
}
} catch (error) {
console.error('Error cleaning up file during shutdown:', error);
}
}
this.fileTransfers.clear();
// Close all client connections first
this.figmaConnections.forEach(ws => {
if (ws.readyState === ws.OPEN) {
ws.close();
}
});
this.figmaConnections.clear();
// Close the WebSocket server and wait for it to fully close
this.wss.close((error) => {
if (error) {
console.error('Error closing WebSocket server:', error);
} else {
console.error('Figma WebSocket server shut down completely');
}
this.wss = null;
resolve();
});
} else {
resolve();
}
});
}
}