fix(vscode): stop the LSP restarting on unrelated configuration changes

vscode/extension/src/extension.ts

@@ -24,8 +24,10 @@ import {
   traceError,
 } from './utilities/common/log'
 import { onDidChangePythonInterpreter } from './utilities/common/python'
+import { requiresLspRestart } from './utilities/common/configurationChange'
+import { coalesceAsync } from './utilities/coalesceAsync'
 import { sleep } from './utilities/sleep'
-import { handleError } from './utilities/errors'
+import { ErrorType, handleError } from './utilities/errors'
 
 import { selector, completionProvider } from './completion/completion'
 import { LineagePanel } from './webviews/lineagePanel'
@@ -65,20 +67,32 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  const restartLsp = async (invokedByUser = false): Promise<void> => {
+  /**
+   * Set when a restart was asked for explicitly, so that a user-invoked restart
+   * coalesced together with an automatic one is still treated as user-invoked.
+   */
+  let restartInvokedByUser = false
+
+  /**
+   * Set by a failed run and handled by the caller once the run has finished.
+   * Handling it inside the run would deadlock: the not_signed_in handler waits
+   * on a sign-in that restarts the client again, and that restart would wait on
+   * the run that is still waiting on the handler.
+   */
+  let restartError: ErrorType | undefined
+
+  const runRestart = async (): Promise<void> => {
+    const invokedByUser = restartInvokedByUser
+    restartInvokedByUser = false
+
     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',
-      )
+      restartError = result.error
       return
     }
 
@@ -95,6 +109,25 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(testControllerDisposable)
   }
 
+  /**
+   * Restarts are serialized: a client disposing while the next one starts up
+   * leaves colliding command registrations and requests aimed at a disposed
+   * client, which is what surfaced as constant crashes.
+   */
+  const restartLspSerialized = coalesceAsync(runRestart)
+
+  const restartLsp = async (invokedByUser = false): Promise<void> => {
+    restartInvokedByUser = restartInvokedByUser || invokedByUser
+    await restartLspSerialized()
+
+    // Claimed so that callers coalesced into the same run don't each report it.
+    const error = restartError
+    restartError = undefined
+    if (error) {
+      await handleError(authProvider, restartLsp, error, 'LSP restart failed')
+    }
+  }
+
   // commands needing the restart helper
   context.subscriptions.push(
     vscode.commands.registerCommand(
@@ -191,7 +224,13 @@ export async function activate(context: vscode.ExtensionContext) {
 
   context.subscriptions.push(
     onDidChangePythonInterpreter(() => restartLsp()),
-    onDidChangeConfiguration(() => restartLsp()),
+    // Only restart for settings the server actually reads. This event fires for
+    // every setting in the editor, including ones written by other extensions.
+    onDidChangeConfiguration(event => {
+      if (requiresLspRestart(event)) {
+        void restartLsp()
+      }
+    }),
   )
 
   if (!lspClient.hasCompletionCapability()) {

vscode/extension/src/utilities/coalesceAsync.test.ts

@@ -0,0 +1,159 @@
+// SPDX-License-Identifier: Apache-2.0
+
+import { describe, expect, it } from 'vitest'
+import { coalesceAsync } from './coalesceAsync'
+
+/**
+ * Let pending microtasks run. A rerun is scheduled off the promise of the run
+ * that precedes it, so it starts a few hops after that run settles.
+ */
+const flush = () => new Promise<void>(resolve => setTimeout(resolve, 0))
+
+/** A task that only settles when the test releases it. */
+function deferredTask() {
+  let release: (() => void) | undefined
+  let rejectWith: ((error: Error) => void) | undefined
+  let calls = 0
+
+  const run = () => {
+    calls += 1
+    return new Promise<void>((resolve, reject) => {
+      release = resolve
+      rejectWith = reject
+    })
+  }
+
+  return {
+    run,
+    get calls() {
+      return calls
+    },
+    release: () => release?.(),
+    reject: (error: Error) => rejectWith?.(error),
+  }
+}
+
+describe('coalesceAsync', () => {
+  it('runs the task immediately when idle', async () => {
+    const task = deferredTask()
+    const run = coalesceAsync(task.run)
+
+    const first = run()
+    expect(task.calls).toBe(1)
+
+    task.release()
+    await first
+  })
+
+  // A burst of triggers used to start a restart per event, which left clients
+  // disposing and starting concurrently. See #5642.
+  it('collapses every call made while running into a single rerun', async () => {
+    const task = deferredTask()
+    const run = coalesceAsync(task.run)
+
+    const first = run()
+    expect(task.calls).toBe(1)
+
+    const queued = [run(), run(), run(), run()]
+    expect(task.calls).toBe(1)
+
+    task.release()
+    await first
+    await flush()
+
+    // Exactly one rerun is scheduled, no matter how many calls arrived.
+    expect(task.calls).toBe(2)
+
+    task.release()
+    await Promise.all(queued)
+    await flush()
+    expect(task.calls).toBe(2)
+  })
+
+  it('runs again for calls made after the previous run finished', async () => {
+    const task = deferredTask()
+    const run = coalesceAsync(task.run)
+
+    const first = run()
+    task.release()
+    await first
+    expect(task.calls).toBe(1)
+
+    const second = run()
+    expect(task.calls).toBe(2)
+    task.release()
+    await second
+  })
+
+  it('resolves the callers that were coalesced together', async () => {
+    const task = deferredTask()
+    const run = coalesceAsync(task.run)
+
+    const first = run()
+    const queued = [run(), run()]
+
+    task.release()
+    await first
+    // The rerun has to be in flight before it can be released.
+    await flush()
+    task.release()
+
+    await expect(Promise.all(queued)).resolves.toEqual([undefined, undefined])
+  })
+
+  it('surfaces a failure without wedging later calls', async () => {
+    const task = deferredTask()
+    const run = coalesceAsync(task.run)
+
+    const first = run()
+    task.reject(new Error('restart failed'))
+    await expect(first).rejects.toThrow('restart failed')
+
+    const second = run()
+    expect(task.calls).toBe(2)
+    task.release()
+    await second
+  })
+
+  // A failed restart is reported to the caller, which handles it after the run
+  // has finished. The not_signed_in handler signs the user in and restarts the
+  // client again, so that handler must not run inside the task: the restart it
+  // triggers would wait on the run that is still waiting on the handler, and
+  // neither would ever settle. See #5920.
+  it('lets a failure handler trigger another run without deadlocking', async () => {
+    let failNextRun = true
+    let reportedError: string | undefined
+    let runs = 0
+
+    const runRestart = async (): Promise<void> => {
+      runs += 1
+      // The real run awaits the client restart before it knows the outcome.
+      // Without that await the re-entrant call below lands in the synchronous
+      // prefix of the run, before it has been recorded as in flight, and the
+      // deadlock this covers cannot happen.
+      const failed = await Promise.resolve(failNextRun)
+      if (failed) {
+        failNextRun = false
+        reportedError = 'not_signed_in'
+      }
+    }
+
+    const runRestartSerialized = coalesceAsync(runRestart)
+
+    const restart = async (): Promise<void> => {
+      await runRestartSerialized()
+
+      const error = reportedError
+      reportedError = undefined
+      if (error === 'not_signed_in') {
+        // Stands in for the sign-in flow, which restarts once signed in.
+        await restart()
+      }
+    }
+
+    await restart()
+
+    expect(runs).toBe(2)
+    expect(reportedError).toBeUndefined()
+  }, 2_000)
+})

vscode/extension/src/utilities/coalesceAsync.ts

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Serialize an async task so it never overlaps with itself, collapsing any
+ * calls that arrive while it is running into a single rerun.
+ *
+ * Restarting the language server means stopping a client and starting another.
+ * Letting two of those interleave leaves commands registered by the outgoing
+ * client colliding with the incoming one, and leaves requests addressed to a
+ * client that has already been disposed. Queueing one run per trigger would
+ * only spread the same problem out over time, so queued triggers collapse into
+ * one rerun: all the caller wants is for the task to have run after its
+ * request, not for it to run once per request.
+ *
+ * @param task The task to serialize.
+ * @returns A function that resolves once the task has run for that call.
+ */
+export function coalesceAsync(task: () => Promise<void>): () => Promise<void> {
+  let running: Promise<void> | undefined
+  let queued: Promise<void> | undefined
+
+  const start = async (): Promise<void> => {
+    try {
+      await task()
+    } finally {
+      running = undefined
+    }
+  }
+
+  return (): Promise<void> => {
+    if (!running) {
+      running = start()
+      return running
+    }
+
+    // A rerun is already scheduled, so this call is satisfied by that one.
+    if (!queued) {
+      queued = running
+        // A failed run must not stop the rerun that was asked for; the caller
+        // waiting on the failed run is the one that sees the error.
+        .catch(() => undefined)
+        .then(() => {
+          queued = undefined
+          running = start()
+          return running
+        })
+    }
+
+    return queued
+  }
+}

vscode/extension/src/utilities/common/configurationChange.test.ts

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+
+import { describe, expect, it } from 'vitest'
+import { requiresLspRestart } from './configurationChange'
+
+/**
+ * Build a stand-in for `vscode.ConfigurationChangeEvent` from the settings that
+ * changed. VS Code reports a section as affected when the changed key is the
+ * section itself or sits underneath it.
+ */
+const changed = (...keys: string[]) => ({
+  affectsConfiguration: (section: string) =>
+    keys.some(key => key === section || key.startsWith(`${section}.`)),
+})
+
+describe('requiresLspRestart', () => {
+  it('restarts when a sqlmesh setting changes', () => {
+    expect(requiresLspRestart(changed('sqlmesh.projectPaths'))).toBe(true)
+    expect(requiresLspRestart(changed('sqlmesh.lspEntrypoint'))).toBe(true)
+  })
+
+  it('restarts when the python interpreter changes', () => {
+    expect(requiresLspRestart(changed('python.defaultInterpreterPath'))).toBe(
+      true,
+    )
+  })
+
+  // The LSP used to restart on every configuration change in the editor, so
+  // anything that wrote a setting took the extension down with it. See #5920.
+  it('ignores settings the language server does not read', () => {
+    expect(requiresLspRestart(changed('editor.fontSize'))).toBe(false)
+    expect(requiresLspRestart(changed('workbench.colorTheme'))).toBe(false)
+    expect(requiresLspRestart(changed('files.autoSave'))).toBe(false)
+  })
+
+  // Running any python command in a VS Code terminal makes the Python
+  // extension touch its own terminal settings, which is what made the
+  // extension crash whenever sqlmesh was run in the terminal. See #5642.
+  it('ignores python settings unrelated to the interpreter', () => {
+    expect(
+      requiresLspRestart(changed('python.terminal.activateEnvironment')),
+    ).toBe(false)
+    expect(
+      requiresLspRestart(changed('python.analysis.typeCheckingMode')),
+    ).toBe(false)
+    expect(requiresLspRestart(changed('terminal.integrated.env.linux'))).toBe(
+      false,
+    )
+  })
+
+  it('does not restart when nothing relevant changed', () => {
+    expect(requiresLspRestart(changed())).toBe(false)
+  })
+
+  it('restarts when a relevant setting changes alongside irrelevant ones', () => {
+    expect(
+      requiresLspRestart(changed('editor.fontSize', 'sqlmesh.projectPaths')),
+    ).toBe(true)
+  })
+})