From a3b08593d837cde80efaf858657f777e357a3bbf Mon Sep 17 00:00:00 2001 From: Ben King <9087625+benfdking@users.noreply.github.com> Date: Tue, 29 Jul 2025 13:38:45 +0100 Subject: [PATCH] fix(vscode): improve registering of lsp in extension --- vscode/extension/src/commands/renderModel.ts | 5 +- vscode/extension/src/extension.ts | 193 ++++++++-------- vscode/extension/src/lsp/lsp.ts | 229 ++++++++----------- 3 files changed, 204 insertions(+), 223 deletions(-) diff --git a/vscode/extension/src/commands/renderModel.ts b/vscode/extension/src/commands/renderModel.ts index 6b5db3055c..24225c3e45 100644 --- a/vscode/extension/src/commands/renderModel.ts +++ b/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 { 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', { diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts index 74454f8fdb..0d0f6252cb 100644 --- a/vscode/extension/src/extension.ts +++ b/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 => { + 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}`) } } diff --git a/vscode/extension/src/lsp/lsp.ts b/vscode/extension/src/lsp/lsp.ts index f5c7532ae4..0b7a5b4b62 100644 --- a/vscode/extension/src/lsp/lsp.ts +++ b/vscode/extension/src/lsp/lsp.ts @@ -20,47 +20,43 @@ import { CustomLSPMethods } from './custom' type SupportedMethodsState = | { type: 'not-fetched' } | { type: 'fetched'; methods: Set } - // 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 + /** 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. + /** + * 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> { @@ -70,10 +66,19 @@ export class LSPClient implements Disposable { ) return ok(undefined) } + + // Guard against duplicate initialisation + if (this.isRunning) { + traceInfo('LSP client already running – start() is a no‑op.') + return ok(undefined) + } + + // Ensure we have an output channel if (!outputChannel) { outputChannel = window.createOutputChannel('sqlmesh-lsp') } + // Resolve sqlmesh executable const sqlmesh = await sqlmeshLspExec() if (isErr(sqlmesh)) { traceError( @@ -81,114 +86,110 @@ export class LSPClient implements Disposable { ) return sqlmesh } - const workspaceFolders = getWorkspaceFolders() - if (workspaceFolders.length === 0) { - traceError(`No workspace folders found`) - return err({ - type: 'generic', - message: 'No workspace folders found', - }) + + // 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> { - 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 { if (this.client) { - await this.client.stop() + // Shut down the JSON‑RPC connection + await this.client + .stop() + .catch(err => traceError(`Error while stopping LSP: ${err}`)) + + // Unregister commands, code lenses, etc. + await this.client.dispose() + this.client = undefined - // Reset supported methods state when the client stops this.supportedMethodsState = { type: 'not-fetched' } + traceInfo('SQLMesh LSP client disposed.') } + if (stoppedByUser) { this.explicitlyStopped = true traceInfo('SQLMesh LSP client stopped by user.') } } - public async dispose() { + public async dispose(): Promise { await this.stop() } private async fetchSupportedMethods(): Promise { - 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< @@ -208,79 +209,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['request'], Response extends Extract['response'], >(method: Method, request: Request): Promise> { - 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(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)) } }