vibe-coder-mcp
Version:
Production-ready MCP server with complete agent integration, multi-transport support, and comprehensive development automation tools for AI-assisted workflows.
443 lines (442 loc) • 18.8 kB
JavaScript
import fs from 'fs/promises';
import path from 'path';
import os from 'os';
import ParserFromPackage from 'web-tree-sitter';
import logger from '../../../logger.js';
import { resolveProjectPath } from '../utils/pathUtils.enhanced.js';
const GRAMMARS_BASE_DIR = resolveProjectPath('src/tools/code-map-generator/grammars');
export class GrammarManager {
initialized = false;
parser = null;
grammars = new Map();
grammarConfigs = {};
lruList = [];
options;
grammarsBaseDir;
grammarSizes = new Map();
lastUsedTimestamps = new Map();
totalMemoryUsage = 0;
static DEFAULT_OPTIONS = {
maxGrammars: 20,
preloadCommonGrammars: false,
preloadExtensions: ['.js', '.ts', '.py', '.html', '.css'],
grammarsBaseDir: GRAMMARS_BASE_DIR,
maxMemoryUsage: 100 * 1024 * 1024,
grammarIdleTimeout: 5 * 60 * 1000,
enableIncrementalParsing: true,
incrementalParsingThreshold: 1 * 1024 * 1024
};
constructor(grammarConfigs, options = {}) {
this.grammarConfigs = grammarConfigs;
this.options = {
...GrammarManager.DEFAULT_OPTIONS,
...options
};
this.grammarsBaseDir = this.options.grammarsBaseDir;
logger.info(`Grammar files directory: ${this.grammarsBaseDir}`);
logger.debug(`GrammarManager created with max grammars: ${this.options.maxGrammars}`);
}
async initialize() {
if (this.initialized) {
return;
}
try {
await ParserFromPackage.init();
this.parser = new ParserFromPackage();
this.initialized = true;
if (this.options.preloadCommonGrammars) {
await this.preloadGrammars();
}
logger.info('GrammarManager initialized successfully.');
}
catch (error) {
logger.error({ err: error }, 'Failed to initialize GrammarManager.');
throw error;
}
}
async preloadGrammars() {
logger.debug(`Preloading grammars for extensions: ${this.options.preloadExtensions.join(', ')}`);
const preloadPromises = this.options.preloadExtensions.map(ext => this.loadGrammar(ext).catch(error => {
logger.warn({ err: error, extension: ext }, `Failed to preload grammar for ${ext}`);
}));
await Promise.all(preloadPromises);
}
async loadGrammar(extension) {
this.ensureInitialized();
if (this.grammars.has(extension)) {
this.updateLRU(extension);
this.lastUsedTimestamps.set(extension, Date.now());
return this.grammars.get(extension);
}
const langConfig = this.grammarConfigs[extension];
if (!langConfig) {
throw new Error(`No language configuration found for extension: ${extension}`);
}
await this.checkMemoryUsage();
try {
const wasmPath = path.join(this.grammarsBaseDir, langConfig.wasmPath);
try {
await fs.access(wasmPath, fs.constants.F_OK);
logger.debug(`Grammar file found: ${wasmPath}`);
}
catch (accessError) {
logger.error({
err: accessError,
grammarName: langConfig.name,
wasmPath: wasmPath,
cwd: process.cwd()
}, `File not found: Tree-sitter grammar for ${langConfig.name}. Ensure '${langConfig.wasmPath}' exists in '${this.grammarsBaseDir}'.`);
throw new Error(`Grammar file not found: ${wasmPath}`);
}
const stats = await fs.stat(wasmPath);
const fileSize = stats.size;
const estimatedMemoryUsage = fileSize * 4;
if (this.totalMemoryUsage + estimatedMemoryUsage > this.options.maxMemoryUsage) {
const freedMemory = await this.unloadLeastRecentlyUsedGrammars(estimatedMemoryUsage);
if (this.totalMemoryUsage + estimatedMemoryUsage - freedMemory > this.options.maxMemoryUsage) {
logger.warn({
grammarName: langConfig.name,
extension,
estimatedMemoryUsage: this.formatBytes(estimatedMemoryUsage),
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage),
maxMemoryUsage: this.formatBytes(this.options.maxMemoryUsage)
}, `Loading grammar for ${langConfig.name} may exceed memory limits`);
}
}
const language = await ParserFromPackage.Language.load(wasmPath);
if (this.grammars.size >= this.options.maxGrammars) {
this.evictLRU();
}
this.grammars.set(extension, language);
this.updateLRU(extension);
this.lastUsedTimestamps.set(extension, Date.now());
this.grammarSizes.set(extension, estimatedMemoryUsage);
this.totalMemoryUsage += estimatedMemoryUsage;
logger.info({
grammarName: langConfig.name,
extension,
memoryUsage: this.formatBytes(estimatedMemoryUsage),
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage)
}, `Successfully loaded Tree-sitter grammar for ${langConfig.name} (${extension})`);
return language;
}
catch (error) {
logger.error({
err: error,
grammarName: langConfig.name,
extension
}, `Failed to load Tree-sitter grammar for ${langConfig.name}`);
throw error;
}
}
updateLRU(extension) {
const index = this.lruList.indexOf(extension);
if (index !== -1) {
this.lruList.splice(index, 1);
}
this.lruList.unshift(extension);
}
evictLRU() {
if (this.lruList.length === 0) {
return;
}
const lruExtension = this.lruList.pop();
if (lruExtension) {
const memoryUsage = this.grammarSizes.get(lruExtension) || 0;
this.grammars.delete(lruExtension);
this.grammarSizes.delete(lruExtension);
this.lastUsedTimestamps.delete(lruExtension);
this.totalMemoryUsage = Math.max(0, this.totalMemoryUsage - memoryUsage);
logger.debug({
extension: lruExtension,
freedMemory: this.formatBytes(memoryUsage),
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage)
}, `Evicted grammar for extension ${lruExtension} due to LRU policy`);
}
}
async checkMemoryUsage() {
if (this.totalMemoryUsage > this.options.maxMemoryUsage * 0.8) {
logger.info({
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage),
maxMemoryUsage: this.formatBytes(this.options.maxMemoryUsage)
}, 'Grammar memory usage is high, unloading unused grammars');
await this.unloadUnusedGrammars();
}
const now = Date.now();
const idleExtensions = [];
for (const [extension, lastUsed] of this.lastUsedTimestamps.entries()) {
if (now - lastUsed > this.options.grammarIdleTimeout) {
idleExtensions.push(extension);
}
}
if (idleExtensions.length > 0) {
for (const extension of idleExtensions) {
this.unloadGrammar(extension);
}
logger.debug({
count: idleExtensions.length,
extensions: idleExtensions
}, `Unloaded ${idleExtensions.length} idle grammars`);
}
}
async unloadUnusedGrammars() {
const sortedExtensions = Array.from(this.lastUsedTimestamps.entries())
.sort((a, b) => a[1] - b[1])
.map(([extension]) => extension);
const extensionsToUnload = sortedExtensions.slice(0, -5);
for (const extension of extensionsToUnload) {
this.unloadGrammar(extension);
}
logger.info({
count: extensionsToUnload.length,
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage)
}, `Unloaded ${extensionsToUnload.length} unused grammars`);
}
async unloadLeastRecentlyUsedGrammars(requiredMemory) {
const sortedExtensions = Array.from(this.lastUsedTimestamps.entries())
.sort((a, b) => a[1] - b[1])
.map(([extension]) => extension);
const candidateExtensions = sortedExtensions.slice(0, -3);
let freedMemory = 0;
for (const extension of candidateExtensions) {
const memoryUsage = this.grammarSizes.get(extension) || 0;
this.unloadGrammar(extension);
freedMemory += memoryUsage;
if (freedMemory >= requiredMemory) {
break;
}
}
logger.debug({
requiredMemory: this.formatBytes(requiredMemory),
freedMemory: this.formatBytes(freedMemory)
}, `Freed ${this.formatBytes(freedMemory)} of memory by unloading grammars`);
return freedMemory;
}
unloadGrammar(extension) {
if (!this.grammars.has(extension)) {
return;
}
const memoryUsage = this.grammarSizes.get(extension) || 0;
this.grammars.delete(extension);
this.grammarSizes.delete(extension);
this.lastUsedTimestamps.delete(extension);
const index = this.lruList.indexOf(extension);
if (index !== -1) {
this.lruList.splice(index, 1);
}
this.totalMemoryUsage = Math.max(0, this.totalMemoryUsage - memoryUsage);
logger.debug({
extension,
freedMemory: this.formatBytes(memoryUsage),
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage)
}, `Unloaded grammar for extension ${extension}`);
}
formatBytes(bytes) {
if (bytes === 0)
return '0 Bytes';
const k = 1024;
const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB'];
const i = Math.floor(Math.log(bytes) / Math.log(k));
return parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + ' ' + sizes[i];
}
async getParserForExtension(extension) {
this.ensureInitialized();
const language = await this.loadGrammar(extension);
this.parser.setLanguage(language);
return this.parser;
}
async getParserForExtensionWithMemoryAwareness(extension) {
this.ensureInitialized();
const language = await this.loadGrammarWithMemoryAwareness(extension);
this.parser.setLanguage(language);
return this.parser;
}
ensureInitialized() {
if (!this.initialized) {
throw new Error('GrammarManager is not initialized. Call initialize() first.');
}
}
async getMemoryStats() {
const memoryUsage = process.memoryUsage();
const systemTotal = os.totalmem();
const memoryUsagePercentage = memoryUsage.rss / systemTotal;
return {
heapUsed: memoryUsage.heapUsed,
heapTotal: memoryUsage.heapTotal,
rss: memoryUsage.rss,
systemTotal,
memoryUsagePercentage,
formatted: {
heapUsed: this.formatBytes(memoryUsage.heapUsed),
heapTotal: this.formatBytes(memoryUsage.heapTotal),
rss: this.formatBytes(memoryUsage.rss),
systemTotal: this.formatBytes(systemTotal)
}
};
}
estimateGrammarSize(langConfig) {
const baseSize = 3 * 1024 * 1024;
const grammarSizeMultipliers = {
'JavaScript': 1.2,
'TypeScript': 1.5,
'Python': 1.0,
'HTML': 0.8,
'CSS': 0.7,
'C++': 1.8,
'Java': 1.3,
'Ruby': 1.1,
'Go': 0.9,
'Rust': 1.4,
};
const multiplier = grammarSizeMultipliers[langConfig.name] || 1.0;
return baseSize * multiplier;
}
async loadGrammarWithMemoryAwareness(extension) {
this.ensureInitialized();
if (this.grammars.has(extension)) {
this.updateLRU(extension);
this.lastUsedTimestamps.set(extension, Date.now());
return this.grammars.get(extension);
}
const langConfig = this.grammarConfigs[extension];
if (!langConfig) {
throw new Error(`No language configuration found for extension: ${extension}`);
}
const memoryStats = await this.getMemoryStats();
const estimatedGrammarSize = this.estimateGrammarSize(langConfig);
logger.debug({
extension,
grammarName: langConfig.name,
estimatedSize: this.formatBytes(estimatedGrammarSize),
currentMemoryUsage: this.formatBytes(this.totalMemoryUsage),
maxMemoryUsage: this.formatBytes(this.options.maxMemoryUsage)
}, `Preparing to load grammar for ${langConfig.name}`);
if (memoryStats.memoryUsagePercentage > 0.7) {
logger.info(`Memory usage high (${memoryStats.memoryUsagePercentage.toFixed(2)}%), performing aggressive cleanup before loading new grammar`);
const requiredMemory = estimatedGrammarSize * 1.5;
await this.unloadLeastRecentlyUsedGrammars(requiredMemory);
if (global.gc) {
global.gc();
}
}
else if (this.grammars.size >= this.options.maxGrammars) {
this.evictLRU();
}
try {
const wasmPath = path.join(this.grammarsBaseDir, langConfig.wasmPath);
try {
await fs.access(wasmPath, fs.constants.F_OK);
logger.debug(`Grammar file found: ${wasmPath}`);
}
catch {
throw new Error(`Grammar file not found: ${wasmPath}`);
}
const startTime = performance.now();
const language = await ParserFromPackage.Language.load(wasmPath);
const loadTime = performance.now() - startTime;
this.grammars.set(extension, language);
this.updateLRU(extension);
this.lastUsedTimestamps.set(extension, Date.now());
const estimatedMemoryUsage = this.estimateGrammarSize(langConfig);
this.grammarSizes.set(extension, estimatedMemoryUsage);
this.totalMemoryUsage += estimatedMemoryUsage;
logger.info({
extension,
grammarName: langConfig.name,
loadTimeMs: loadTime.toFixed(2),
estimatedSize: this.formatBytes(estimatedMemoryUsage),
totalMemoryUsage: this.formatBytes(this.totalMemoryUsage),
totalGrammars: this.grammars.size
}, `Successfully loaded grammar for ${langConfig.name}`);
return language;
}
catch (error) {
logger.error({
err: error,
grammarName: langConfig.name,
extension
}, `Failed to load Tree-sitter grammar for ${langConfig.name}`);
throw error;
}
}
isInitialized() {
return this.initialized;
}
getLoadedGrammars() {
return Array.from(this.grammars.keys());
}
getStats() {
const grammarStats = [];
for (const [extension] of this.grammars.entries()) {
const size = this.grammarSizes.get(extension) || 0;
const lastUsed = this.lastUsedTimestamps.get(extension) || 0;
const lruIndex = this.lruList.indexOf(extension);
grammarStats.push({
extension,
size,
sizeFormatted: this.formatBytes(size),
lastUsed: new Date(lastUsed).toISOString(),
idleTime: Date.now() - lastUsed,
lruIndex: lruIndex === -1 ? 'not in LRU' : lruIndex
});
}
grammarStats.sort((a, b) => b.size - a.size);
return {
loadedGrammars: this.grammars.size,
maxGrammars: this.options.maxGrammars,
totalMemoryUsage: this.totalMemoryUsage,
totalMemoryUsageFormatted: this.formatBytes(this.totalMemoryUsage),
maxMemoryUsage: this.options.maxMemoryUsage,
maxMemoryUsageFormatted: this.formatBytes(this.options.maxMemoryUsage),
memoryUsagePercentage: (this.totalMemoryUsage / this.options.maxMemoryUsage) * 100,
lruList: [...this.lruList],
initialized: this.initialized,
grammars: grammarStats
};
}
getOptions() {
return { ...this.options };
}
async prepareGrammarsForBatch(fileExtensions) {
const extensionCounts = new Map();
for (const ext of fileExtensions) {
extensionCounts.set(ext, (extensionCounts.get(ext) || 0) + 1);
}
const sortedExtensions = Array.from(extensionCounts.entries())
.sort((a, b) => b[1] - a[1])
.map(([ext]) => ext);
const availableMemory = this.options.maxMemoryUsage - this.totalMemoryUsage;
const estimatedSizes = new Map();
let totalEstimatedSize = 0;
const extensionsToLoad = [];
for (const ext of sortedExtensions) {
if (this.grammars.has(ext))
continue;
const langConfig = this.grammarConfigs[ext];
if (!langConfig)
continue;
const estimatedSize = this.estimateGrammarSize(langConfig);
estimatedSizes.set(ext, estimatedSize);
if (totalEstimatedSize + estimatedSize <= availableMemory) {
extensionsToLoad.push(ext);
totalEstimatedSize += estimatedSize;
}
}
if (extensionsToLoad.length < sortedExtensions.length) {
const additionalMemoryNeeded = totalEstimatedSize - availableMemory;
if (additionalMemoryNeeded > 0) {
await this.unloadLeastRecentlyUsedGrammars(additionalMemoryNeeded);
}
}
const loadPromises = extensionsToLoad.map(ext => this.loadGrammarWithMemoryAwareness(ext).catch(error => {
logger.warn({ err: error, extension: ext }, `Failed to preload grammar for batch`);
return null;
}));
await Promise.all(loadPromises);
logger.info({
batchSize: fileExtensions.length,
uniqueExtensions: sortedExtensions.length,
loadedExtensions: extensionsToLoad.length
}, `Prepared grammars for batch processing`);
}
}