UNPKG

code-context-mcp

Version:

MCP server for semantic code search powered by MongoDB Atlas Vector Search and Voyage AI embeddings

644 lines โ€ข 34.7 kB
import * as fs from "fs"; import * as path from "path"; import { COLLECTION_LIMIT_MESSAGE } from "@zilliz/claude-context-core"; import { ensureAbsolutePath, truncateContent, trackCodebasePath } from "./utils.js"; export class ToolHandlers { context; snapshotManager; indexingStats = null; currentWorkspace; constructor(context, snapshotManager) { this.context = context; this.snapshotManager = snapshotManager; this.currentWorkspace = process.cwd(); console.log(`[WORKSPACE] Current workspace: ${this.currentWorkspace}`); } /** * Sync indexed codebases from Zilliz Cloud collections * This method fetches all collections from the vector database, * gets the first document from each collection to extract codebasePath from metadata, * and updates the snapshot with discovered codebases. * * Logic: Compare mcp-codebase-snapshot.json with zilliz cloud collections * - If local snapshot has extra directories (not in cloud), remove them * - If local snapshot is missing directories (exist in cloud), ignore them */ async syncIndexedCodebasesFromCloud() { try { console.log(`[SYNC-CLOUD] ๐Ÿ”„ Syncing indexed codebases from Zilliz Cloud...`); // Get all collections using the interface method const vectorDb = this.context.getVectorDatabase(); // Use the new listCollections method from the interface const collections = await vectorDb.listCollections(); console.log(`[SYNC-CLOUD] ๐Ÿ“‹ Found ${collections.length} collections in Zilliz Cloud`); if (collections.length === 0) { console.log(`[SYNC-CLOUD] โœ… No collections found in cloud`); // If no collections in cloud, remove all local codebases const localCodebases = this.snapshotManager.getIndexedCodebases(); if (localCodebases.length > 0) { console.log(`[SYNC-CLOUD] ๐Ÿงน Removing ${localCodebases.length} local codebases as cloud has no collections`); for (const codebasePath of localCodebases) { this.snapshotManager.removeIndexedCodebase(codebasePath); console.log(`[SYNC-CLOUD] โž– Removed local codebase: ${codebasePath}`); } this.snapshotManager.saveCodebaseSnapshot(); console.log(`[SYNC-CLOUD] ๐Ÿ’พ Updated snapshot to match empty cloud state`); } return; } const cloudCodebases = new Set(); // Check each collection for codebase path for (const collectionName of collections) { try { // Skip collections that don't match the code_chunks pattern (support both legacy and new collections) if (!collectionName.startsWith('code_chunks_') && !collectionName.startsWith('hybrid_code_chunks_')) { console.log(`[SYNC-CLOUD] โญ๏ธ Skipping non-code collection: ${collectionName}`); continue; } console.log(`[SYNC-CLOUD] ๐Ÿ” Checking collection: ${collectionName}`); // Query the first document to get metadata const results = await vectorDb.query(collectionName, '', // Empty filter to get all results ['metadata'], // Only fetch metadata field 1 // Only need one result to extract codebasePath ); if (results && results.length > 0) { const firstResult = results[0]; const metadataStr = firstResult.metadata; if (metadataStr) { try { const metadata = JSON.parse(metadataStr); const codebasePath = metadata.codebasePath; if (codebasePath && typeof codebasePath === 'string') { console.log(`[SYNC-CLOUD] ๐Ÿ“ Found codebase path: ${codebasePath} in collection: ${collectionName}`); cloudCodebases.add(codebasePath); } else { console.warn(`[SYNC-CLOUD] โš ๏ธ No codebasePath found in metadata for collection: ${collectionName}`); } } catch (parseError) { console.warn(`[SYNC-CLOUD] โš ๏ธ Failed to parse metadata JSON for collection ${collectionName}:`, parseError); } } else { console.warn(`[SYNC-CLOUD] โš ๏ธ No metadata found in collection: ${collectionName}`); } } else { console.log(`[SYNC-CLOUD] โ„น๏ธ Collection ${collectionName} is empty`); } } catch (collectionError) { console.warn(`[SYNC-CLOUD] โš ๏ธ Error checking collection ${collectionName}:`, collectionError.message || collectionError); // Continue with next collection } } console.log(`[SYNC-CLOUD] ๐Ÿ“Š Found ${cloudCodebases.size} valid codebases in cloud`); // Get current local codebases const localCodebases = new Set(this.snapshotManager.getIndexedCodebases()); console.log(`[SYNC-CLOUD] ๐Ÿ“Š Found ${localCodebases.size} local codebases in snapshot`); let hasChanges = false; // Remove local codebases that don't exist in cloud for (const localCodebase of localCodebases) { if (!cloudCodebases.has(localCodebase)) { this.snapshotManager.removeIndexedCodebase(localCodebase); hasChanges = true; console.log(`[SYNC-CLOUD] โž– Removed local codebase (not in cloud): ${localCodebase}`); } } // Note: We don't add cloud codebases that are missing locally (as per user requirement) console.log(`[SYNC-CLOUD] โ„น๏ธ Skipping addition of cloud codebases not present locally (per sync policy)`); if (hasChanges) { this.snapshotManager.saveCodebaseSnapshot(); console.log(`[SYNC-CLOUD] ๐Ÿ’พ Updated snapshot to match cloud state`); } else { console.log(`[SYNC-CLOUD] โœ… Local snapshot already matches cloud state`); } console.log(`[SYNC-CLOUD] โœ… Cloud sync completed successfully`); } catch (error) { console.error(`[SYNC-CLOUD] โŒ Error syncing codebases from cloud:`, error.message || error); // Don't throw - this is not critical for the main functionality } } async handleIndexCodebase(args) { const { path: codebasePath, force, splitter, customExtensions, ignorePatterns } = args; const forceReindex = force || false; const splitterType = splitter || 'ast'; // Default to AST const customFileExtensions = customExtensions || []; const customIgnorePatterns = ignorePatterns || []; try { // Sync indexed codebases from cloud first await this.syncIndexedCodebasesFromCloud(); // Validate splitter parameter if (splitterType !== 'ast' && splitterType !== 'langchain') { return { content: [{ type: "text", text: `Error: Invalid splitter type '${splitterType}'. Must be 'ast' or 'langchain'.` }], isError: true }; } // Force absolute path resolution - warn if relative path provided const absolutePath = ensureAbsolutePath(codebasePath); // Validate path exists if (!fs.existsSync(absolutePath)) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' does not exist. Original input: '${codebasePath}'` }], isError: true }; } // Check if it's a directory const stat = fs.statSync(absolutePath); if (!stat.isDirectory()) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' is not a directory` }], isError: true }; } // Check if already indexing if (this.snapshotManager.getIndexingCodebases().includes(absolutePath)) { return { content: [{ type: "text", text: `Codebase '${absolutePath}' is already being indexed in the background. Please wait for completion.` }], isError: true }; } //Check if the snapshot and cloud index are in sync if (this.snapshotManager.getIndexedCodebases().includes(absolutePath) !== await this.context.hasIndex(absolutePath)) { console.warn(`[INDEX-VALIDATION] โŒ Snapshot and cloud index mismatch: ${absolutePath}`); } // Check if already indexed (unless force is true) if (!forceReindex && this.snapshotManager.getIndexedCodebases().includes(absolutePath)) { return { content: [{ type: "text", text: `Codebase '${absolutePath}' is already indexed. Use force=true to re-index.` }], isError: true }; } // If force reindex and codebase is already indexed, remove it if (forceReindex) { if (this.snapshotManager.getIndexedCodebases().includes(absolutePath)) { console.log(`[FORCE-REINDEX] ๐Ÿ”„ Removing '${absolutePath}' from indexed list for re-indexing`); this.snapshotManager.removeIndexedCodebase(absolutePath); } if (await this.context.hasIndex(absolutePath)) { console.log(`[FORCE-REINDEX] ๐Ÿ”„ Clearing index for '${absolutePath}'`); await this.context.clearIndex(absolutePath); } } // CRITICAL: Pre-index collection creation validation try { console.log(`[INDEX-VALIDATION] ๐Ÿ” Validating collection creation capability`); //dummy collection name const collectionName = `dummy_collection_${Date.now()}`; await this.context.getVectorDatabase().createCollection(collectionName, 128); if (await this.context.getVectorDatabase().hasCollection(collectionName)) { console.log(`[INDEX-VALIDATION] โ„น๏ธ Dummy collection created successfully`); await this.context.getVectorDatabase().dropCollection(collectionName); } else { console.log(`[INDEX-VALIDATION] โŒ Dummy collection creation failed`); } console.log(`[INDEX-VALIDATION] โœ… Collection creation validation completed`); } catch (validationError) { const errorMessage = typeof validationError === 'string' ? validationError : (validationError instanceof Error ? validationError.message : String(validationError)); if (errorMessage === COLLECTION_LIMIT_MESSAGE || errorMessage.includes(COLLECTION_LIMIT_MESSAGE)) { console.error(`[INDEX-VALIDATION] โŒ Collection limit validation failed: ${absolutePath}`); // CRITICAL: Immediately return the COLLECTION_LIMIT_MESSAGE to MCP client return { content: [{ type: "text", text: COLLECTION_LIMIT_MESSAGE }], isError: true }; } else { // Handle other collection creation errors console.error(`[INDEX-VALIDATION] โŒ Collection creation validation failed:`, validationError); return { content: [{ type: "text", text: `Error validating collection creation: ${validationError.message || validationError}` }], isError: true }; } } // Add custom extensions if provided if (customFileExtensions.length > 0) { console.log(`[CUSTOM-EXTENSIONS] Adding ${customFileExtensions.length} custom extensions: ${customFileExtensions.join(', ')}`); this.context.addCustomExtensions(customFileExtensions); } // Add custom ignore patterns if provided (before loading file-based patterns) if (customIgnorePatterns.length > 0) { console.log(`[IGNORE-PATTERNS] Adding ${customIgnorePatterns.length} custom ignore patterns: ${customIgnorePatterns.join(', ')}`); this.context.addCustomIgnorePatterns(customIgnorePatterns); } // Add to indexing list and save snapshot immediately this.snapshotManager.addIndexingCodebase(absolutePath); this.snapshotManager.saveCodebaseSnapshot(); // Track the codebase path for syncing trackCodebasePath(absolutePath); // Start background indexing - now safe to proceed this.startBackgroundIndexing(absolutePath, forceReindex, splitterType); const pathInfo = codebasePath !== absolutePath ? `\nNote: Input path '${codebasePath}' was resolved to absolute path '${absolutePath}'` : ''; const extensionInfo = customFileExtensions.length > 0 ? `\nUsing ${customFileExtensions.length} custom extensions: ${customFileExtensions.join(', ')}` : ''; const ignoreInfo = customIgnorePatterns.length > 0 ? `\nUsing ${customIgnorePatterns.length} custom ignore patterns: ${customIgnorePatterns.join(', ')}` : ''; return { content: [{ type: "text", text: `Started background indexing for codebase '${absolutePath}' using ${splitterType.toUpperCase()} splitter.${pathInfo}${extensionInfo}${ignoreInfo}\n\nIndexing is running in the background. You can search the codebase while indexing is in progress, but results may be incomplete until indexing completes.` }] }; } catch (error) { // Enhanced error handling to prevent MCP service crash console.error('Error in handleIndexCodebase:', error); // Ensure we always return a proper MCP response, never throw return { content: [{ type: "text", text: `Error starting indexing: ${error.message || error}` }], isError: true }; } } async startBackgroundIndexing(codebasePath, forceReindex, splitterType) { const absolutePath = codebasePath; let lastSaveTime = 0; // Track last save timestamp try { console.log(`[BACKGROUND-INDEX] Starting background indexing for: ${absolutePath}`); // Note: If force reindex, collection was already cleared during validation phase if (forceReindex) { console.log(`[BACKGROUND-INDEX] โ„น๏ธ Force reindex mode - collection was already cleared during validation`); } // Use the existing Context instance for indexing. let contextForThisTask = this.context; if (splitterType !== 'ast') { console.warn(`[BACKGROUND-INDEX] Non-AST splitter '${splitterType}' requested; falling back to AST splitter`); } // Load ignore patterns from files first (including .ignore, .gitignore, etc.) await this.context.getLoadedIgnorePatterns(absolutePath); // Initialize file synchronizer with proper ignore patterns (including project-specific patterns) const { FileSynchronizer } = await import("@zilliz/claude-context-core"); const ignorePatterns = this.context.getIgnorePatterns() || []; console.log(`[BACKGROUND-INDEX] Using ignore patterns: ${ignorePatterns.join(', ')}`); const synchronizer = new FileSynchronizer(absolutePath, ignorePatterns); await synchronizer.initialize(); // Store synchronizer in the context (let context manage collection names) await this.context.getPreparedCollection(absolutePath); const collectionName = this.context.getCollectionName(absolutePath); this.context.setSynchronizer(collectionName, synchronizer); if (contextForThisTask !== this.context) { contextForThisTask.setSynchronizer(collectionName, synchronizer); } console.log(`[BACKGROUND-INDEX] Starting indexing with ${splitterType} splitter for: ${absolutePath}`); // Log embedding provider information before indexing const embeddingProvider = this.context.getEmbedding(); console.log(`[BACKGROUND-INDEX] ๐Ÿง  Using embedding provider: ${embeddingProvider.getProvider()} with dimension: ${embeddingProvider.getDimension()}`); // Start indexing with the appropriate context and progress tracking console.log(`[BACKGROUND-INDEX] ๐Ÿš€ Beginning codebase indexing process...`); const stats = await contextForThisTask.indexCodebase(absolutePath, (progress) => { // Update progress in snapshot manager this.snapshotManager.updateIndexingProgress(absolutePath, progress.percentage); // Save snapshot periodically (every 2 seconds to avoid too frequent saves) const currentTime = Date.now(); if (currentTime - lastSaveTime >= 2000) { // 2 seconds = 2000ms this.snapshotManager.saveCodebaseSnapshot(); lastSaveTime = currentTime; console.log(`[BACKGROUND-INDEX] ๐Ÿ’พ Saved progress snapshot at ${progress.percentage.toFixed(1)}%`); } console.log(`[BACKGROUND-INDEX] Progress: ${progress.phase} - ${progress.percentage}% (${progress.current}/${progress.total})`); }); console.log(`[BACKGROUND-INDEX] โœ… Indexing completed successfully! Files: ${stats.indexedFiles}, Chunks: ${stats.totalChunks}`); // Move from indexing to indexed list this.snapshotManager.moveFromIndexingToIndexed(absolutePath); this.indexingStats = { indexedFiles: stats.indexedFiles, totalChunks: stats.totalChunks }; // Save snapshot after updating codebase lists this.snapshotManager.saveCodebaseSnapshot(); let message = `Background indexing completed for '${absolutePath}' using ${splitterType.toUpperCase()} splitter.\nIndexed ${stats.indexedFiles} files, ${stats.totalChunks} chunks.`; if (stats.status === 'limit_reached') { message += `\nโš ๏ธ Warning: Indexing stopped because the chunk limit (450,000) was reached. The index may be incomplete.`; } console.log(`[BACKGROUND-INDEX] ${message}`); } catch (error) { console.error(`[BACKGROUND-INDEX] Error during indexing for ${absolutePath}:`, error); // Remove from indexing list on error this.snapshotManager.removeIndexingCodebase(absolutePath); this.snapshotManager.saveCodebaseSnapshot(); // Log error but don't crash MCP service - indexing errors are handled gracefully console.error(`[BACKGROUND-INDEX] Indexing failed for ${absolutePath}: ${error.message || error}`); } } async handleSearchCode(args) { const { path: codebasePath, query, limit = 10 } = args; const resultLimit = limit || 10; try { // Sync indexed codebases from cloud first await this.syncIndexedCodebasesFromCloud(); // Force absolute path resolution - warn if relative path provided const absolutePath = ensureAbsolutePath(codebasePath); // Validate path exists if (!fs.existsSync(absolutePath)) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' does not exist. Original input: '${codebasePath}'` }], isError: true }; } // Check if it's a directory const stat = fs.statSync(absolutePath); if (!stat.isDirectory()) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' is not a directory` }], isError: true }; } trackCodebasePath(absolutePath); // Check if this codebase is indexed or being indexed const isIndexed = this.snapshotManager.getIndexedCodebases().includes(absolutePath); const isIndexing = this.snapshotManager.getIndexingCodebases().includes(absolutePath); if (!isIndexed && !isIndexing) { return { content: [{ type: "text", text: `Error: Codebase '${absolutePath}' is not indexed. Please index it first using the index_codebase tool.` }], isError: true }; } // Show indexing status if codebase is being indexed let indexingStatusMessage = ''; if (isIndexing) { indexingStatusMessage = `\nโš ๏ธ **Indexing in Progress**: This codebase is currently being indexed in the background. Search results may be incomplete until indexing completes.`; } console.log(`[SEARCH] Searching in codebase: ${absolutePath}`); console.log(`[SEARCH] Query: "${query}"`); console.log(`[SEARCH] Indexing status: ${isIndexing ? 'In Progress' : 'Completed'}`); // Log embedding provider information before search const embeddingProvider = this.context.getEmbedding(); console.log(`[SEARCH] ๐Ÿง  Using embedding provider: ${embeddingProvider.getProvider()} for search`); console.log(`[SEARCH] ๐Ÿ” Generating embeddings for query using ${embeddingProvider.getProvider()}...`); // Search in the specified codebase const searchResults = await this.context.semanticSearch(absolutePath, query, Math.min(resultLimit, 50), 0.3); console.log(`[SEARCH] โœ… Search completed! Found ${searchResults.length} results using ${embeddingProvider.getProvider()} embeddings`); if (searchResults.length === 0) { let noResultsMessage = `No results found for query: "${query}" in codebase '${absolutePath}'`; if (isIndexing) { noResultsMessage += `\n\nNote: This codebase is still being indexed. Try searching again after indexing completes, or the query may not match any indexed content.`; } return { content: [{ type: "text", text: noResultsMessage }] }; } // Format results const formattedResults = searchResults.map((result, index) => { const location = `${result.relativePath}:${result.startLine}-${result.endLine}`; const context = truncateContent(result.content, 5000); const codebaseInfo = path.basename(absolutePath); return `${index + 1}. Code snippet (${result.language}) [${codebaseInfo}]\n` + ` Location: ${location}\n` + ` Rank: ${index + 1}\n` + ` Context: \n\`\`\`${result.language}\n${context}\n\`\`\`\n`; }).join('\n'); let resultMessage = `Found ${searchResults.length} results for query: "${query}" in codebase '${absolutePath}'${indexingStatusMessage}\n\n${formattedResults}`; if (isIndexing) { resultMessage += `\n\n๐Ÿ’ก **Tip**: This codebase is still being indexed. More results may become available as indexing progresses.`; } return { content: [{ type: "text", text: resultMessage }] }; } catch (error) { // Check if this is the collection limit error // Handle both direct string throws and Error objects containing the message const errorMessage = typeof error === 'string' ? error : (error instanceof Error ? error.message : String(error)); if (errorMessage === COLLECTION_LIMIT_MESSAGE || errorMessage.includes(COLLECTION_LIMIT_MESSAGE)) { // Return the collection limit message as a successful response // This ensures LLM treats it as final answer, not as retryable error return { content: [{ type: "text", text: COLLECTION_LIMIT_MESSAGE }] }; } return { content: [{ type: "text", text: `Error searching code: ${errorMessage} Please check if the codebase has been indexed first.` }], isError: true }; } } async handleClearIndex(args) { const { path: codebasePath } = args; if (this.snapshotManager.getIndexedCodebases().length === 0 && this.snapshotManager.getIndexingCodebases().length === 0) { return { content: [{ type: "text", text: "No codebases are currently indexed or being indexed." }] }; } try { // Force absolute path resolution - warn if relative path provided const absolutePath = ensureAbsolutePath(codebasePath); // Validate path exists if (!fs.existsSync(absolutePath)) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' does not exist. Original input: '${codebasePath}'` }], isError: true }; } // Check if it's a directory const stat = fs.statSync(absolutePath); if (!stat.isDirectory()) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' is not a directory` }], isError: true }; } // Check if this codebase is indexed or being indexed const isIndexed = this.snapshotManager.getIndexedCodebases().includes(absolutePath); const isIndexing = this.snapshotManager.getIndexingCodebases().includes(absolutePath); if (!isIndexed && !isIndexing) { return { content: [{ type: "text", text: `Error: Codebase '${absolutePath}' is not indexed or being indexed.` }], isError: true }; } console.log(`[CLEAR] Clearing codebase: ${absolutePath}`); try { await this.context.clearIndex(absolutePath); console.log(`[CLEAR] Successfully cleared index for: ${absolutePath}`); } catch (error) { const errorMsg = `Failed to clear ${absolutePath}: ${error.message}`; console.error(`[CLEAR] ${errorMsg}`); return { content: [{ type: "text", text: errorMsg }], isError: true }; } // Remove the cleared codebase from both lists this.snapshotManager.removeIndexedCodebase(absolutePath); this.snapshotManager.removeIndexingCodebase(absolutePath); // Reset indexing stats if this was the active codebase this.indexingStats = null; // Save snapshot after clearing index this.snapshotManager.saveCodebaseSnapshot(); let resultText = `Successfully cleared codebase '${absolutePath}'`; const remainingIndexed = this.snapshotManager.getIndexedCodebases().length; const remainingIndexing = this.snapshotManager.getIndexingCodebases().length; if (remainingIndexed > 0 || remainingIndexing > 0) { resultText += `\n${remainingIndexed} other indexed codebase(s) and ${remainingIndexing} indexing codebase(s) remain`; } return { content: [{ type: "text", text: resultText }] }; } catch (error) { // Check if this is the collection limit error // Handle both direct string throws and Error objects containing the message const errorMessage = typeof error === 'string' ? error : (error instanceof Error ? error.message : String(error)); if (errorMessage === COLLECTION_LIMIT_MESSAGE || errorMessage.includes(COLLECTION_LIMIT_MESSAGE)) { // Return the collection limit message as a successful response // This ensures LLM treats it as final answer, not as retryable error return { content: [{ type: "text", text: COLLECTION_LIMIT_MESSAGE }] }; } return { content: [{ type: "text", text: `Error clearing index: ${errorMessage}` }], isError: true }; } } async handleGetIndexingStatus(args) { const { path: codebasePath } = args; try { // Force absolute path resolution const absolutePath = ensureAbsolutePath(codebasePath); // Validate path exists if (!fs.existsSync(absolutePath)) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' does not exist. Original input: '${codebasePath}'` }], isError: true }; } // Check if it's a directory const stat = fs.statSync(absolutePath); if (!stat.isDirectory()) { return { content: [{ type: "text", text: `Error: Path '${absolutePath}' is not a directory` }], isError: true }; } // Check indexing status const isIndexed = this.snapshotManager.getIndexedCodebases().includes(absolutePath); const isIndexing = this.snapshotManager.getIndexingCodebases().includes(absolutePath); const indexingProgress = this.snapshotManager.getIndexingProgress(absolutePath); let statusMessage = ''; if (isIndexed) { statusMessage = `โœ… Codebase '${absolutePath}' is fully indexed and ready for search.`; } else if (isIndexing) { const progressPercentage = indexingProgress !== undefined ? indexingProgress : 0; statusMessage = `๐Ÿ”„ Codebase '${absolutePath}' is currently being indexed. Progress: ${progressPercentage.toFixed(1)}%`; // Add more detailed status based on progress if (progressPercentage < 10) { statusMessage += ' (Preparing and scanning files...)'; } else if (progressPercentage < 100) { statusMessage += ' (Processing files and generating embeddings...)'; } } else { statusMessage = `โŒ Codebase '${absolutePath}' is not indexed. Please use the index_codebase tool to index it first.`; } const pathInfo = codebasePath !== absolutePath ? `\nNote: Input path '${codebasePath}' was resolved to absolute path '${absolutePath}'` : ''; return { content: [{ type: "text", text: statusMessage + pathInfo }] }; } catch (error) { return { content: [{ type: "text", text: `Error getting indexing status: ${error.message || error}` }], isError: true }; } } } //# sourceMappingURL=handlers.js.map