fix: stopped server would just restart

- stopped server would just restart if something needed it
- this adds explicit flag which indicates that it was stopped by a user

Author: benfdking

vscode/extension/src/commands/stop.ts

@@ -11,7 +11,7 @@ export const stop = (lspClient: LSPClient | undefined) => {
       return
     }
 
-    await lspClient.stop()
+    await lspClient.stop(true)
     await window.showInformationMessage('LSP server stopped')
     traceInfo('LSP server stopped successfully')
   }

vscode/extension/src/extension.ts

@@ -115,10 +115,10 @@ export async function activate(context: vscode.ExtensionContext) {
     }),
   )
 
-  const restart = async () => {
+  const restart = async (invokedByUser = false) => {
     if (lspClient) {
       traceVerbose('Restarting LSP client')
-      const restartResult = await lspClient.restart()
+      const restartResult = await lspClient.restart(invokedByUser)
       if (isErr(restartResult)) {
         return handleError(
           authProvider,
@@ -130,7 +130,7 @@ export async function activate(context: vscode.ExtensionContext) {
       context.subscriptions.push(lspClient)
     } else {
       lspClient = new LSPClient()
-      const result = await lspClient.start()
+      const result = await lspClient.start(invokedByUser)
       if (isErr(result)) {
         return handleError(
           authProvider,
@@ -159,7 +159,7 @@ export async function activate(context: vscode.ExtensionContext) {
       await restart()
     }),
     registerCommand(`sqlmesh.restart`, async () => {
-      await restart()
+      await restart(true)
     }),
     registerCommand(`sqlmesh.stop`, stop(lspClient)),
     registerCommand(`sqlmesh.printEnvironment`, printEnvironment()),

vscode/extension/src/lsp/lsp.ts

@@ -36,6 +36,13 @@ export class LSPClient implements Disposable {
    */
   private supportedMethodsState: SupportedMethodsState = { type: 'not-fetched' }
 
+  /**
+   * Explicitly stopped remembers whether the LSP client has been explicitly stopped
+   * by the user. This is used to prevent the client from being restarted unless the user
+   * explicitly calls the `restart` method.
+   */
+  private explicitlyStopped = false
+
   constructor() {
     this.client = undefined
   }
@@ -54,7 +61,15 @@ export class LSPClient implements Disposable {
     return completion !== undefined
   }
 
-  public async start(): Promise<Result<undefined, ErrorType>> {
+  public async start(
+    overrideStoppedByUser = false,
+  ): Promise<Result<undefined, ErrorType>> {
+    if (this.explicitlyStopped && !overrideStoppedByUser) {
+      traceInfo(
+        'LSP client has been explicitly stopped by user, not starting again.',
+      )
+      return ok(undefined)
+    }
     if (!outputChannel) {
       outputChannel = window.createOutputChannel('sqlmesh-lsp')
     }
@@ -124,18 +139,24 @@ export class LSPClient implements Disposable {
     return ok(undefined)
   }
 
-  public async restart(): Promise<Result<undefined, ErrorType>> {
+  public async restart(
+    overrideByUser = false,
+  ): Promise<Result<undefined, ErrorType>> {
     await this.stop()
-    return await this.start()
+    return await this.start(overrideByUser)
   }
 
-  public async stop() {
+  public async stop(stoppedByUser = false): Promise<void> {
     if (this.client) {
       await this.client.stop()
       this.client = undefined
       // Reset supported methods state when the client stops
       this.supportedMethodsState = { type: 'not-fetched' }
     }
+    if (stoppedByUser) {
+      this.explicitlyStopped = true
+      traceInfo('SQLMesh LSP client stopped by user.')
+    }
   }
 
   public async dispose() {

vscode/extension/tests/stop.spec.ts

@@ -18,7 +18,6 @@ test('Stop server works', async ({ page, sharedCodeServer }) => {
 
   // Navigate to code-server instance
   await openServerPage(page, tempDir, sharedCodeServer)
-  await page.waitForSelector('[role="application"]', { timeout: 10000 })
 
   // Wait for the models folder to be visible in the file explorer
   await page.waitForSelector('text=models')
@@ -52,3 +51,84 @@ test('Stop server works', async ({ page, sharedCodeServer }) => {
     'text="Failed to render model: LSP client not ready."',
   )
 })
+
+test('Stopped server only restarts when explicitly requested', async ({
+  page,
+  sharedCodeServer,
+}) => {
+  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'))
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+
+  await createPythonInterpreterSettingsSpecifier(tempDir)
+
+  // Navigate to code-server instance
+  await openServerPage(page, tempDir, sharedCodeServer)
+
+  // Wait for the models folder to be visible in the file explorer
+  await page.waitForSelector('text=models')
+
+  // Click on the models folder, excluding external_models
+  await page
+    .getByRole('treeitem', { name: 'models', exact: true })
+    .locator('a')
+    .click()
+
+  // Open the customers.sql model
+  await page
+    .getByRole('treeitem', { name: 'marketing.sql', exact: true })
+    .locator('a')
+    .click()
+  await page.waitForSelector('text=grain')
+  await waitForLoadedSQLMesh(page)
+
+  // Click on sushi.raw_marketing
+  await page.getByText('sushi.raw_marketing;').click()
+
+  // Open the preview hover
+  await runCommand(page, 'Show Definition Preview Hover')
+
+  // Assert that the hover is visible with text "Table of marketing status."
+  await page.waitForSelector('text=Table of marketing status.', {
+    timeout: 5_000,
+    state: 'visible',
+  })
+
+  // Hit Esc to close the hover
+  await page.keyboard.press('Escape')
+
+  // Assert that the hover is no longer visible
+  await page.waitForSelector('text=Table of marketing status.', {
+    timeout: 5_000,
+    state: 'hidden',
+  })
+
+  // Stop the server
+  await runCommand(page, 'SQLMesh: Stop Server')
+
+  // Await LSP server stopped message
+  await page.waitForSelector('text=LSP server stopped')
+
+  // Open the preview hover again
+  await runCommand(page, 'Show Definition Preview Hover')
+
+  // Assert that the hover is not visible
+  await page.waitForSelector('text=Table of marketing status.', {
+    timeout: 5_000,
+    state: 'hidden',
+  })
+
+  // Restart the server explicitly
+  await runCommand(page, 'SQLMesh: Restart Server')
+
+  // Await LSP server started message
+  await waitForLoadedSQLMesh(page)
+
+  // Open the preview hover again
+  await runCommand(page, 'Show Definition Preview Hover')
+
+  // Assert that the hover is visible with text "Table of marketing status."
+  await page.waitForSelector('text=Table of marketing status.', {
+    timeout: 5_000,
+    state: 'visible',
+  })
+})