fix(vscode): improve registering of lsp in extension

Author: benfdking

vscode/extension/src/commands/renderModel.ts

@@ -6,13 +6,16 @@ import { RenderedModelProvider } from '../providers/renderedModelProvider'
 
 export async function reRenderModelForSourceFile(
   sourceUri: string,
-  lspClient: LSPClient,
+  lspClient: LSPClient | undefined,
   renderedModelProvider: RenderedModelProvider,
 ): Promise<void> {
   const renderedUri = renderedModelProvider.getRenderedUriForSource(sourceUri)
   if (!renderedUri) {
     return // No rendered model exists for this source file
   }
+  if (!lspClient) {
+    return
+  }
 
   // Call the render model API
   const result = await lspClient.call_custom_method('sqlmesh/render_model', {

vscode/extension/src/extension.ts

@@ -1,35 +1,51 @@
-import { format } from './commands/format'
+/**********************************************************************
+ *  Extension entry point                                              *
+ *********************************************************************/
+
 import * as vscode from 'vscode'
-import {
-  createOutputChannel,
-  onDidChangeConfiguration,
-  registerCommand,
-} from './utilities/common/vscodeapi'
-import { registerLogger, traceInfo, traceVerbose } from './utilities/common/log'
-import { onDidChangePythonInterpreter } from './utilities/common/python'
-import { LSPClient } from './lsp/lsp'
-import { AuthenticationProviderTobikoCloud } from './auth/auth'
+
+import { format } from './commands/format'
 import { signOut } from './commands/signout'
 import { signIn } from './commands/signin'
 import { signInSpecifyFlow } from './commands/signinSpecifyFlow'
 import { renderModel, reRenderModelForSourceFile } from './commands/renderModel'
 import { stop } from './commands/stop'
 import { printEnvironment } from './commands/printEnvironment'
-import { isErr } from '@bus/result'
+
+import {
+  createOutputChannel,
+  onDidChangeConfiguration,
+  registerCommand,
+} from './utilities/common/vscodeapi'
+import {
+  registerLogger,
+  traceInfo,
+  traceVerbose,
+  traceError,
+} from './utilities/common/log'
+import { onDidChangePythonInterpreter } from './utilities/common/python'
+import { sleep } from './utilities/sleep'
 import { handleError } from './utilities/errors'
+
 import { selector, completionProvider } from './completion/completion'
 import { LineagePanel } from './webviews/lineagePanel'
 import { RenderedModelProvider } from './providers/renderedModelProvider'
-import { sleep } from './utilities/sleep'
+
 import {
   controller as testController,
   setupTestController,
 } from './tests/tests'
 
+import { isErr } from '@bus/result'
+import { AuthenticationProviderTobikoCloud } from './auth/auth'
+import { LSPClient } from './lsp/lsp'
+
+/** Singleton LSP client for the extension. */
 let lspClient: LSPClient | undefined
 
-// This method is called when your extension is activated
-// Your extension is activated the very first time the command is executed
+/** Handle to the (single) test controller disposable so we can replace it on restart. */
+let testControllerDisposable: vscode.Disposable | undefined
+
 export async function activate(context: vscode.ExtensionContext) {
   const extensionOutputChannel = createOutputChannel('sqlmesh')
   context.subscriptions.push(
@@ -38,7 +54,6 @@ export async function activate(context: vscode.ExtensionContext) {
   )
   traceInfo('Activating SQLMesh extension')
 
-  traceInfo('Registering authentication provider')
   const authProvider = new AuthenticationProviderTobikoCloud()
   context.subscriptions.push(
     vscode.authentication.registerAuthenticationProvider(
@@ -48,33 +63,70 @@ export async function activate(context: vscode.ExtensionContext) {
       { supportsMultipleAccounts: false },
     ),
   )
-  traceInfo('Authentication provider registered')
 
+  const restartLsp = async (invokedByUser = false): Promise<void> => {
+    if (!lspClient) {
+      lspClient = new LSPClient()
+    }
+
+    traceVerbose('Restarting SQLMesh LSP client')
+    const result = await lspClient.restart(invokedByUser)
+    if (isErr(result)) {
+      await handleError(
+        authProvider,
+        restartLsp,
+        result.error,
+        'LSP restart failed',
+      )
+      return
+    }
+
+    // push once to avoid duplicate disposables on multiple restarts
+    if (!context.subscriptions.includes(lspClient)) {
+      context.subscriptions.push(lspClient)
+    }
+
+    /* Replace the test controller each time we restart the client */
+    if (testControllerDisposable) {
+      testControllerDisposable.dispose()
+    }
+    testControllerDisposable = setupTestController(lspClient)
+    context.subscriptions.push(testControllerDisposable)
+  }
+
+  // commands needing the restart helper
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'sqlmesh.signin',
-      signIn(authProvider, async () => {
-        traceInfo('Restarting LSP after sign-in')
-        await restart()
-      }),
+      signIn(authProvider, () => restartLsp()),
     ),
-  )
-  context.subscriptions.push(
     vscode.commands.registerCommand(
       'sqlmesh.signinSpecifyFlow',
-      signInSpecifyFlow(authProvider, async () => {
-        traceInfo('Restarting LSP after sign-in')
-        await restart()
-      }),
+      signInSpecifyFlow(authProvider, () => restartLsp()),
     ),
-  )
-  context.subscriptions.push(
     vscode.commands.registerCommand('sqlmesh.signout', signOut(authProvider)),
   )
 
+  //  Instantiate the LSP client (once)
   lspClient = new LSPClient()
+  const startResult = await lspClient.start()
+  if (isErr(startResult)) {
+    await handleError(
+      authProvider,
+      restartLsp,
+      startResult.error,
+      'Failed to start LSP',
+    )
+    return // abort activation – nothing else to do
+  }
+
+  context.subscriptions.push(lspClient)
 
-  // Create and register the rendered model provider
+  // Initialize the test controller
+  testControllerDisposable = setupTestController(lspClient)
+  context.subscriptions.push(testControllerDisposable, testController)
+
+  // Register the rendered model provider
   const renderedModelProvider = new RenderedModelProvider()
   context.subscriptions.push(
     vscode.workspace.registerTextDocumentContentProvider(
@@ -91,7 +143,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  // Register the webview
   const lineagePanel = new LineagePanel(context.extensionUri, lspClient)
   context.subscriptions.push(
     vscode.window.registerWebviewViewProvider(
@@ -100,11 +151,10 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  // Add file save listener for auto-rerendering models
+  // Re‑render model automatically when its source file is saved
   context.subscriptions.push(
     vscode.workspace.onDidSaveTextDocument(async document => {
       if (
-        lspClient &&
         renderedModelProvider.hasRenderedModelForSource(
           document.uri.toString(true),
         )
@@ -119,73 +169,23 @@ export async function activate(context: vscode.ExtensionContext) {
     }),
   )
 
-  const restart = async (invokedByUser = false) => {
-    if (lspClient) {
-      traceVerbose('Restarting LSP client')
-      const restartResult = await lspClient.restart(invokedByUser)
-      if (isErr(restartResult)) {
-        return handleError(
-          authProvider,
-          restart,
-          restartResult.error,
-          'LSP restart failed',
-        )
-      }
-      context.subscriptions.push(lspClient)
-      context.subscriptions.push(setupTestController(lspClient))
-    } else {
-      lspClient = new LSPClient()
-      const result = await lspClient.start(invokedByUser)
-      if (isErr(result)) {
-        return handleError(
-          authProvider,
-          restart,
-          result.error,
-          'Failed to start LSP',
-        )
-      } else {
-        context.subscriptions.push(lspClient)
-        context.subscriptions.push(setupTestController(lspClient))
-      }
-    }
-  }
-
+  // miscellaneous commands
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'sqlmesh.format',
-      format(authProvider, lspClient, restart),
+      format(authProvider, lspClient, restartLsp),
     ),
+    registerCommand('sqlmesh.restart', () => restartLsp(true)),
+    registerCommand('sqlmesh.stop', stop(lspClient)),
+    registerCommand('sqlmesh.printEnvironment', printEnvironment()),
   )
 
   context.subscriptions.push(
-    onDidChangePythonInterpreter(async () => {
-      await restart()
-    }),
-    onDidChangeConfiguration(async () => {
-      await restart()
-    }),
-    registerCommand(`sqlmesh.restart`, async () => {
-      await restart(true)
-    }),
-    registerCommand(`sqlmesh.stop`, stop(lspClient)),
-    registerCommand(`sqlmesh.printEnvironment`, printEnvironment()),
+    onDidChangePythonInterpreter(() => restartLsp()),
+    onDidChangeConfiguration(() => restartLsp()),
   )
 
-  const result = await lspClient.start()
-  if (isErr(result)) {
-    return handleError(
-      authProvider,
-      restart,
-      result.error,
-      'Failed to start LSP',
-    )
-  } else {
-    context.subscriptions.push(lspClient)
-    context.subscriptions.push(setupTestController(lspClient))
-    context.subscriptions.push(testController)
-  }
-
-  if (lspClient && !lspClient.hasCompletionCapability()) {
+  if (!lspClient.hasCompletionCapability()) {
     context.subscriptions.push(
       vscode.languages.registerCompletionItemProvider(
         selector,
@@ -194,12 +194,19 @@ export async function activate(context: vscode.ExtensionContext) {
     )
   }
 
-  traceInfo('Extension activated')
+  traceInfo('SQLMesh extension activated')
 }
 
 // This method is called when your extension is deactivated
 export async function deactivate() {
-  if (lspClient) {
-    await lspClient.dispose()
+  try {
+    if (testControllerDisposable) {
+      testControllerDisposable.dispose()
+    }
+    if (lspClient) {
+      await lspClient.dispose()
+    }
+  } catch (e) {
+    traceError(`Error during deactivate: ${e}`)
   }
 }