simple-init-service
Version:
MCP server for semantic codebase search using Transformers.js
569 lines (482 loc) • 17.4 kB
JavaScript
#!/usr/bin/env node
/**
* MCP Codebase Search Server - MCP Only Version
*
* Servidor que funciona exclusivamente via MCP stdio protocol
* Para uso com Trae IDE e outros clientes MCP
*
* Uso: node mcp-server.js --stdio
*/
const fs = require('fs-extra');
const path = require('path');
const { pipeline } = require('@xenova/transformers');
const CodebaseIndexer = require('./index-codebase.cjs');
// __dirname já está disponível em CommonJS
class CodebaseSearchEngine {
constructor() {
this.projectPath = process.cwd(); // Usar diretório atual
this.indexPath = path.join(this.projectPath, '.trae_index'); // Índice no projeto
this.documents = [];
this.embeddings = [];
this.metadata = null;
this.embedder = null;
this.isLoaded = false;
}
/**
* Inicializa o mecanismo de busca
*/
async initialize() {
try {
console.log('[LOG] Iniciando mecanismo de busca...');
// Carregar modelo de embeddings
console.log('[LOG] Carregando modelo de embeddings...');
await this.loadEmbedder();
console.log('[LOG] Modelo de embeddings carregado com sucesso');
// Carregar índice
console.log('[LOG] Carregando índice...');
await this.loadIndex();
console.log('[LOG] Índice carregado com sucesso');
this.isLoaded = true;
console.log('[LOG] Mecanismo de busca inicializado com sucesso');
} catch (error) {
console.error('[ERROR] Erro na inicialização:', error.message);
throw error;
}
}
/**
* Carrega o modelo de embeddings
*/
async loadEmbedder() {
try {
this.embedder = await pipeline(
'feature-extraction',
'Xenova/all-MiniLM-L6-v2',
{
quantized: true
}
);
} catch (error) {
throw error;
}
}
/**
* Carrega o índice do disco (cria automaticamente se não existir)
*/
async loadIndex() {
try {
// Verificar se o índice existe
const indexFile = path.join(this.indexPath, 'index.json');
const indexExists = await fs.pathExists(indexFile);
if (!indexExists) {
console.log('[LOG] Índice não encontrado, criando automaticamente...');
// Criar indexador e executar indexação
const indexer = new CodebaseIndexer(this.projectPath);
await indexer.run();
console.log('[LOG] Índice criado com sucesso');
}
// Carregar índice completo
const indexData = await fs.readJson(indexFile);
this.metadata = indexData.metadata;
this.documents = indexData.documents;
this.embeddings = indexData.embeddings;
} catch (error) {
throw error;
}
}
/**
* Calcula similaridade de cosseno entre dois vetores
*/
cosineSimilarity(vecA, vecB) {
if (vecA.length !== vecB.length) {
throw new Error('Vetores devem ter o mesmo tamanho');
}
let dotProduct = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < vecA.length; i++) {
dotProduct += vecA[i] * vecB[i];
normA += vecA[i] * vecA[i];
normB += vecB[i] * vecB[i];
}
return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}
/**
* Gera embedding para uma consulta
*/
async generateQueryEmbedding(query) {
try {
const output = await this.embedder(query, {
pooling: 'mean',
normalize: true
});
return Array.from(output.data);
} catch (error) {
throw error;
}
}
/**
* Realiza busca semântica
*/
async search(query, options = {}) {
if (!this.isLoaded) {
throw new Error('Mecanismo de busca não foi inicializado');
}
const {
limit = 10,
threshold = 0.3,
file_types = null
} = options;
try {
// Gerar embedding da consulta
const queryEmbedding = await this.generateQueryEmbedding(query);
// Calcular similaridades
const similarities = this.embeddings.map((embedding, index) => ({
index,
similarity: this.cosineSimilarity(queryEmbedding, embedding),
document: this.documents[index]
}));
// Filtrar por limiar de similaridade
let filteredResults = similarities.filter(result =>
result.similarity >= threshold
);
// Filtrar por tipos de arquivo se especificado
if (file_types && file_types.length > 0) {
filteredResults = filteredResults.filter(result =>
file_types.some(type =>
result.document.extension.toLowerCase().includes(type.toLowerCase()) ||
result.document.path.toLowerCase().includes(type.toLowerCase())
)
);
}
// Ordenar por similaridade (maior primeiro)
filteredResults.sort((a, b) => b.similarity - a.similarity);
// Limitar resultados
const limitedResults = filteredResults.slice(0, limit);
// Formatar resultados
const formattedResults = limitedResults.map(result => ({
path: result.document.path,
content: result.document.content,
similarity: Math.round(result.similarity * 100) / 100,
size: result.document.size,
extension: result.document.extension,
lastModified: result.document.lastModified,
// Extrair snippet relevante (primeiras 3 linhas)
snippet: result.document.content.split('\n').slice(0, 3).join('\n') +
(result.document.content.split('\n').length > 3 ? '...' : '')
}));
return {
query,
results: formattedResults,
total_results: formattedResults.length,
search_time: Date.now(),
metadata: {
threshold,
limit,
file_types,
index_info: {
total_documents: this.documents.length,
indexed_at: this.metadata.created,
project_path: this.metadata.projectPath
}
}
};
} catch (error) {
throw error;
}
}
/**
* Re-indexa o projeto atual
*/
async reindex(force = false) {
try {
const projectPath = this.projectPath;
// Verificar se já existe um índice (se não forçar)
const indexFile = path.join(this.indexPath, 'index.json');
const indexExists = await fs.pathExists(indexFile);
if (indexExists && !force) {
throw new Error('Índice já existe. Use force: true para sobrescrever.');
}
// Criar novo indexador para o projeto atual
const indexer = new CodebaseIndexer(projectPath);
// Executar indexação
await indexer.run();
// Recarregar o índice no mecanismo de busca
await this.loadIndex();
return {
success: true,
message: 'Re-indexação concluída com sucesso',
project_path: projectPath,
index_path: this.indexPath,
total_documents: this.documents.length,
timestamp: new Date().toISOString()
};
} catch (error) {
throw error;
}
}
/**
* Obtém estatísticas do índice
*/
getStats() {
if (!this.isLoaded) {
return { error: 'Índice não carregado' };
}
const extensionCounts = {};
this.documents.forEach(doc => {
const ext = doc.extension || 'sem extensão';
extensionCounts[ext] = (extensionCounts[ext] || 0) + 1;
});
return {
total_documents: this.documents.length,
embedding_dimension: this.embeddings.length > 0 ? this.embeddings[0].length : 0,
indexed_at: this.metadata.created,
project_path: this.metadata.projectPath,
model_name: 'Xenova/all-MiniLM-L6-v2',
file_types: extensionCounts,
index_size_mb: Math.round((JSON.stringify(this.embeddings).length / 1024 / 1024) * 100) / 100
};
}
}
// Inicializar mecanismo de busca
const searchEngine = new CodebaseSearchEngine();
// Servidor MCP via stdio
async function startMCPServer() {
try {
console.error('[LOG] Iniciando servidor MCP...');
// Inicializar mecanismo de busca em background (não bloquear)
console.error('[LOG] Inicializando mecanismo de busca em background...');
searchEngine.initialize().then(() => {
console.error('[LOG] Mecanismo de busca inicializado com sucesso');
}).catch(error => {
console.error('[ERROR] Erro na inicialização do mecanismo de busca:', error.message);
});
// Processar mensagens MCP via stdin/stdout
console.error('[LOG] Configurando stdin/stdout...');
process.stdin.resume(); // Garantir que stdin está ativo
let buffer = Buffer.alloc(0);
const processedIds = new Set();
// Manter o processo vivo
const keepAlive = setInterval(() => {
// Apenas manter o processo rodando
}, 30000); // 30 segundos
// Limpar interval ao sair
process.on('exit', () => {
clearInterval(keepAlive);
});
process.stdin.on('data', (chunk) => {
console.error('[LOG] Recebendo dados via stdin:', chunk.toString().trim());
buffer = Buffer.concat([buffer, chunk]);
parseMessages();
});
function parseMessages() {
while (true) {
// Procurar por Content-Length header
const headerEnd = buffer.indexOf('\r\n\r\n');
if (headerEnd === -1) break;
const headerStr = buffer.slice(0, headerEnd).toString();
const contentLengthMatch = headerStr.match(/Content-Length: (\d+)/);
if (!contentLengthMatch) {
console.error('[ERROR] Content-Length header não encontrado');
buffer = buffer.slice(headerEnd + 4);
continue;
}
const contentLength = parseInt(contentLengthMatch[1]);
const messageStart = headerEnd + 4;
if (buffer.length < messageStart + contentLength) {
// Mensagem incompleta, aguardar mais dados
break;
}
const messageBuffer = buffer.slice(messageStart, messageStart + contentLength);
const messageStr = messageBuffer.toString();
try {
console.error('[LOG] Processando mensagem:', messageStr);
const message = JSON.parse(messageStr);
console.error('[LOG] Mensagem parseada:', JSON.stringify(message, null, 2));
// Processar mensagem de forma não-bloqueante
handleMCPMessage(message).then(() => {
console.error('[LOG] Mensagem processada com sucesso');
}).catch(error => {
console.error('[ERROR] Erro ao processar mensagem:', error.message);
sendMCPError(-32603, 'Internal error', message?.id || null);
});
} catch (error) {
console.error('[ERROR] Erro ao fazer parse da mensagem:', error.message);
sendMCPError(-32700, 'Parse error', null);
}
buffer = buffer.slice(messageStart + contentLength);
}
}
process.stdin.on('end', () => {
console.error('[LOG] stdin finalizado, encerrando servidor');
process.exit(0);
});
process.on('SIGTERM', () => {
console.error('[LOG] Recebido SIGTERM, encerrando servidor');
process.exit(0);
});
process.on('SIGINT', () => {
console.error('[LOG] Recebido SIGINT, encerrando servidor');
process.exit(0);
});
console.error('[LOG] Servidor MCP pronto para receber mensagens');
} catch (error) {
console.error('[ERROR] Erro fatal no servidor MCP:', error.message);
process.exit(1);
}
}
// Manipular mensagens MCP
async function handleMCPMessage(message) {
try {
console.error('[LOG] Manipulando mensagem MCP:', message.method);
const { method, params, id } = message;
switch (method) {
case 'initialize':
console.error('[LOG] Processando inicialização MCP');
sendMCPResponse({
capabilities: {
tools: {
search: {
description: 'Busca semântica no codebase',
inputSchema: {
type: 'object',
properties: {
query: { type: 'string', description: 'Consulta de busca' },
limit: { type: 'number', default: 10, description: 'Máximo de resultados' },
threshold: { type: 'number', default: 0.3, description: 'Limiar de similaridade' }
},
required: ['query']
}
},
reindex: {
description: 'Re-indexa o projeto atual',
inputSchema: {
type: 'object',
properties: {
force: { type: 'boolean', default: false, description: 'Forçar re-indexação' }
}
}
},
stats: {
description: 'Obtém estatísticas do índice',
inputSchema: {
type: 'object',
properties: {}
}
}
}
}
}, id);
console.error('[LOG] Resposta de inicialização enviada');
break;
case 'tools/call':
if (params?.name === 'search') {
const { query, limit = 10, threshold = 0.3 } = params.arguments || {};
if (!query) {
sendMCPError(-32602, 'Invalid params: query is required', id);
return;
}
if (!searchEngine.isLoaded) {
sendMCPError(-32603, 'Search engine is still initializing. Please try again in a moment.', id);
return;
}
const results = await searchEngine.search(query, { limit, threshold });
sendMCPResponse({ content: [{ type: 'text', text: JSON.stringify(results, null, 2) }] }, id);
} else if (params?.name === 'reindex') {
const { force = false } = params.arguments || {};
if (!searchEngine.isLoaded) {
sendMCPError(-32603, 'Search engine is still initializing. Please try again in a moment.', id);
return;
}
const results = await searchEngine.reindex(force);
sendMCPResponse({ content: [{ type: 'text', text: JSON.stringify(results, null, 2) }] }, id);
} else if (params?.name === 'stats') {
if (!searchEngine.isLoaded) {
sendMCPResponse({ content: [{ type: 'text', text: JSON.stringify({ status: 'initializing', message: 'Search engine is still loading...' }, null, 2) }] }, id);
return;
}
const stats = searchEngine.getStats();
sendMCPResponse({ content: [{ type: 'text', text: JSON.stringify(stats, null, 2) }] }, id);
} else {
sendMCPError(-32601, `Method not found: ${params?.name}`, id);
}
break;
case 'tools/list':
sendMCPResponse({
tools: [
{
name: 'search',
description: 'Busca semântica no codebase',
inputSchema: {
type: 'object',
properties: {
query: { type: 'string', description: 'Consulta de busca' },
limit: { type: 'number', default: 10, description: 'Máximo de resultados' },
threshold: { type: 'number', default: 0.3, description: 'Limiar de similaridade' }
},
required: ['query']
}
},
{
name: 'reindex',
description: 'Re-indexa o projeto atual',
inputSchema: {
type: 'object',
properties: {
force: { type: 'boolean', default: false, description: 'Forçar re-indexação' }
}
}
},
{
name: 'stats',
description: 'Obtém estatísticas do índice',
inputSchema: {
type: 'object',
properties: {}
}
}
]
}, id);
break;
default:
sendMCPError(-32601, `Method not found: ${method}`, id);
}
} catch (error) {
sendMCPError(-32603, 'Internal error', message?.id || null);
}
}
// Enviar resposta MCP
function sendMCPResponse(result, id) {
const response = {
jsonrpc: '2.0',
result,
id
};
sendMessage(response);
}
// Enviar erro MCP
function sendMCPError(code, message, id) {
const response = {
jsonrpc: '2.0',
error: { code, message },
id
};
sendMessage(response);
}
// Enviar mensagem com Content-Length header
function sendMessage(message) {
const messageStr = JSON.stringify(message);
const contentLength = Buffer.byteLength(messageStr, 'utf8');
const header = `Content-Length: ${contentLength}\r\n\r\n`;
console.error('[LOG] Enviando mensagem MCP:', messageStr);
process.stdout.write(header + messageStr);
}
// Iniciar se executado diretamente
if (require.main === module) {
console.error('[LOG] Iniciando servidor MCP diretamente');
console.error('[LOG] Argumentos:', process.argv);
startMCPServer();
}
module.exports = {
startMCPServer,
CodebaseSearchEngine
};