UNPKG

@moikas/code-audit-mcp

Version:

AI-powered code auditing via MCP using local Ollama models for security, performance, and quality analysis

746 lines 27.7 kB
/** * Main MCP server implementation for code auditing */ import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { CallToolRequestSchema, ListToolsRequestSchema, ErrorCode, McpError, } from '@modelcontextprotocol/sdk/types.js'; import { OllamaClient } from './ollama/client.js'; import { ModelManager, DefaultModelSelectionStrategy, } from './ollama/models.js'; import { AuditorFactory } from './auditors/index.js'; import { logger } from './utils/mcp-logger.js'; /** * Default server configuration */ const DEFAULT_CONFIG = { name: 'code-audit-mcp', version: '1.0.0', ollama: { host: 'http://localhost:11434', timeout: 30000, retryAttempts: 3, retryDelay: 1000, healthCheckInterval: 60000, }, models: {}, auditors: { security: { enabled: true, severity: ['critical', 'high', 'medium'], rules: {}, thresholds: {}, }, completeness: { enabled: true, severity: ['critical', 'high', 'medium'], rules: {}, thresholds: {}, }, performance: { enabled: true, severity: ['critical', 'high', 'medium'], rules: {}, thresholds: {}, }, quality: { enabled: true, severity: ['high', 'medium', 'low'], rules: {}, thresholds: {}, }, architecture: { enabled: true, severity: ['high', 'medium', 'low'], rules: {}, thresholds: {}, }, testing: { enabled: true, severity: ['high', 'medium', 'low'], rules: {}, thresholds: {}, }, documentation: { enabled: true, severity: ['medium', 'low', 'info'], rules: {}, thresholds: {}, }, all: { enabled: true, severity: ['critical', 'high', 'medium', 'low'], rules: {}, thresholds: {}, }, }, languages: {}, logging: { level: 'info', enableMetrics: true, enableTracing: false, }, performance: { maxConcurrentAudits: 3, cacheEnabled: false, cacheTtl: 300, }, }; /** * Main MCP server class */ export class CodeAuditServer { server; config; ollamaClient; modelManager; auditors = {}; activeAudits = new Map(); constructor(config = {}) { this.config = { ...DEFAULT_CONFIG, ...config }; this.server = new Server({ name: this.config.name, version: this.config.version, }, { capabilities: { tools: {}, }, }); this.ollamaClient = new OllamaClient(this.config.ollama); this.modelManager = new ModelManager(new DefaultModelSelectionStrategy()); this.setupToolHandlers(); this.setupErrorHandling(); } /** * Initialize the server */ async initialize() { try { logger.log('Initializing Code Audit MCP Server...'); // Create auditors first (doesn't require Ollama connection) this.auditors = AuditorFactory.createAllAuditors(this.config.auditors, this.ollamaClient, this.modelManager); logger.log(`Server initialized with ${Object.keys(this.auditors).length} auditors`); // Initialize Ollama client in the background (non-blocking) this.initializeOllamaAsync(); } catch (error) { logger.error('Failed to initialize server:', error); throw error; } } /** * Initialize Ollama client asynchronously */ async initializeOllamaAsync() { try { await this.ollamaClient.initialize(); logger.log('Ollama client connected successfully'); // Perform health check after Ollama is ready const health = await this.healthCheck(); if (health.status === 'unhealthy') { logger.warn('Ollama health check failed - some features may be limited'); } } catch (error) { logger.warn('Failed to connect to Ollama:', error instanceof Error ? error.message : error); logger.log('Server will continue running - Ollama connection will be retried on demand'); // Ensure the error doesn't propagate as an unhandled rejection return; } } /** * Setup tool handlers */ setupToolHandlers() { this.server.setRequestHandler(ListToolsRequestSchema, async () => { return { tools: [ { name: 'audit_code', description: 'Perform comprehensive code audit using AI models', inputSchema: { type: 'object', properties: { code: { type: 'string', description: 'The code to audit', }, language: { type: 'string', description: 'Programming language of the code', }, auditType: { type: 'string', enum: [ 'security', 'completeness', 'performance', 'quality', 'architecture', 'testing', 'documentation', 'all', ], description: 'Type of audit to perform', default: 'all', }, file: { type: 'string', description: 'Optional file path for context', }, context: { type: 'object', description: 'Additional context for the audit', properties: { framework: { type: 'string' }, environment: { type: 'string', enum: ['production', 'development', 'testing'], }, performanceCritical: { type: 'boolean' }, teamSize: { type: 'number' }, projectType: { type: 'string', enum: [ 'web', 'api', 'cli', 'library', 'mobile', 'desktop', ], }, }, }, priority: { type: 'string', enum: ['fast', 'thorough'], description: 'Audit priority (fast = security + completeness only)', default: 'thorough', }, maxIssues: { type: 'number', description: 'Maximum number of issues to return', default: 50, }, includeFixSuggestions: { type: 'boolean', description: 'Include fix suggestions in the response', default: true, }, }, required: ['code', 'language'], }, }, { name: 'health_check', description: 'Check the health status of the audit server', inputSchema: { type: 'object', properties: {}, }, }, { name: 'list_models', description: 'List available AI models for auditing', inputSchema: { type: 'object', properties: {}, }, }, { name: 'update_config', description: 'Update server configuration', inputSchema: { type: 'object', properties: { auditors: { type: 'object', description: 'Auditor configurations to update', }, ollama: { type: 'object', description: 'Ollama client configuration', }, }, }, }, ], }; }); this.server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; try { switch (name) { case 'audit_code': return await this.handleAuditCode(args); case 'health_check': return await this.handleHealthCheck(); case 'list_models': return await this.handleListModels(); case 'update_config': return await this.handleUpdateConfig(args); default: throw new McpError(ErrorCode.MethodNotFound, `Unknown tool: ${name}`); } } catch (error) { logger.error(`Tool ${name} failed:`, error); if (error) logger.error('Error stack:', error.stack); if (error instanceof McpError) { throw error; } const errorMessage = error instanceof Error ? error.message : typeof error === 'string' ? error : 'Truly unknown error - check server logs'; throw new McpError(ErrorCode.InternalError, `Tool execution failed: ${errorMessage}`); } }); } /** * Handle audit_code tool requests */ async handleAuditCode(request) { // Validate request this.validateAuditRequest(request); // Check for duplicate audits const requestKey = this.generateRequestKey(request); const existingAudit = this.activeAudits.get(requestKey); if (existingAudit) { logger.log('Audit already in progress, waiting for completion...'); const result = await existingAudit.promise; return { content: [ { type: 'text', text: JSON.stringify(result, null, 2), }, ], }; } // Create audit promise with timeout const auditPromise = this.performAudit(request); const auditTimeout = 300000; // 5 minutes // Create timeout that will clean up the audit const timeout = setTimeout(() => { const audit = this.activeAudits.get(requestKey); if (audit) { logger.error(`Audit ${requestKey} timed out after ${auditTimeout / 1000} seconds`); this.activeAudits.delete(requestKey); } }, auditTimeout); // Store audit metadata this.activeAudits.set(requestKey, { promise: auditPromise, startTime: Date.now(), timeout, }); try { const result = await auditPromise; return { content: [ { type: 'text', text: JSON.stringify(result, null, 2), }, ], }; } finally { // Clean up audit and cancel timeout const audit = this.activeAudits.get(requestKey); if (audit) { clearTimeout(audit.timeout); this.activeAudits.delete(requestKey); } } } /** * Perform the actual audit */ async performAudit(request) { const startTime = Date.now(); try { // Handle 'all' audit type if (request.auditType === 'all') { return await this.performMultipleAudits(request); } // Handle fast mode (security + completeness only) if (request.priority === 'fast') { return await this.performFastModeAudit(request); } // Handle single audit type const auditor = this.auditors[request.auditType]; if (!auditor) { throw new Error(`Auditor not available for type: ${request.auditType}`); } const result = await auditor.audit(request); // Log metrics this.logAuditMetrics(request, result, Date.now() - startTime); return result; } catch (error) { logger.error('Audit failed:', error); throw error; } } /** * Perform multiple audits for 'all' type */ async performMultipleAudits(request) { const auditTypes = [ 'security', 'completeness', 'performance', 'quality', 'architecture', 'testing', 'documentation', ]; const results = []; // Run audits in parallel (up to maxConcurrentAudits) const concurrentLimit = this.config.performance.maxConcurrentAudits; const chunks = this.chunkArray(auditTypes, concurrentLimit); for (const chunk of chunks) { const chunkPromises = chunk.map(async (auditType) => { const auditor = this.auditors[auditType]; if (auditor) { const auditRequest = { ...request, auditType }; return await auditor.audit(auditRequest); } return null; }); const chunkResults = await Promise.all(chunkPromises); results.push(...chunkResults.filter((result) => result !== null)); } // Merge results return this.mergeAuditResults(results, request); } /** * Perform fast mode audit (security + completeness only) */ async performFastModeAudit(request) { const fastAuditTypes = ['security', 'completeness']; const results = []; for (const auditType of fastAuditTypes) { const auditor = this.auditors[auditType]; if (auditor) { const auditRequest = { ...request, auditType, priority: 'fast', }; const result = await auditor.audit(auditRequest); results.push(result); } } return this.mergeAuditResults(results, request); } /** * Handle health_check tool requests */ async handleHealthCheck() { const health = await this.healthCheck(); return { content: [ { type: 'text', text: JSON.stringify(health, null, 2), }, ], }; } /** * Handle list_models tool requests */ async handleListModels() { const availableModels = this.ollamaClient.getAvailableModels(); const allModels = this.modelManager.getAllModels(); const modelMetrics = this.ollamaClient.getAllMetrics(); const modelInfo = allModels.map((config) => ({ ...config, available: availableModels.includes(config.name), metrics: modelMetrics[config.name] || null, })); return { content: [ { type: 'text', text: JSON.stringify({ models: modelInfo }, null, 2), }, ], }; } /** * Handle update_config tool requests */ async handleUpdateConfig(args) { try { if (args.auditors) { for (const [auditType, config] of Object.entries(args.auditors)) { if (this.auditors[auditType]) { this.auditors[auditType].updateConfig(config); } } } if (args.ollama) { // Note: Ollama config updates would require client recreation logger.log('Ollama config update requested (requires restart)'); } return { content: [ { type: 'text', text: JSON.stringify({ message: 'Configuration updated successfully' }, null, 2), }, ], }; } catch (error) { throw new McpError(ErrorCode.InternalError, `Failed to update configuration: ${error instanceof Error ? error.message : 'Unknown error'}`); } } /** * Validate audit request */ validateAuditRequest(request) { if (!request.code?.trim()) { throw new McpError(ErrorCode.InvalidParams, 'Code is required and cannot be empty'); } if (!request.language?.trim()) { throw new McpError(ErrorCode.InvalidParams, 'Language is required'); } if (request.code.length > 100000) { throw new McpError(ErrorCode.InvalidParams, 'Code size exceeds limit (100KB)'); } const validAuditTypes = [ 'security', 'completeness', 'performance', 'quality', 'architecture', 'testing', 'documentation', 'all', ]; if (!validAuditTypes.includes(request.auditType)) { throw new McpError(ErrorCode.InvalidParams, `Invalid audit type: ${request.auditType}`); } } /** * Generate request key for deduplication */ generateRequestKey(request) { const hash = this.simpleHash(request.code); return `${request.language}_${request.auditType}_${request.priority || 'thorough'}_${hash}`; } /** * Simple hash function for code deduplication */ simpleHash(str) { let hash = 0; for (let i = 0; i < str.length; i++) { const char = str.charCodeAt(i); hash = (hash << 5) - hash + char; hash = hash & hash; // Convert to 32-bit integer } return Math.abs(hash).toString(36); } /** * Chunk array into smaller arrays */ chunkArray(array, chunkSize) { const chunks = []; for (let i = 0; i < array.length; i += chunkSize) { chunks.push(array.slice(i, i + chunkSize)); } return chunks; } /** * Merge multiple audit results */ mergeAuditResults(results, request) { if (results.length === 0) { throw new Error('No audit results to merge'); } if (results.length === 1) { return results[0]; } const merged = { requestId: results[0].requestId, issues: [], summary: { total: 0, critical: 0, high: 0, medium: 0, low: 0, info: 0, byCategory: {}, byType: {}, }, coverage: { linesAnalyzed: 0, functionsAnalyzed: 0, classesAnalyzed: 0, complexity: 0, }, suggestions: { autoFixable: [], priorityFixes: [], quickWins: [], technicalDebt: [], }, metrics: { duration: 0, modelResponseTime: 0, parsingTime: 0, postProcessingTime: 0, }, model: results.map((r) => r.model).join(', '), timestamp: new Date().toISOString(), version: results[0].version, }; // Merge all issues for (const result of results) { merged.issues.push(...result.issues); // Update summary merged.summary.total += result.summary.total; merged.summary.critical += result.summary.critical; merged.summary.high += result.summary.high; merged.summary.medium += result.summary.medium; merged.summary.low += result.summary.low; merged.summary.info += result.summary.info; // Merge categories for (const [category, count] of Object.entries(result.summary.byCategory)) { merged.summary.byCategory[category] = (merged.summary.byCategory[category] || 0) + count; } // Merge types for (const [type, count] of Object.entries(result.summary.byType)) { merged.summary.byType[type] = (merged.summary.byType[type] || 0) + count; } // Merge suggestions merged.suggestions.autoFixable.push(...result.suggestions.autoFixable); merged.suggestions.priorityFixes.push(...result.suggestions.priorityFixes); merged.suggestions.quickWins.push(...result.suggestions.quickWins); merged.suggestions.technicalDebt.push(...result.suggestions.technicalDebt); // Update metrics merged.metrics.duration += result.metrics.duration; merged.metrics.modelResponseTime += result.metrics.modelResponseTime; merged.metrics.parsingTime += result.metrics.parsingTime; merged.metrics.postProcessingTime += result.metrics.postProcessingTime; // Use max coverage values merged.coverage.linesAnalyzed = Math.max(merged.coverage.linesAnalyzed, result.coverage.linesAnalyzed); merged.coverage.functionsAnalyzed = Math.max(merged.coverage.functionsAnalyzed, result.coverage.functionsAnalyzed); merged.coverage.classesAnalyzed = Math.max(merged.coverage.classesAnalyzed, result.coverage.classesAnalyzed); merged.coverage.complexity = Math.max(merged.coverage.complexity, result.coverage.complexity); } // Sort merged issues by severity and line number merged.issues.sort((a, b) => { const severityOrder = { critical: 0, high: 1, medium: 2, low: 3, info: 4, }; const severityDiff = severityOrder[a.severity] - severityOrder[b.severity]; if (severityDiff !== 0) return severityDiff; return a.location.line - b.location.line; }); // Apply max issues limit if (request.maxIssues && merged.issues.length > request.maxIssues) { merged.issues = merged.issues.slice(0, request.maxIssues); } return merged; } /** * Perform health check */ async healthCheck() { const health = await this.ollamaClient.healthCheck(); // Add auditor health checks for (const [auditType, auditor] of Object.entries(this.auditors)) { health.checks.auditors[auditType] = auditor !== undefined; } return health; } /** * Log audit metrics */ logAuditMetrics(request, result, totalTime) { if (!this.config.logging.enableMetrics) { return; } logger.log(`Audit completed: ${request.auditType} (${request.language}) - ${result.issues.length} issues found in ${totalTime}ms`); } /** * Setup error handling */ setupErrorHandling() { this.server.onerror = (error) => { logger.error('MCP Server error:', error); }; process.on('SIGINT', async () => { logger.log('Shutting down server...'); await this.cleanup(); process.exit(0); }); process.on('SIGTERM', async () => { logger.log('Shutting down server...'); await this.cleanup(); process.exit(0); }); } /** * Start the server */ async start() { await this.initialize(); const transport = new StdioServerTransport(); await this.server.connect(transport); logger.log('Code Audit MCP Server is running'); } /** * Cleanup resources */ async cleanup() { try { await this.ollamaClient.cleanup(); logger.log('Server cleanup completed'); } catch (error) { logger.error('Error during cleanup:', error); } } } /** * Start server if this file is run directly */ if (import.meta.url === `file://${process.argv[1]}`) { // Add unhandled rejection handlers process.on('unhandledRejection', (reason, promise) => { logger.error('Unhandled Rejection at:', promise, 'reason:', reason); // Don't exit in stdio mode to allow proper error reporting if (process.env.MCP_STDIO_MODE !== 'true') { process.exit(1); } }); process.on('uncaughtException', (error) => { logger.error('Uncaught Exception:', error); // Don't exit in stdio mode to allow proper error reporting if (process.env.MCP_STDIO_MODE !== 'true') { process.exit(1); } }); const server = new CodeAuditServer(); server.start().catch((error) => { logger.error('Failed to start server:', error); process.exit(1); }); } //# sourceMappingURL=index.js.map