testing new lsp.ts

Author: benfdking

vscode/extension/src/lsp/lsp.ts

@@ -20,177 +20,183 @@ import { CustomLSPMethods } from './custom'
 type SupportedMethodsState =
   | { type: 'not-fetched' }
   | { type: 'fetched'; methods: Set<string> }
-  // TODO: This state is used when the `sqlmesh/supported_methods` endpoint is
-  //  not supported by the LSP server. This is in order to be backward compatible
-  //  with older versions of SQLMesh that do not support this endpoint. At some point
-  //  we should remove this state and always fetch the supported methods.
-  | { type: 'endpoint-not-supported' }
+  | { type: 'endpoint-not-supported' } // fallback for very old servers
 
 let outputChannel: OutputChannel | undefined
 
 export class LSPClient implements Disposable {
   private client: LanguageClient | undefined
-  /**
-   * State to track whether the supported methods have been fetched. These are used to determine if a method is supported
-   * by the LSP server and return an error if not.
-   */
+
+  /** Caches which custom methods the server supports */
   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.
+   * Remember whether the user explicitly stopped the client so that we do not
+   * auto‑start again until they ask for it.
    */
   private explicitlyStopped = false
 
-  constructor() {
-    this.client = undefined
+  // ───────────────────────────────────────────── helpers
+
+  /** True when a LanguageClient instance is alive. */
+  private get isRunning(): boolean {
+    return this.client !== undefined
   }
 
-  // TODO: This method is used to check if the LSP client has completion capability
-  //  in order to be backward compatible with older versions of SQLMesh that do not
-  //  support completion. At some point we should remove this method and always assume
-  //  that the LSP client has completion capability.
+  // ───────────────────────────────────────────── public api
+
+  /**
+   * Query whether the connected server advertises completion capability.
+   * (Transient helper kept for backwards‑compat reasons.)
+   */
   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
+    return (
+      this.client.initializeResult?.capabilities?.completionProvider !==
+      undefined
+    )
   }
 
+  /** Start the Language Client unless it is already running. */
   public async start(
     overrideStoppedByUser = false,
   ): Promise<Result<undefined, ErrorType>> {
+    // 0. Respect the “user stopped” flag
     if (this.explicitlyStopped && !overrideStoppedByUser) {
       traceInfo(
         'LSP client has been explicitly stopped by user, not starting again.',
       )
       return ok(undefined)
     }
+
+    // 1. Guard against duplicate initialisation
+    if (this.isRunning) {
+      traceInfo('LSP client already running – start() is a no‑op.')
+      return ok(undefined)
+    }
+
+    // 2. Ensure we have an output channel
     if (!outputChannel) {
       outputChannel = window.createOutputChannel('sqlmesh-lsp')
     }
 
+    // 3. Resolve sqlmesh executable
     const sqlmesh = await sqlmeshLspExec()
     if (isErr(sqlmesh)) {
       traceError(
         `Failed to get sqlmesh_lsp_exec, ${JSON.stringify(sqlmesh.error)}`,
       )
       return sqlmesh
     }
-    const workspaceFolders = getWorkspaceFolders()
-    if (workspaceFolders.length === 0) {
-      traceError(`No workspace folders found`)
-      return err({
-        type: 'generic',
-        message: 'No workspace folders found',
-      })
+
+    // 4. We need at least one workspace
+    if (getWorkspaceFolders().length === 0) {
+      const msg = 'No workspace folders found'
+      traceError(msg)
+      return err({ type: 'generic', message: msg })
     }
+
     const workspacePath = sqlmesh.value.workspacePath
     const serverOptions: ServerOptions = {
       run: {
         command: sqlmesh.value.bin,
         transport: TransportKind.stdio,
-        options: {
-          cwd: workspacePath,
-          env: sqlmesh.value.env,
-        },
+        options: { cwd: workspacePath, env: sqlmesh.value.env },
         args: sqlmesh.value.args,
       },
       debug: {
         command: sqlmesh.value.bin,
         transport: TransportKind.stdio,
-        options: {
-          cwd: workspacePath,
-          env: sqlmesh.value.env,
-        },
+        options: { cwd: workspacePath, env: sqlmesh.value.env },
         args: sqlmesh.value.args,
       },
     }
     const clientOptions: LanguageClientOptions = {
       documentSelector: [
-        { scheme: 'file', pattern: `**/*.sql` },
-        {
-          scheme: 'file',
-          pattern: '**/external_models.yaml',
-        },
-        {
-          scheme: 'file',
-          pattern: '**/external_models.yml',
-        },
+        { scheme: 'file', pattern: '**/*.sql' },
+        { scheme: 'file', pattern: '**/external_models.yaml' },
+        { scheme: 'file', pattern: '**/external_models.yml' },
       ],
       diagnosticCollectionName: 'sqlmesh',
-      outputChannel: outputChannel,
+      outputChannel,
     }
 
     traceInfo(
-      `Starting SQLMesh Language Server with workspace path: ${workspacePath} with server options ${JSON.stringify(serverOptions)} and client options ${JSON.stringify(clientOptions)}`,
+      `Starting SQLMesh LSP (cwd=${workspacePath})\n` +
+        `  serverOptions=${JSON.stringify(serverOptions)}\n` +
+        `  clientOptions=${JSON.stringify(clientOptions)}`,
     )
+
     this.client = new LanguageClient(
       'sqlmesh-lsp',
       'SQLMesh Language Server',
       serverOptions,
       clientOptions,
     )
+    this.explicitlyStopped = false // user wanted it running again
     await this.client.start()
     return ok(undefined)
   }
 
+  /** Restart = stop + start. */
   public async restart(
-    overrideByUser = false,
+    overrideStoppedByUser = false,
   ): Promise<Result<undefined, ErrorType>> {
-    await this.stop()
-    return await this.start(overrideByUser)
+    await this.stop() // this also disposes
+    return this.start(overrideStoppedByUser)
   }
 
+  /**
+   * Stop the client (if running) and clean up all VS Code resources so that a
+   * future `start()` registers its commands without collisions.
+   */
   public async stop(stoppedByUser = false): Promise<void> {
     if (this.client) {
-      traceInfo('Stopping SQLMesh Language Server')
-      await this.client.stop()
+      // Shut down the JSON‑RPC connection
+      await this.client
+        .stop()
+        .catch(err => traceError(`Error while stopping LSP: ${err}`))
+
+      // **Important** – unregister commands, code lenses, etc.
+      this.client.dispose()
+
       this.client = undefined
-      // Reset supported methods state when the client stops
       this.supportedMethodsState = { type: 'not-fetched' }
-      traceInfo('SQLMesh Language Server stopped')
+      traceInfo('SQLMesh LSP client disposed.')
     }
+
     if (stoppedByUser) {
       this.explicitlyStopped = true
       traceInfo('SQLMesh LSP client stopped by user.')
     }
   }
 
-  public async dispose() {
+  public async dispose(): Promise<void> {
     await this.stop()
   }
 
+  // ───────────────────────────────────────────── custom‑method helpers
+
   private async fetchSupportedMethods(): Promise<void> {
-    if (!this.client || this.supportedMethodsState.type !== 'not-fetched') {
+    if (!this.client || this.supportedMethodsState.type !== 'not-fetched')
       return
-    }
-    try {
-      const result = await this.internal_call_custom_method(
-        'sqlmesh/supported_methods',
-        {},
-      )
-      if (isErr(result)) {
-        traceError(`Failed to fetch supported methods: ${result.error}`)
-        this.supportedMethodsState = { type: 'endpoint-not-supported' }
-        return
-      }
-      const methodNames = new Set(result.value.methods.map(m => m.name))
-      this.supportedMethodsState = { type: 'fetched', methods: methodNames }
-      traceInfo(
-        `Fetched supported methods: ${Array.from(methodNames).join(', ')}`,
-      )
-    } catch {
-      // If the supported_methods endpoint doesn't exist, mark it as not supported
+
+    const result = await this.internal_call_custom_method(
+      'sqlmesh/supported_methods',
+      {},
+    )
+    if (isErr(result)) {
+      traceError(`Failed to fetch supported methods: ${result.error}`)
       this.supportedMethodsState = { type: 'endpoint-not-supported' }
-      traceInfo(
-        'Supported methods endpoint not available, proceeding without validation',
-      )
+      return
     }
+
+    const methodNames = new Set(result.value.methods.map(m => m.name))
+    this.supportedMethodsState = { type: 'fetched', methods: methodNames }
+    traceInfo(`Fetched supported methods: ${[...methodNames].join(', ')}`)
   }
 
   public async call_custom_method<
@@ -210,79 +216,49 @@ export class LSPClient implements Disposable {
     >
   > {
     if (!this.client) {
-      return err({
-        type: 'generic',
-        message: 'LSP client not ready.',
-      })
+      return err({ type: 'generic', message: 'LSP client not ready.' })
     }
+
     await this.fetchSupportedMethods()
 
     const supportedState = this.supportedMethodsState
-    switch (supportedState.type) {
-      case 'not-fetched':
-        return err({
-          type: 'invalid_state',
-          message: 'Supported methods not fetched yet whereas they should.',
-        })
-      case 'fetched': {
-        // If we have fetched the supported methods, we can check if the method is supported
-        if (!supportedState.methods.has(method)) {
-          return err({
-            type: 'sqlmesh_outdated',
-            message: `Method '${method}' is not supported by this LSP server.`,
-          })
-        }
-        const response = await this.internal_call_custom_method(
-          method,
-          request as any,
-        )
-        if (isErr(response)) {
-          return err({
-            type: 'generic',
-            message: response.error,
-          })
-        }
-        return ok(response.value as Response)
-      }
-      case 'endpoint-not-supported': {
-        const response = await this.internal_call_custom_method(
-          method,
-          request as any,
-        )
-        if (isErr(response)) {
-          return err({
-            type: 'generic',
-            message: response.error,
-          })
-        }
-        return ok(response.value as Response)
-      }
+    if (
+      supportedState.type === 'fetched' &&
+      !supportedState.methods.has(method)
+    ) {
+      return err({
+        type: 'sqlmesh_outdated',
+        message: `Method '${method}' is not supported by this LSP server.`,
+      })
+    }
+
+    const response = await this.internal_call_custom_method(
+      method,
+      request as any,
+    )
+    if (isErr(response)) {
+      return err({ type: 'generic', message: response.error })
     }
+    return ok(response.value as Response)
   }
 
   /**
-   * Internal method to call a custom LSP method without checking if the method is supported. It is used for
-   * the class whereas the `call_custom_method` checks if the method is supported.
+   * Low‑level helper that sends a raw JSON‑RPC request without any feature checks.
    */
   public async internal_call_custom_method<
     Method extends CustomLSPMethods['method'],
     Request extends Extract<CustomLSPMethods, { method: Method }>['request'],
     Response extends Extract<CustomLSPMethods, { method: Method }>['response'],
   >(method: Method, request: Request): Promise<Result<Response, string>> {
-    if (!this.client) {
-      return err('lsp client not ready')
-    }
+    if (!this.client) return err('lsp client not ready')
 
     try {
       const result = await this.client.sendRequest<Response>(method, request)
-      if (result.response_error) {
-        return err(result.response_error)
-      }
+      if ((result as any).response_error)
+        return err((result as any).response_error)
       return ok(result)
     } catch (error) {
-      traceError(
-        `lsp '${method}' request ${JSON.stringify(request)} failed: ${JSON.stringify(error)}`,
-      )
+      traceError(`LSP '${method}' request failed: ${JSON.stringify(error)}`)
       return err(JSON.stringify(error))
     }
   }