From 22c02ba30375b0bd4713178b97204126180e2bf4 Mon Sep 17 00:00:00 2001 From: Alexandre Vilain Date: Sun, 31 Aug 2025 14:53:11 +0200 Subject: [PATCH] feat(completion): add support for LSP to get better completions --- src/autocomplete/defaultHoleFiller.ts | 90 +++++++ src/autocomplete/holeFiller.ts | 2 + src/extension.ts | 11 +- src/services/lspProvider.ts | 332 ++++++++++++++++++++++++++ src/types/index.ts | 3 +- src/types/lsp.ts | 24 ++ src/vscode/completionProvider.ts | 52 +++- 7 files changed, 511 insertions(+), 3 deletions(-) create mode 100644 src/services/lspProvider.ts create mode 100644 src/types/lsp.ts diff --git a/src/autocomplete/defaultHoleFiller.ts b/src/autocomplete/defaultHoleFiller.ts index bda032c..e70ad36 100644 --- a/src/autocomplete/defaultHoleFiller.ts +++ b/src/autocomplete/defaultHoleFiller.ts @@ -1,4 +1,6 @@ import { HoleFiller, PromptArgs, AutoCompleteContext } from "./holeFiller"; +import { LSPDefinition } from "../types/lsp"; +import * as vscode from 'vscode'; // Source: continue/core/autocomplete/templating/AutocompleteTemplate.ts (holeFillerTemplate) export class DefaultHoleFiller implements HoleFiller { @@ -99,9 +101,97 @@ function hypothenuse(a, b) { if (ctx.language !== '') { context += `// Programming language: "${ctx.language}" \n`; } + + // Add LSP context if available + if (ctx.lspContext && ctx.lspContext.definitions.length > 0) { + context += this.formatLSPContext(ctx.lspContext.definitions); + } + return `${context}\n${ctx.textBeforeCursor}{{FILL_HERE}}${ctx.textAfterCursor}\n\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n`; } + private formatLSPContext(definitions: LSPDefinition[]): string { + if (definitions.length === 0) {return '';} + + let context = '// Available definitions and symbols:\n'; + + // Group definitions by type for better organization + const functions = definitions.filter(d => [vscode.SymbolKind.Function, vscode.SymbolKind.Method].includes(d.kind)); + const classes = definitions.filter(d => [vscode.SymbolKind.Class, vscode.SymbolKind.Interface].includes(d.kind)); + const variables = definitions.filter(d => [vscode.SymbolKind.Variable, vscode.SymbolKind.Property, vscode.SymbolKind.Field].includes(d.kind)); + const types = definitions.filter(d => [vscode.SymbolKind.Enum, vscode.SymbolKind.Struct, vscode.SymbolKind.TypeParameter].includes(d.kind)); + + // Add functions and methods + if (functions.length > 0) { + context += '// Functions/Methods:\n'; + functions.slice(0, 10).forEach(def => { + const signature = this.formatDefinitionSignature(def); + context += `// ${signature}\n`; + }); + } + + // Add classes and interfaces + if (classes.length > 0) { + context += '// Classes/Interfaces:\n'; + classes.slice(0, 8).forEach(def => { + const signature = this.formatDefinitionSignature(def); + context += `// ${signature}\n`; + }); + } + + // Add variables and properties + if (variables.length > 0) { + context += '// Variables/Properties:\n'; + variables.slice(0, 8).forEach(def => { + const signature = this.formatDefinitionSignature(def); + context += `// ${signature}\n`; + }); + } + + // Add types + if (types.length > 0) { + context += '// Types:\n'; + types.slice(0, 5).forEach(def => { + const signature = this.formatDefinitionSignature(def); + context += `// ${signature}\n`; + }); + } + + return context; + } + + private formatDefinitionSignature(def: LSPDefinition): string { + const kindName = this.getSymbolKindName(def.kind); + let signature = `${def.name}`; + + if (def.detail) { + signature += `: ${def.detail}`; + } + + if (def.containerName) { + signature = `${def.containerName}.${signature}`; + } + + return `${kindName} ${signature}`; + } + + private getSymbolKindName(kind: vscode.SymbolKind): string { + switch (kind) { + case vscode.SymbolKind.Function: return 'function'; + case vscode.SymbolKind.Method: return 'method'; + case vscode.SymbolKind.Class: return 'class'; + case vscode.SymbolKind.Interface: return 'interface'; + case vscode.SymbolKind.Variable: return 'var'; + case vscode.SymbolKind.Property: return 'prop'; + case vscode.SymbolKind.Field: return 'field'; + case vscode.SymbolKind.Enum: return 'enum'; + case vscode.SymbolKind.Struct: return 'struct'; + case vscode.SymbolKind.TypeParameter: return 'type'; + case vscode.SymbolKind.Constant: return 'const'; + default: return 'symbol'; + } + } + prompt(params: AutoCompleteContext): PromptArgs { return { messages: [ diff --git a/src/autocomplete/holeFiller.ts b/src/autocomplete/holeFiller.ts index f62d8bc..3941a1c 100644 --- a/src/autocomplete/holeFiller.ts +++ b/src/autocomplete/holeFiller.ts @@ -1,5 +1,6 @@ import { Prompt } from 'ai'; import { ProviderOptions } from '@ai-sdk/provider-utils'; +import { LSPContext } from '../types/lsp'; export interface HoleFiller { prompt(params: AutoCompleteContext): PromptArgs @@ -15,4 +16,5 @@ export type AutoCompleteContext = { currentLineText: string, filename?: string, language?: string, + lspContext?: LSPContext, } \ No newline at end of file diff --git a/src/extension.ts b/src/extension.ts index dd8ff07..47b646a 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -5,6 +5,7 @@ import { TabCoderStatusBarProvider } from './vscode/statusBarProvider'; import { ProfileCommandProvider } from './vscode/profileCommandProvider'; import { logger } from './utils/logger'; import { ProfileService } from './services/profileService'; +import { LSPProvider } from './services/lspProvider'; export function activate(context: vscode.ExtensionContext) { // Initialize ConfigurationProvider with context for secure storage. @@ -23,8 +24,16 @@ export function activate(context: vscode.ExtensionContext) { vscode.commands.registerCommand('tabcoder.statusBarClicked', () => statusBarProvider.handleStatusBarClick()) ); + // Create LSP provider for gathering code definitions + const lspProvider = new LSPProvider({ + maxDefinitions: 50, + includeWorkspaceSymbols: true, + cacheTimeout: 30000, // 30 seconds + relevanceThreshold: 0.3 + }); + // Register the inline completion provider for all languages. - const inlineCompletionProvider = new TabCoderInlineCompletionProvider(profileService, statusBarProvider); + const inlineCompletionProvider = new TabCoderInlineCompletionProvider(profileService, statusBarProvider, lspProvider); context.subscriptions.push( vscode.languages.registerInlineCompletionItemProvider( '*', // Register for all languages diff --git a/src/services/lspProvider.ts b/src/services/lspProvider.ts new file mode 100644 index 0000000..83eab2e --- /dev/null +++ b/src/services/lspProvider.ts @@ -0,0 +1,332 @@ +import * as vscode from 'vscode'; +import { LSPDefinition, LSPContext, LSPProviderOptions } from '../types/lsp'; +import { logger } from '../utils/logger'; + +export class LSPProvider { + private cache = new Map(); + private pendingRequests = new Map>(); + private options: Required; + + constructor(options: LSPProviderOptions = {}) { + this.options = { + maxDefinitions: options.maxDefinitions ?? 50, + includeWorkspaceSymbols: options.includeWorkspaceSymbols ?? true, + cacheTimeout: options.cacheTimeout ?? 30000, // 30 seconds + relevanceThreshold: options.relevanceThreshold ?? 0.3 + }; + } + + /** + * Get LSP context for a document. This method is optimized for performance: + * - Returns cached data immediately if available and fresh + * - Starts async update in background if cache is stale + * - Returns empty context if no cache exists (async update will populate it) + */ + public async getLSPContext(document: vscode.TextDocument, position: vscode.Position): Promise { + const cacheKey = this.getCacheKey(document.uri); + const cached = this.cache.get(cacheKey); + const now = Date.now(); + + // Return cached data if fresh + if (cached && (now - cached.lastUpdated) < this.options.cacheTimeout) { + return cached; + } + + // Check if we already have a pending request for this document + const pendingRequest = this.pendingRequests.get(cacheKey); + if (pendingRequest) { + // Return cached data if available, otherwise wait for pending request + return cached || await pendingRequest; + } + + // Start async update + const updatePromise = this.updateLSPContext(document, position); + this.pendingRequests.set(cacheKey, updatePromise); + + try { + const result = await updatePromise; + this.cache.set(cacheKey, result); + return result; + } catch (error) { + logger.error('Failed to update LSP context:', error); + // Return cached data if available, otherwise empty context + return cached || this.createEmptyContext(); + } finally { + this.pendingRequests.delete(cacheKey); + } + } + + /** + * Get LSP context synchronously from cache only (for performance-critical paths) + */ + public getLSPContextSync(document: vscode.TextDocument): LSPContext | null { + const cacheKey = this.getCacheKey(document.uri); + const cached = this.cache.get(cacheKey); + const now = Date.now(); + + if (cached && (now - cached.lastUpdated) < this.options.cacheTimeout) { + return cached; + } + + return null; + } + + /** + * Preload LSP context for a document (fire and forget) + */ + public preloadLSPContext(document: vscode.TextDocument, position: vscode.Position): void { + const cacheKey = this.getCacheKey(document.uri); + + // Don't preload if we already have fresh data or a pending request + if (this.cache.has(cacheKey) || this.pendingRequests.has(cacheKey)) { + return; + } + + // Start async update without waiting + this.getLSPContext(document, position).catch(error => { + logger.debug('Preload LSP context failed:', error); + }); + } + + private async updateLSPContext(document: vscode.TextDocument, position: vscode.Position): Promise { + const definitions: LSPDefinition[] = []; + + try { + // Get document symbols (functions, classes, etc. in current file) + const documentSymbols = await this.getDocumentSymbols(document); + definitions.push(...documentSymbols); + + // Get workspace symbols if enabled (symbols from other files) + if (this.options.includeWorkspaceSymbols) { + const workspaceSymbols = await this.getRelevantWorkspaceSymbols(document, position); + definitions.push(...workspaceSymbols); + } + + // Get definitions at current position (go-to-definition results) + const positionDefinitions = await this.getDefinitionsAtPosition(document, position); + definitions.push(...positionDefinitions); + + // Sort by relevance and limit results + const sortedDefinitions = this.sortDefinitionsByRelevance(definitions, document, position); + const limitedDefinitions = sortedDefinitions.slice(0, this.options.maxDefinitions); + + return { + definitions: limitedDefinitions, + lastUpdated: Date.now(), + workspaceSymbols: definitions.filter(d => d.uri !== document.uri.toString()) + }; + } catch (error) { + logger.error('Error updating LSP context:', error); + return this.createEmptyContext(); + } + } + + private async getDocumentSymbols(document: vscode.TextDocument): Promise { + try { + const symbols = await vscode.commands.executeCommand( + 'vscode.executeDocumentSymbolProvider', + document.uri + ); + + if (!symbols) {return [];} + + return this.flattenDocumentSymbols(symbols, document.uri.toString()); + } catch (error) { + logger.debug('Failed to get document symbols:', error); + return []; + } + } + + private async getRelevantWorkspaceSymbols(document: vscode.TextDocument, position: vscode.Position): Promise { + try { + // Get text around cursor to find relevant symbols + const line = document.lineAt(position.line); + const wordRange = document.getWordRangeAtPosition(position); + const currentWord = wordRange ? document.getText(wordRange) : ''; + + // Search for symbols that might be relevant + const query = currentWord || line.text.trim().split(/\s+/).pop() || ''; + + if (!query || query.length < 2) {return [];} + + const symbols = await vscode.commands.executeCommand( + 'vscode.executeWorkspaceSymbolProvider', + query + ); + + if (!symbols) {return [];} + + return symbols + .filter(symbol => symbol.location.uri.toString() !== document.uri.toString()) + .slice(0, 20) // Limit workspace symbols + .map(symbol => this.convertSymbolInformation(symbol)); + } catch (error) { + logger.debug('Failed to get workspace symbols:', error); + return []; + } + } + + private async getDefinitionsAtPosition(document: vscode.TextDocument, position: vscode.Position): Promise { + try { + const definitions = await vscode.commands.executeCommand( + 'vscode.executeDefinitionProvider', + document.uri, + position + ); + + if (!definitions) {return [];} + + const results: LSPDefinition[] = []; + for (const definition of definitions) { + const defDocument = await vscode.workspace.openTextDocument(definition.uri); + const symbol = await this.getSymbolAtPosition(defDocument, definition.range.start); + if (symbol) { + results.push(symbol); + } + } + + return results; + } catch (error) { + logger.debug('Failed to get definitions at position:', error); + return []; + } + } + + private async getSymbolAtPosition(document: vscode.TextDocument, position: vscode.Position): Promise { + try { + const symbols = await vscode.commands.executeCommand( + 'vscode.executeDocumentSymbolProvider', + document.uri + ); + + if (!symbols) {return null;} + + const symbol = this.findSymbolAtPosition(symbols, position); + if (!symbol) {return null;} + + return this.convertDocumentSymbol(symbol, document.uri.toString()); + } catch (error) { + logger.debug('Failed to get symbol at position:', error); + return null; + } + } + + private flattenDocumentSymbols(symbols: vscode.DocumentSymbol[], uri: string): LSPDefinition[] { + const result: LSPDefinition[] = []; + + for (const symbol of symbols) { + result.push(this.convertDocumentSymbol(symbol, uri)); + + // Recursively add children + if (symbol.children && symbol.children.length > 0) { + result.push(...this.flattenDocumentSymbols(symbol.children, uri)); + } + } + + return result; + } + + private convertDocumentSymbol(symbol: vscode.DocumentSymbol, uri: string): LSPDefinition { + return { + name: symbol.name, + kind: symbol.kind, + detail: symbol.detail, + range: symbol.range, + uri + }; + } + + private convertSymbolInformation(symbol: vscode.SymbolInformation): LSPDefinition { + return { + name: symbol.name, + kind: symbol.kind, + containerName: symbol.containerName, + range: symbol.location.range, + uri: symbol.location.uri.toString() + }; + } + + private findSymbolAtPosition(symbols: vscode.DocumentSymbol[], position: vscode.Position): vscode.DocumentSymbol | null { + for (const symbol of symbols) { + if (symbol.range.contains(position)) { + // Check children first (more specific) + if (symbol.children) { + const childSymbol = this.findSymbolAtPosition(symbol.children, position); + if (childSymbol) {return childSymbol;} + } + return symbol; + } + } + return null; + } + + private sortDefinitionsByRelevance(definitions: LSPDefinition[], document: vscode.TextDocument, position: vscode.Position): LSPDefinition[] { + const currentUri = document.uri.toString(); + + return definitions.sort((a, b) => { + let scoreA = 0; + let scoreB = 0; + + // Prefer symbols from the same file + if (a.uri === currentUri) {scoreA += 10;} + if (b.uri === currentUri) {scoreB += 10;} + + // Prefer functions and classes over variables + if ([vscode.SymbolKind.Function, vscode.SymbolKind.Method, vscode.SymbolKind.Class, vscode.SymbolKind.Interface].includes(a.kind)) { + scoreA += 5; + } + if ([vscode.SymbolKind.Function, vscode.SymbolKind.Method, vscode.SymbolKind.Class, vscode.SymbolKind.Interface].includes(b.kind)) { + scoreB += 5; + } + + // Prefer symbols closer to current position (same file only) + if (a.uri === currentUri && b.uri === currentUri) { + const distanceA = Math.abs(a.range.start.line - position.line); + const distanceB = Math.abs(b.range.start.line - position.line); + scoreA += Math.max(0, 100 - distanceA); + scoreB += Math.max(0, 100 - distanceB); + } + + return scoreB - scoreA; + }); + } + + private getCacheKey(uri: vscode.Uri): string { + return uri.toString(); + } + + private createEmptyContext(): LSPContext { + return { + definitions: [], + lastUpdated: Date.now(), + workspaceSymbols: [] + }; + } + + /** + * Clear cache for a specific document + */ + public clearCache(uri: vscode.Uri): void { + const cacheKey = this.getCacheKey(uri); + this.cache.delete(cacheKey); + this.pendingRequests.delete(cacheKey); + } + + /** + * Clear all cache + */ + public clearAllCache(): void { + this.cache.clear(); + this.pendingRequests.clear(); + } + + /** + * Get cache statistics for debugging + */ + public getCacheStats(): { size: number; pendingRequests: number } { + return { + size: this.cache.size, + pendingRequests: this.pendingRequests.size + }; + } +} \ No newline at end of file diff --git a/src/types/index.ts b/src/types/index.ts index d90cc54..ea2d0c6 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -1,4 +1,5 @@ export * from './configuration'; export * from './model'; export * from './profile'; -export * from './provider'; \ No newline at end of file +export * from './provider'; +export * from './lsp'; \ No newline at end of file diff --git a/src/types/lsp.ts b/src/types/lsp.ts new file mode 100644 index 0000000..0422e87 --- /dev/null +++ b/src/types/lsp.ts @@ -0,0 +1,24 @@ +import * as vscode from 'vscode'; + +export interface LSPDefinition { + name: string; + kind: vscode.SymbolKind; + detail?: string; + documentation?: string; + range: vscode.Range; + uri: string; + containerName?: string; +} + +export interface LSPContext { + definitions: LSPDefinition[]; + lastUpdated: number; + workspaceSymbols: LSPDefinition[]; +} + +export interface LSPProviderOptions { + maxDefinitions?: number; + includeWorkspaceSymbols?: boolean; + cacheTimeout?: number; + relevanceThreshold?: number; +} \ No newline at end of file diff --git a/src/vscode/completionProvider.ts b/src/vscode/completionProvider.ts index c8933d2..1a90012 100644 --- a/src/vscode/completionProvider.ts +++ b/src/vscode/completionProvider.ts @@ -7,10 +7,12 @@ import { TabCoderStatusBarProvider } from './statusBarProvider'; import { logger } from '../utils/logger'; import { LanguageModelV2 } from '@ai-sdk/provider'; import { ProfileWithAPIKey } from '../types'; +import { LSPProvider } from '../services/lspProvider'; export class TabCoderInlineCompletionProvider implements vscode.InlineCompletionItemProvider { private profileService: ProfileService; private statusBarProvider: TabCoderStatusBarProvider; + private lspProvider: LSPProvider; private debounceTimeout: NodeJS.Timeout | undefined; private currentAbortController: AbortController | undefined; @@ -24,11 +26,15 @@ export class TabCoderInlineCompletionProvider implements vscode.InlineCompletion private lastChangeTimestamp: number = 0; - constructor(profileService: ProfileService, statusBarProvider: TabCoderStatusBarProvider) { + constructor(profileService: ProfileService, statusBarProvider: TabCoderStatusBarProvider, lspProvider: LSPProvider) { this.profileService = profileService; this.profileService.onDidActiveProfileChange(this.handleProfileChange, this); this.statusBarProvider = statusBarProvider; + this.lspProvider = lspProvider; + + // Set up document change listeners to preload LSP context + this.setupLSPPreloading(); } async provideInlineCompletionItems( @@ -224,12 +230,20 @@ export class TabCoderInlineCompletionProvider implements vscode.InlineCompletion new vscode.Position(document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length) )); + // Get LSP context (sync first, then async if needed) + let lspContext = this.lspProvider.getLSPContextSync(document); + if (!lspContext) { + // Start async LSP context gathering (fire and forget for this request) + this.lspProvider.preloadLSPContext(document, position); + } + const params: AutoCompleteContext = { textBeforeCursor, textAfterCursor, filename: document.fileName, language: document.languageId, currentLineText: document.lineAt(position.line).text, + lspContext: lspContext || undefined, }; logger.info(`Request ${requestId}: Using profile: ${profile.name} (ID: ${profile.id})`); @@ -335,4 +349,40 @@ export class TabCoderInlineCompletionProvider implements vscode.InlineCompletion this.lastUsedProfileId = undefined; logger.info('Active profile changed, clearing cached model'); } + + private setupLSPPreloading(): void { + // Preload LSP context when documents are opened or changed + vscode.workspace.onDidOpenTextDocument((document) => { + if (document.uri.scheme === 'file') { + // Preload LSP context for newly opened documents + const position = new vscode.Position(0, 0); + this.lspProvider.preloadLSPContext(document, position); + } + }); + + vscode.workspace.onDidChangeTextDocument((event) => { + if (event.document.uri.scheme === 'file' && event.contentChanges.length > 0) { + // Clear cache when document changes significantly + if (event.contentChanges.some(change => change.text.includes('\n') || change.rangeLength > 10)) { + this.lspProvider.clearCache(event.document.uri); + } + } + }); + + vscode.workspace.onDidSaveTextDocument((document) => { + if (document.uri.scheme === 'file') { + // Clear cache and preload fresh LSP context after save + this.lspProvider.clearCache(document.uri); + const position = new vscode.Position(0, 0); + this.lspProvider.preloadLSPContext(document, position); + } + }); + } + + /** + * Get LSP provider instance for external access + */ + public getLSPProvider(): LSPProvider { + return this.lspProvider; + } } \ No newline at end of file