chore(lsp): move autocomplete to the lsp

Author: benfdking

sqlmesh/lsp/main.py

@@ -434,6 +434,50 @@ def workspace_diagnostic(
                 )
                 return types.WorkspaceDiagnosticReport(items=[])
 
+        @self.server.feature(types.TEXT_DOCUMENT_COMPLETION)
+        def completion(
+            ls: LanguageServer, params: types.CompletionParams
+        ) -> t.Optional[types.CompletionList]:
+            """Handle completion requests from the client."""
+            try:
+                uri = URI(params.text_document.uri)
+                context = self._context_get_or_load(uri)
+                
+                # Get completions using the existing completions module
+                completion_response = get_sql_completions(context, uri)
+                
+                # Convert to LSP completion items
+                completion_items = []
+                
+                # Add model completions
+                for model in completion_response.models:
+                    completion_items.append(
+                        types.CompletionItem(
+                            label=model,
+                            kind=types.CompletionItemKind.Reference,
+                            detail="SQLMesh Model",
+                        )
+                    )
+                
+                # Add keyword completions
+                for keyword in completion_response.keywords:
+                    completion_items.append(
+                        types.CompletionItem(
+                            label=keyword,
+                            kind=types.CompletionItemKind.Keyword,
+                            detail="SQL Keyword",
+                        )
+                    )
+                
+                return types.CompletionList(
+                    is_incomplete=False,
+                    items=completion_items,
+                )
+                
+            except Exception as e:
+                ls.show_message(f"Error getting completions: {e}", types.MessageType.Error)
+                return None
+
     def _get_diagnostics_for_uri(self, uri: URI) -> t.Tuple[t.List[types.Diagnostic], int]:
         """Get diagnostics for a specific URI, returning (diagnostics, result_id).
 

vscode/extension/src/extension.ts

@@ -83,13 +83,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  context.subscriptions.push(
-    vscode.languages.registerCompletionItemProvider(
-      selector,
-      completionProvider(lspClient),
-    ),
-  )
-
   // Register the webview
   const lineagePanel = new LineagePanel(context.extensionUri, lspClient)
   context.subscriptions.push(
@@ -165,6 +158,15 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(lspClient)
   }
 
+  if (lspClient && !lspClient.hasCompletionCapability()) {
+    context.subscriptions.push(
+      vscode.languages.registerCompletionItemProvider(
+        selector,
+        completionProvider(lspClient),
+      ),
+    )
+  }
+
   traceInfo('Extension activated')
 }
 

vscode/extension/src/lsp/lsp.ts

@@ -21,6 +21,16 @@ export class LSPClient implements Disposable {
     this.client = undefined
   }
 
+  public hasCompletionCapability(): boolean {
+    if (!this.client) {
+      traceError('LSP client is not initialized')
+      return false
+    }
+    const capabilities = this.client.initializeResult?.capabilities
+    const completion = capabilities?.completionProvider
+    return completion !== undefined
+  }
+
   public async start(): Promise<Result<undefined, ErrorType>> {
     if (!outputChannel) {
       outputChannel = window.createOutputChannel('sqlmesh-lsp')

vscode/extension/tests/completions.spec.ts

@@ -44,7 +44,7 @@ test('Autocomplete for model names', async () => {
     await window.waitForTimeout(500)
 
     // Check if the autocomplete suggestion for sushi.customers is visible
-    await expect(window.locator('text=sushi.customers')).toBeVisible()
+    expect(await window.locator('text=sushi.customers').count()).toBe(1)
 
     await close()
   } finally {