How to Implement Semantic Search in VS Code Extensions

Introduction to Semantic Search in VS Code Extensions

Traditional search in code editors relies on exact keyword matching or regular expressions. While effective for finding specific identifiers, it falls short when developers search by intent — for example, "where do we handle user authentication failures" won't match a function named validateCredentials using lexical search. Semantic search bridges this gap by understanding the meaning behind both the query and the indexed content, returning results based on conceptual similarity rather than string overlap.

In this tutorial, you'll learn how to build a VS Code extension that implements semantic search over a workspace. We'll cover embedding generation, vector storage, similarity computation, and integration with the VS Code API. By the end, your extension will let users search their codebase using natural language queries and receive ranked, context-aware results.

Why Semantic Search Matters in VS Code Extensions

Developers spend a significant portion of their time navigating and understanding unfamiliar code. Semantic search transforms this experience in several ways:

• Intent-based discovery: Users can search for concepts like "database connection retry logic" and find relevant code even if those exact words never appear together.
• Cross-language relevance: A query in English can match code written in any programming language, since embeddings capture semantic meaning across modalities.
• Onboarding acceleration: New team members can explore a codebase without knowing naming conventions or architectural patterns in advance.
• Reduced cognitive load: Instead of crafting precise regex patterns, developers express what they're looking for naturally.

For extension authors, adding semantic search differentiates your tool from the dozens of keyword-based search extensions already available, and it opens the door to more advanced features like AI-assisted code navigation and contextual documentation retrieval.

Prerequisites and Setup

Before we begin, ensure you have the following:

• Node.js 18 or later installed
• Visual Studio Code (version 1.80 or later recommended)
• Basic familiarity with TypeScript and the VS Code Extension API
• The yo code generator installed (npm install -g yo generator-code)

Generate a new extension project and install the dependencies we'll need for embedding generation and vector operations:

yo code
# Choose: New Extension (TypeScript)
# Name: semantic-search-demo

cd semantic-search-demo
npm install @xenova/transformers
npm install vectordb
npm install cosine-similarity

We'll use @xenova/transformers to run a lightweight embedding model locally within the extension host process. This avoids requiring users to supply API keys and keeps all data on their machine, which is critical for code that may be proprietary.

Understanding the Architecture

A semantic search system has three core components: an embedding model that converts text into high-dimensional vectors, an index that stores these vectors efficiently, and a query pipeline that embeds the user's search string and compares it against the index. In a VS Code extension, all of this runs inside the extension host, a Node.js process that VS Code spawns when your extension activates.

The workflow is straightforward. When the extension activates or when the user triggers indexing, we read source files from the workspace, split them into meaningful chunks (such as individual functions or classes), generate an embedding vector for each chunk, and store these vectors along with metadata like file path and line numbers. At query time, we embed the user's query, compute cosine similarity against all stored vectors, and return the top matches.

Choosing an Embedding Model

The embedding model determines the quality of your search results. For a VS Code extension, you need a model that balances accuracy, size, and inference speed, since it runs locally on the user's machine. The all-MiniLM-L6-v2 model from sentence-transformers is an excellent default choice. It produces 384-dimensional vectors, is roughly 80 MB, and runs efficiently on CPU.

For code-specific use cases, you might consider models like codebert-base or unixcoder-base, which are pre-trained on source code. However, general-purpose sentence embedding models often perform surprisingly well on code because they capture structural and semantic patterns that transfer across text and code domains. Start with all-MiniLM-L6-v2 and experiment with alternatives if search quality needs improvement.

Implementing the Embedding Service

Create a file named src/embeddings.ts that wraps the transformers.js pipeline. We'll lazy-load the model so it only initializes when first needed, avoiding startup overhead.

import { pipeline, Pipeline } from '@xenova/transformers';

let embedder: Pipeline | null = null;

export async function getEmbedder(): Promise<Pipeline> {
    if (!embedder) {
        embedder = await pipeline(
            'feature-extraction',
            'Xenova/all-MiniLM-L6-v2'
        );
    }
    return embedder;
}

export async function generateEmbedding(text: string): Promise<number[]> {
    const extractor = await getEmbedder();
    const output = await extractor(text, {
        pooling: 'mean',
        normalize: true,
    });
    return Array.from(output.data as Float32Array);
}

export async function generateEmbeddings(
    texts: string[]
): Promise<number[][]> {
    const results: number[][] = [];
    for (const text of texts) {
        const embedding = await generateEmbedding(text);
        results.push(embedding);
    }
    return results;
}

The pooling: 'mean' option averages token embeddings into a single sentence-level vector, and normalize: true scales the vector to unit length, which simplifies similarity calculations to a simple dot product.

Chunking Source Files

Embedding an entire file as one vector produces poor results because the vector must represent too much information. Instead, we split files into smaller, semantically meaningful chunks. For code, a good strategy is to split on function and class boundaries using a simple regex-based parser.

Create src/chunker.ts:

export interface CodeChunk {
    filePath: string;
    startLine: number;
    endLine: number;
    content: string;
    language: string;
}

export function chunkFile(
    filePath: string,
    content: string,
    language: string
): CodeChunk[] {
    const lines = content.split('\n');
    const chunks: CodeChunk[] = [];
    const maxChunkLines = 60;

    // Simple heuristic: split on top-level function/class declarations
    const boundaryPattern =
        /^(export\s+)?(async\s+)?(function|class|const|let|var|def)\s+\w+/;

    let currentStart = 0;
    let currentLines: string[] = [];

    for (let i = 0; i < lines.length; i++) {
        const line = lines[i];

        if (boundaryPattern.test(line.trim()) && currentLines.length > 0) {
            chunks.push(createChunk(filePath, currentStart, i - 1, currentLines, language));
            currentStart = i;
            currentLines = [];
        }

        currentLines.push(line);

        if (currentLines.length >= maxChunkLines) {
            chunks.push(createChunk(filePath, currentStart, i, currentLines, language));
            currentStart = i + 1;
            currentLines = [];
        }
    }

    if (currentLines.length > 0) {
        chunks.push(createChunk(filePath, currentStart, lines.length - 1, currentLines, language));
    }

    return chunks;
}

function createChunk(
    filePath: string,
    startLine: number,
    endLine: number,
    lines: string[],
    language: string
): CodeChunk {
    return {
        filePath,
        startLine,
        endLine,
        content: lines.join('\n'),
        language,
    };
}

This is a deliberately simple chunker. For production use, consider integrating a proper AST parser for each language you want to support, which will produce more accurate boundaries. The maxChunkLines limit prevents excessively large chunks that dilute embedding quality.

Building the Vector Index

Now we need a storage layer that holds our embeddings and supports fast similarity queries. For workspaces with up to a few thousand chunks, an in-memory brute-force search is perfectly adequate and avoids external dependencies. Create src/vectorStore.ts:

import { CodeChunk } from './chunker';

export interface IndexedChunk extends CodeChunk {
    embedding: number[];
}

export class VectorStore {
    private chunks: IndexedChunk[] = [];

    addChunk(chunk: CodeChunk, embedding: number[]): void {
        this.chunks.push({ ...chunk, embedding });
    }

    clear(): void {
        this.chunks = [];
    }

    size(): number {
        return this.chunks.length;
    }

    search(queryEmbedding: number[], topK: number = 10): SearchResult[] {
        const scores = this.chunks.map((chunk) => ({
            chunk,
            score: cosineSimilarity(queryEmbedding, chunk.embedding),
        }));

        scores.sort((a, b) => b.score - a.score);
        return scores.slice(0, topK).map((s) => ({
            chunk: s.chunk,
            score: s.score,
        }));
    }
}

export interface SearchResult {
    chunk: IndexedChunk;
    score: number;
}

function cosineSimilarity(a: number[], b: number[]): number {
    let dotProduct = 0;
    for (let i = 0; i < a.length; i++) {
        dotProduct += a[i] * b[i];
    }
    return dotProduct; // Vectors are pre-normalized, so dot product = cosine
}

Because we normalized embeddings during generation, cosine similarity reduces to a simple dot product, which is fast even for thousands of vectors. If your workspace grows beyond tens of thousands of chunks, consider integrating a proper vector database like ChromaDB or LanceDB, which support approximate nearest neighbor search for sub-linear query times.

Indexing the Workspace

With the embedding service, chunker, and vector store in place, we can build the indexing pipeline. Create src/indexer.ts:

import * as vscode from 'vscode';
import * as path from 'path';
import { chunkFile, CodeChunk } from './chunker';
import { generateEmbedding } from './embeddings';
import { VectorStore } from './vectorStore';

const SUPPORTED_EXTENSIONS = [
    '.ts', '.js', '.py', '.java', '.go', '.rs',
    '.cpp', '.c', '.rb', '.php', '.cs', '.swift',
    '.kt', '.scala', '.jsx', '.tsx', '.vue', '.svelte',
];

const IGNORED_DIRS = [
    'node_modules', '.git', 'dist', 'build', '.vscode',
    'out', 'coverage', '__pycache__', '.next', 'vendor',
];

export class WorkspaceIndexer {
    constructor(private store: VectorStore) {}

    async indexWorkspace(
        progress?: vscode.Progress<{ message?: string; increment?: number }>,
        token?: vscode.CancellationToken
    ): Promise<number> {
        this.store.clear();
        const workspaceFolders = vscode.workspace.workspaceFolders;
        if (!workspaceFolders) {
            return 0;
        }

        const files = await this.collectFiles(workspaceFolders[0].uri);
        let processed = 0;

        for (const fileUri of files) {
            if (token?.isCancellationRequested) {
                break;
            }

            try {
                const document = await vscode.workspace.openTextDocument(fileUri);
                const content = document.getText();
                const language = this.getLanguageId(document.languageId);
                const chunks = chunkFile(fileUri.fsPath, content, language);

                for (const chunk of chunks) {
                    const embedding = await generateEmbedding(chunk.content);
                    this.store.addChunk(chunk, embedding);
                }

                processed++;
                if (progress) {
                    progress.report({
                        message: `Indexed ${processed}/${files.length} files`,
                        increment: (1 / files.length) * 100,
                    });
                }
            } catch (error) {
                console.error(`Failed to index ${fileUri.fsPath}:`, error);
            }
        }

        return this.store.size();
    }

    private async collectFiles(rootUri: vscode.Uri): Promise<vscode.Uri[]> {
        const files: vscode.Uri[] = [];
        await this.walkDirectory(rootUri, files);
        return files;
    }

    private async walkDirectory(uri: vscode.Uri, files: vscode.Uri[]): Promise<void> {
        const entries = await vscode.workspace.fs.readDirectory(uri);

        for (const [name, type] of entries) {
            const childUri = vscode.Uri.joinPath(uri, name);

            if (type === vscode.FileType.Directory) {
                if (IGNORED_DIRS.includes(name)) continue;
                await this.walkDirectory(childUri, files);
            } else if (type === vscode.FileType.File) {
                const ext = path.extname(name).toLowerCase();
                if (SUPPORTED_EXTENSIONS.includes(ext)) {
                    files.push(childUri);
                }
            }
        }
    }

    private getLanguageId(vscodeLangId: string): string {
        return vscodeLangId;
    }
}

The indexer walks the workspace directory tree, skips common directories that shouldn't be indexed, and processes each supported file. The progress callback integrates with VS Code's progress API so users see real-time feedback during indexing.

Implementing the Search Command

Now let's wire everything together in the extension's main entry point. We'll register a command that prompts the user for a query, performs semantic search, and displays results in a QuickPick. Create or update src/extension.ts:

import * as vscode from 'vscode';
import { VectorStore } from './vectorStore';
import { WorkspaceIndexer } from './indexer';
import { generateEmbedding } from './embeddings';

const store = new VectorStore();
let indexer: WorkspaceIndexer;

export function activate(context: vscode.ExtensionContext) {
    indexer = new WorkspaceIndexer(store);

    const indexCommand = vscode.commands.registerCommand(
        'semanticSearch.indexWorkspace',
        async () => {
            await vscode.window.withProgress(
                {
                    location: vscode.ProgressLocation.Notification,
                    title: 'Building semantic index',
                    cancellable: true,
                },
                async (progress, token) => {
                    const chunkCount = await indexer.indexWorkspace(progress, token);
                    vscode.window.showInformationMessage(
                        `Semantic index built with ${chunkCount} chunks.`
                    );
                }
            );
        }
    );

    const searchCommand = vscode.commands.registerCommand(
        'semanticSearch.search',
        async () => {
            if (store.size() === 0) {
                const choice = await vscode.window.showWarningMessage(
                    'No index found. Build the index first?',
                    'Build Index',
                    'Cancel'
                );
                if (choice !== 'Build Index') return;
                await vscode.commands.executeCommand('semanticSearch.indexWorkspace');
            }

            const query = await vscode.window.showInputBox({
                prompt: 'Enter a semantic search query',
                placeHolder: 'e.g., where do we handle user authentication',
            });

            if (!query) return;

            const queryEmbedding = await generateEmbedding(query);
            const results = store.search(queryEmbedding, 15);

            if (results.length === 0) {
                vscode.window.showInformationMessage('No results found.');
                return;
            }

            const items = results.map((result) => ({
                label: result.chunk.filePath.split('/').pop() || result.chunk.filePath,
                description: `Score: ${result.score.toFixed(3)}`,
                detail: `Lines ${result.chunk.startLine + 1}-${result.chunk.endLine + 1}`,
                filePath: result.chunk.filePath,
                startLine: result.chunk.startLine,
                endLine: result.chunk.endLine,
            }));

            const selected = await vscode.window.showQuickPick(items, {
                placeHolder: 'Semantic search results',
                matchOnDescription: true,
                matchOnDetail: true,
            });

            if (selected) {
                const uri = vscode.Uri.file(selected.filePath);
                const document = await vscode.workspace.openTextDocument(uri);
                const editor = await vscode.window.showTextDocument(document);

                const range = new vscode.Range(
                    selected.startLine,
                    0,
                    selected.endLine,
                    0
                );
                editor.revealRange(range, vscode.TextEditorRevealType.InCenter);
                editor.selection = new vscode.Selection(range.start, range.start);
            }
        }
    );

    context.subscriptions.push(indexCommand, searchCommand);
}

export function deactivate() {}

Don't forget to register the commands in your package.json under the contributes.commands section:

{
    "contributes": {
        "commands": [
            {
                "command": "semanticSearch.indexWorkspace",
                "title": "Semantic Search: Index Workspace"
            },
            {
                "command": "semanticSearch.search",
                "title": "Semantic Search: Search"
            }
        ]
    }
}

Persisting the Index

Rebuilding the index on every activation is wasteful. Let's add persistence so the index loads from disk when the extension starts. We'll serialize the vector store to a JSON file in the extension's global storage directory. Add these methods to VectorStore:

export class VectorStore {
    // ... existing code ...

    async save(filePath: string): Promise<void> {
        const fs = await import('fs/promises');
        const data = JSON.stringify(this.chunks);
        await fs.writeFile(filePath, data, 'utf-8');
    }

    async load(filePath: string): Promise<boolean> {
        try {
            const fs = await import('fs/promises');
            const data = await fs.readFile(filePath, 'utf-8');
            this.chunks = JSON.parse(data);
            return true;
        } catch {
            return false;
        }
    }
}

Then update extension.ts to load and save the index:

export async function activate(context: vscode.ExtensionContext) {
    indexer = new WorkspaceIndexer(store);

    const indexPath = vscode.Uri.joinPath(
        context.globalStorageUri,
        'semantic-index.json'
    ).fsPath;

    // Ensure storage directory exists
    await vscode.workspace.fs.createDirectory(context.globalStorageUri);

    // Try loading existing index
    const loaded = await store.load(indexPath);
    if (loaded) {
        console.log(`Loaded ${store.size()} chunks from disk.`);
    }

    // Update the index command to save after building
    const indexCommand = vscode.commands.registerCommand(
        'semanticSearch.indexWorkspace',
        async () => {
            await vscode.window.withProgress(
                {
                    location: vscode.ProgressLocation.Notification,
                    title: 'Building semantic index',
                    cancellable: true,
                },
                async (progress, token) => {
                    const chunkCount = await indexer.indexWorkspace(progress, token);
                    await store.save(indexPath);
                    vscode.window.showInformationMessage(
                        `Semantic index built with ${chunkCount} chunks.`
                    );
                }
            );
        }
    );

    // ... rest of activation ...
}

Handling File Changes with Watchers

A static index becomes stale as developers edit files. To keep results relevant, register a file system watcher that re-indexes changed files incrementally. Add this to your activation logic:

const watcher = vscode.workspace.createFileSystemWatcher(
    '**/*.{ts,js,py,java,go,rs,jsx,tsx}',
    false,
    false,
    false
);

watcher.onDidChange(async (uri) => {
    await reindexFile(uri);
});

watcher.onDidCreate(async (uri) => {
    await reindexFile(uri);
});

watcher.onDidDelete((uri) => {
    store.removeByFilePath(uri.fsPath);
});

context.subscriptions.push(watcher);

async function reindexFile(uri: vscode.Uri): Promise<void> {
    try {
        const document = await vscode.workspace.openTextDocument(uri);
        const content = document.getText();
        const chunks = chunkFile(uri.fsPath, content, document.languageId);

        store.removeByFilePath(uri.fsPath);

        for (const chunk of chunks) {
            const embedding = await generateEmbedding(chunk.content);
            store.addChunk(chunk, embedding);
        }

        await store.save(indexPath);
    } catch (error) {
        console.error(`Failed to re-index ${uri.fsPath}:`, error);
    }
}

You'll need to add a removeByFilePath method to VectorStore:

removeByFilePath(filePath: string): void {
    this.chunks = this.chunks.filter(c => c.filePath !== filePath);
}

Adding a Search Results Webview

While QuickPick is functional, a webview panel provides a richer experience with syntax highlighting and code previews. Let's add a command that opens results in a webview:

const searchPanelCommand = vscode.commands.registerCommand(
    'semanticSearch.searchPanel',
    async () => {
        const query = await vscode.window.showInputBox({
            prompt: 'Enter a semantic search query',
        });

        if (!query) return;

        const queryEmbedding = await generateEmbedding(query);
        const results = store.search(queryEmbedding, 20);

        const panel = vscode.window.createWebviewPanel(
            'semanticSearchResults',
            `Semantic Search: "${query}"`,
            vscode.ViewColumn.One,
            { enableScripts: true }
        );

        panel.webview.html = renderResultsHtml(query, results);
    }
);

function renderResultsHtml(query: string, results: SearchResult[]): string {
    const resultCards = results.map((r, i) => `
        <div class="result" data-file="${r.chunk.filePath}" data-line="${r.chunk.startLine}">
            <div class="result-header">
                <span class="filename">${r.chunk.filePath.split('/').pop()}</span>
                <span class="score">${(r.score * 100).toFixed(1)}% match</span>
            </div>
            <div class="path">${r.chunk.filePath}:${r.chunk.startLine + 1}</div>
            <pre class="code"><code>${escapeHtml(r.chunk.content.slice(0, 500))}</code></pre>
        </div>
    `).join('');

    return `
        <!DOCTYPE html>
        <html>
        <head>
            <style>
                body { font-family: var(--vscode-font-family); padding: 16px; color: var(--vscode-foreground); }
                .query { font-size: 1.2em; margin-bottom: 16px; color: var(--vscode-textLink-foreground); }
                .result { margin-bottom: 20px; border: 1px solid var(--vscode-panel-border); border-radius: 4px; padding: 12px; cursor: pointer; }
                .result:hover { background: var(--vscode-list-hoverBackground); }
                .result-header { display: flex; justify-content: space-between; margin-bottom: 4px; }
                .filename { font-weight: bold; }
                .score { color: var(--vscode-descriptionForeground); }
                .path { font-size: 0.85em; color: var(--vscode-descriptionForeground); margin-bottom: 8px; }
                .code { background: var(--vscode-textCodeBlock-background); padding: 8px; border-radius: 4px; overflow-x: auto; font-size: 0.85em; }
            </style>
        </head>
        <body>
            <div class="query">Results for: "${escapeHtml(query)}"</div>
            ${resultCards}
            <script>
                const vscode = acquireVsCodeApi();
                document.querySelectorAll('.result').forEach(el => {
                    el.addEventListener('click', () => {
                        vscode.postMessage({
                            command: 'openFile',
                            filePath: el.dataset.file,
                            line: parseInt(el.dataset.line)
                        });
                    });
                });
            </script>
        </body>
        </html>
    `;
}

function escapeHtml(text: string): string {
    return text
        .replace(/&/g, '&')
        .replace(/</g, '<')
        .replace(/>/g, '>')
        .replace(/"/g, '"');
}

Add a message handler so clicking a result opens the file at the correct location:

panel.webview.onDidReceiveMessage(
    async (message) => {
        if (message.command === 'openFile') {
            const uri = vscode.Uri.file(message.filePath);
            const document = await vscode.workspace.openTextDocument(uri);
            const editor = await vscode.window.showTextDocument(document);
            const position = new vscode.Position(message.line, 0);
            editor.revealRange(
                new vscode.Range(position, position),
                vscode.TextEditorRevealType.InCenter
            );
            editor.selection = new vscode.Selection(position, position);
        }
    },
    undefined,
    context.subscriptions
);

Best Practices

As you refine your semantic search extension, keep these principles in mind:

• Lazy-load the model: The embedding model consumes memory and takes time to initialize. Load it only when the user first triggers a search or indexing operation, not at activation time.
• Respect resource constraints: Running a transformer model in the extension host can be CPU-intensive. Consider showing a warning for large workspaces and allowing users to configure which directories to include or exclude.
• Provide incremental feedback: Indexing can take minutes for large codebases. Always use VS Code's progress API and make the operation cancellable.
• Cache embeddings intelligently: If a file hasn't changed, its embeddings haven't either. Store a hash of each file's content alongside its embeddings and skip re-embedding unchanged files during re-indexing.
• Handle edge cases: Empty files, binary files with code extensions, files with encoding issues, and extremely long lines can all cause problems. Wrap file reading and embedding generation in try-catch blocks and log errors without crashing.
• Offer configuration options: Let users tune the number of results returned, the chunk size, and which file types to include through VS Code settings. Different projects have different needs.
• Test with diverse codebases: Semantic search quality varies across languages and coding styles. Test your extension against projects written in multiple languages to ensure the chunker and embedding model perform well broadly.
• Consider privacy: If you ever move to a cloud-based embedding API, make this explicit to users and provide a local-only option. Developers may be searching proprietary code, and sending it to a third-party service without consent is a serious trust violation.

Conclusion

Implementing semantic search in a VS Code extension transforms how developers navigate and understand code. By combining a local embedding model, a simple vector store, and the VS Code Extension API, you can build a tool that understands developer intent rather than just matching strings. The architecture we've covered — chunking, embedding, indexing, querying, and incremental updates — provides a solid foundation that you can extend with features like cross-repository search, natural language code summarization, or integration with AI-powered code assistants. Start with the minimal implementation above, test it against your own projects, and iterate on the chunking strategy and model choice to maximize search quality for your specific use case.