flash-tool-oss/src/mcp-server.js

Fast AI-powered code completion using OpenAI API - No API key setup required!

#!/usr/bin/env node

const { spawn } = require('child_process');
const path = require('path');
const readline = require('readline');

class MCPServer {
    constructor() {
        this.pythonScriptPath = path.join(__dirname, 'extension.py');
    }

    // Function to call Python script
    callPythonScript(args) {
        return new Promise((resolve, reject) => {
            const pythonProcess = spawn('python', [this.pythonScriptPath, ...args]);
            
            let output = '';
            let errorOutput = '';
            
            pythonProcess.stdout.on('data', (data) => {
                output += data.toString();
            });
            
            pythonProcess.stderr.on('data', (data) => {
                errorOutput += data.toString();
            });
            
            pythonProcess.on('close', (code) => {
                if (code === 0) {
                    resolve(output.trim());
                } else {
                    reject(new Error(`Python script failed: ${errorOutput}`));
                }
            });
        });
    }

    // Process "using oss" commands
    async processUsingOss(message) {
        if (!message || typeof message !== 'string') {
            return null;
        }
        
        if (!message.toLowerCase().startsWith('using oss')) {
            return null;
        }
        
        const prompt = message.substring(9).trim(); // Remove "using oss "
        if (!prompt) {
            return "Please provide a prompt after 'using oss'";
        }
        
        try {
            return await this.callPythonScript(['chat', prompt]);
        } catch (error) {
            return `Error: ${error.message}`;
        }
    }

    // MCP Protocol Implementation
    async handleMCPRequest(request) {
        const { method, params } = request;
        
        switch (method) {
            case 'tools/call':
                const { name, arguments: args } = params;
                
                if (name === 'process_using_oss') {
                    const { message } = args;
                    const response = await this.processUsingOss(message);
                    
                    return {
                        content: [
                            {
                                type: 'text',
                                text: response || 'No response generated'
                            }
                        ]
                    };
                }
                break;
                
            case 'tools/list':
                return {
                    tools: [
                        {
                            name: 'process_using_oss',
                            description: 'Process "using oss" commands and return AI responses for code generation',
                            inputSchema: {
                                type: 'object',
                                properties: {
                                    message: {
                                        type: 'string',
                                        description: 'The full message including "using oss" prefix'
                                    }
                                },
                                required: ['message']
                            }
                        }
                    ]
                };
                
            case 'initialize':
                return {
                    protocolVersion: '2024-11-05',
                    capabilities: {
                        tools: {}
                    },
                    serverInfo: {
                        name: 'flash-tool-oss',
                        version: '1.0.0'
                    }
                };
                
            default:
                return { error: `Unknown method: ${method}` };
        }
    }

    // Start MCP server
    start() {
        const rl = readline.createInterface({
            input: process.stdin,
            output: process.stdout
        });

        console.error('MCP Server started for Flash Tool OSS');

        rl.on('line', async (line) => {
            try {
                const request = JSON.parse(line);
                const response = await this.handleMCPRequest(request);
                
                console.log(JSON.stringify({
                    jsonrpc: '2.0',
                    id: request.id,
                    result: response
                }));
            } catch (error) {
                console.error('Error processing request:', error);
                console.log(JSON.stringify({
                    jsonrpc: '2.0',
                    id: request?.id || null,
                    error: {
                        code: -32603,
                        message: error.message
                    }
                }));
            }
        });

        // Handle process termination
        process.on('SIGINT', () => {
            console.error('MCP Server shutting down...');
            process.exit(0);
        });
    }
}

// Start the server if this file is run directly
if (require.main === module) {
    const server = new MCPServer();
    server.start();
}

module.exports = MCPServer;