simple-init-service
Version:
MCP server for semantic codebase search using Transformers.js
391 lines (343 loc) • 11.9 kB
JavaScript
/**
* MCP Codebase Search Server (Trae compatível)
* - Framing via Content-Length (LSP/MCP)
* - stdin binário (sem setEncoding utf8)
* - initialize rápido (indexação assíncrona para evitar timeout)
* - logs em stderr somente
* - tools: search_codebase, echo
*/
const fs = require('fs-extra');
const path = require('path');
const glob = require('glob');
// ----------------------------------------------------
// Framing JSON-RPC: sendMessage / parse por Content-Length
// ----------------------------------------------------
function sendMessage(msg) {
try {
const json = JSON.stringify(msg);
const contentLength = Buffer.byteLength(json, 'utf8');
const frame = `Content-Length: ${contentLength}\r\n\r\n${json}`;
process.stdout.write(frame);
// Debug opcional:
// console.error('[DBG] >>>', json);
} catch (err) {
console.error('[ERROR] sendMessage falhou:', err?.message || err);
}
}
function parseMessages(buffer, onMessage) {
while (true) {
const headerEnd = buffer.indexOf('\r\n\r\n');
if (headerEnd === -1) break;
const header = buffer.slice(0, headerEnd).toString('utf8');
const m = header.match(/Content-Length:\s*(\d+)/i);
if (!m) {
console.error('[ERROR] Cabeçalho sem Content-Length. Header recebido:', header);
// descarta cabeçalho inválido
buffer = buffer.slice(headerEnd + 4);
continue;
}
const contentLength = parseInt(m[1], 10);
const total = headerEnd + 4 + contentLength;
if (buffer.length < total) break; // aguardar completar o body
const body = buffer.slice(headerEnd + 4, total).toString('utf8');
buffer = buffer.slice(total);
try {
const message = JSON.parse(body);
onMessage(message);
} catch (err) {
console.error('[ERROR] JSON inválido:', err?.message || err, 'Body:', body);
}
}
return buffer;
}
// ----------------------------------------------------
// Engine simples de busca em codebase
// ----------------------------------------------------
class SimpleCodebaseSearch {
constructor() {
this.projectPath = process.cwd();
this.documents = [];
this.isLoaded = false;
this.loadingError = null;
}
async initialize() {
console.error('[LOG] Iniciando indexação da codebase…');
const start = Date.now();
try {
await this.loadDocuments();
this.isLoaded = true;
console.error(`[LOG] Indexação concluída em ${(Date.now() - start)}ms com ${this.documents.length} docs`);
} catch (err) {
this.loadingError = err;
console.error('[ERROR] Falha na indexação:', err?.message || err);
}
}
async loadDocuments() {
const patterns = [
'**/*.js', '**/*.ts', '**/*.jsx', '**/*.tsx',
'**/*.py', '**/*.java', '**/*.cpp', '**/*.c',
'**/*.cs', '**/*.php', '**/*.rb', '**/*.go',
'**/*.rs', '**/*.swift', '**/*.kt', '**/*.scala',
'**/*.html', '**/*.css', '**/*.scss', '**/*.sass',
'**/*.json', '**/*.xml', '**/*.yaml', '**/*.yml',
'**/*.md', '**/*.txt', '**/*.sql'
];
const ignorePatterns = [
'node_modules/**', '.git/**', 'dist/**', 'build/**',
'.next/**', '.nuxt/**', 'coverage/**', '.nyc_output/**',
'*.min.js', '*.bundle.js', '.trae_index/**'
];
this.documents = [];
for (const pattern of patterns) {
const files = glob.sync(pattern, {
cwd: this.projectPath,
ignore: ignorePatterns,
absolute: true
});
for (const filePath of files) {
try {
const stats = await fs.stat(filePath);
if (stats.size > 1024 * 1024) continue; // >1MB ignora
const content = await fs.readFile(filePath, 'utf-8');
const relativePath = path.relative(this.projectPath, filePath);
this.documents.push({
path: relativePath,
fullPath: filePath,
content,
size: stats.size,
extension: path.extname(filePath),
lastModified: stats.mtime.toISOString()
});
} catch (e) {
// ignora arquivo não legível
continue;
}
}
}
console.error(`[LOG] Carregados ${this.documents.length} documentos`);
}
search(query, options = {}) {
if (this.loadingError) {
throw new Error(`Falha na indexação: ${this.loadingError.message || this.loadingError}`);
}
if (!this.isLoaded) {
// você pode escolher retornar parcial; aqui, aviso explícito:
throw new Error('Indexação em andamento. Tente novamente em alguns segundos.');
}
const { limit = 10, file_types = null } = options;
const queryLower = (query || '').toLowerCase().trim();
if (!queryLower) {
return { query, results: [], total_results: 0, search_time: Date.now() };
}
let results = this.documents
.map(doc => {
const contentLower = doc.content.toLowerCase();
const pathLower = doc.path.toLowerCase();
let score = 0;
const queryWords = queryLower.split(/\s+/);
for (const word of queryWords) {
if (!word) continue;
if (contentLower.includes(word)) score += 2;
if (pathLower.includes(word)) score += 1;
}
return { ...doc, score };
})
.filter(doc => doc.score > 0);
if (file_types && file_types.length > 0) {
results = results.filter(doc =>
file_types.some(type =>
doc.extension.toLowerCase().includes(type.toLowerCase()) ||
doc.path.toLowerCase().includes(type.toLowerCase())
)
);
}
results.sort((a, b) => b.score - a.score);
results = results.slice(0, limit);
const formattedResults = results.map(doc => ({
path: doc.path,
content: doc.content,
similarity: Math.min(doc.score / 10, 1),
size: doc.size,
extension: doc.extension,
lastModified: doc.lastModified,
snippet: doc.content.split('\n').slice(0, 3).join('\n') +
(doc.content.split('\n').length > 3 ? '...' : '')
}));
return {
query,
results: formattedResults,
total_results: formattedResults.length,
search_time: Date.now()
};
}
}
// ----------------------------------------------------
// Servidor MCP
// ----------------------------------------------------
class MCPServer {
constructor() {
this.searchEngine = new SimpleCodebaseSearch();
this.indexingPromise = null; // para acompanhar status da indexação
}
// Não bloquear o initialize aguardando a indexação
initializeAsync() {
if (!this.indexingPromise) {
this.indexingPromise = this.searchEngine.initialize()
.catch(err => {
console.error('[ERROR] Indexação async falhou:', err?.message || err);
});
}
}
async handleRequest(request) {
const { method, params } = request;
switch (method) {
case 'initialize': {
// start indexação em background
this.initializeAsync();
return {
protocolVersion: '2024-11-05',
capabilities: {
tools: {}, // habilita tools/list e tools/call
prompts: {}, // (opcional)
resources: {} // (opcional)
},
serverInfo: {
name: 'mcp-codebase-search',
version: '1.1.0'
}
};
}
case 'tools/list': {
return {
tools: [
{
name: 'search_codebase',
description: 'Busca textual simples em arquivos do projeto',
inputSchema: {
type: 'object',
properties: {
query: { type: 'string', description: 'Termo(s) de busca' },
limit: { type: 'number', description: 'Máx. de resultados', default: 10 },
file_types: {
type: 'array',
items: { type: 'string' },
description: 'Filtro por extensões/caminhos (ex: .ts, .jsx, src/)'
}
},
required: ['query']
}
},
{
name: 'echo',
description: 'Retorna os argumentos recebidos (health-check/debug)',
inputSchema: {
type: 'object',
properties: {
message: { type: 'string', description: 'Mensagem para ecoar' }
},
required: ['message']
}
}
]
};
}
case 'tools/call': {
const { name, arguments: args = {} } = params || {};
if (name === 'search_codebase') {
const payload = this.searchEngine.search(args.query, {
limit: args.limit || 10,
file_types: args.file_types
});
return {
content: [
{
type: 'text',
text: JSON.stringify(payload, null, 2)
}
]
};
}
if (name === 'echo') {
return {
content: [
{
type: 'text',
text: JSON.stringify({ ok: true, received: args }, null, 2)
}
]
};
}
throw new Error(`Tool desconhecida: ${name}`);
}
default:
throw new Error(`Método não suportado: ${method}`);
}
}
start() {
console.error('[LOG] Servidor MCP iniciado');
console.error('[LOG] Configurando stdin/stdout...');
process.stdin.setEncoding(null); // binário
process.stdin.resume();
console.error('[LOG] Aguardando mensagens…');
let buf = Buffer.alloc(0);
process.stdin.on('data', (chunk) => {
// console.error('[DBG] chunk bytes:', chunk.length);
const chunkBuffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
buf = Buffer.concat([buf, chunkBuffer]);
try {
buf = parseMessages(buf, (request) => {
// console.error('[DBG] <<<', JSON.stringify(request));
// Processar request de forma não-bloqueante
this.handleRequest(request)
.then(result => {
if (Object.prototype.hasOwnProperty.call(request, 'id')) {
sendMessage({ jsonrpc: '2.0', id: request.id, result });
}
})
.catch(err => {
console.error('[ERROR] handleRequest:', err?.message || err);
if (Object.prototype.hasOwnProperty.call(request, 'id')) {
sendMessage({
jsonrpc: '2.0',
id: request.id,
error: { code: -32603, message: String(err?.message || err) }
});
}
});
});
} catch (err) {
console.error('[ERROR] parseMessages:', err?.message || err);
}
});
process.stdin.on('end', () => {
console.error('[LOG] stdin finalizado');
process.exit(0);
});
process.stdin.on('error', (err) => {
console.error('[ERROR] stdin error:', err?.message || err);
process.exit(1);
});
process.on('SIGINT', () => {
console.error('[LOG] SIGINT recebido, encerrando…');
process.exit(0);
});
process.on('SIGTERM', () => {
console.error('[LOG] SIGTERM recebido, encerrando…');
process.exit(0);
});
process.on('uncaughtException', (err) => {
console.error('[ERROR] uncaughtException:', err?.message || err);
process.exit(1);
});
process.on('unhandledRejection', (reason) => {
console.error('[ERROR] unhandledRejection:', reason);
process.exit(1);
});
}
}
// ----------------------------------------------------
if (require.main === module) {
new MCPServer().start();
}
module.exports = { MCPServer, SimpleCodebaseSearch };