fix(vscode): on sign in, restart lsp (#4772)

Author: benfdking

vscode/extension/src/commands/format.ts

@@ -11,12 +11,18 @@ export const format =
   (
     authProvider: AuthenticationProviderTobikoCloud,
     lsp: LSPClient | undefined,
+    restartLSP: () => Promise<void>,
   ) =>
   async (): Promise<void> => {
     traceLog('Calling format')
     const out = await internalFormat(lsp)
     if (isErr(out)) {
-      return handleError(authProvider, out.error, 'Project format failed')
+      return handleError(
+        authProvider,
+        restartLSP,
+        out.error,
+        'Project format failed',
+      )
     }
     vscode.window.showInformationMessage('Project formatted successfully')
   }

vscode/extension/src/commands/signin.ts

@@ -1,13 +1,30 @@
 import { AuthenticationProviderTobikoCloud } from '../auth/auth'
 import * as vscode from 'vscode'
 import { isCodespaces } from '../utilities/isCodespaces'
+import { traceInfo } from '../utilities/common/log'
 
 export const signIn =
-  (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
+  (
+    authenticationProvider: AuthenticationProviderTobikoCloud,
+    onSignInSuccess: () => Promise<void>,
+  ) =>
+  async () => {
     if (isCodespaces()) {
       await authenticationProvider.sign_in_device_flow()
     } else {
       await authenticationProvider.createSession()
     }
-    await vscode.window.showInformationMessage('Signed in successfully')
+
+    // Do not await this, as this will block the thread, you just need to show the message, but not block
+    vscode.window.showInformationMessage('Signed in successfully')
+
+    // Execute callback after successful sign-in
+    if (onSignInSuccess) {
+      traceInfo('Executing post sign-in callback')
+      try {
+        await onSignInSuccess()
+      } catch (error) {
+        traceInfo(`Error in post sign-in callback: ${error}`)
+      }
+    }
   }

vscode/extension/src/commands/signinSpecifyFlow.ts

@@ -3,7 +3,11 @@ import { traceInfo } from '../utilities/common/log'
 import { window } from 'vscode'
 
 export const signInSpecifyFlow =
-  (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
+  (
+    authenticationProvider: AuthenticationProviderTobikoCloud,
+    onSignInSuccess?: () => Promise<void>,
+  ) =>
+  async () => {
     traceInfo('Sign in specify flow')
     const flowOptions = [
       {
@@ -24,11 +28,31 @@ export const signInSpecifyFlow =
       await authenticationProvider.sign_in_oauth_flow()
       await authenticationProvider.getSessions()
       await window.showInformationMessage('Sign in success')
+
+      // Execute callback after successful sign-in
+      if (onSignInSuccess) {
+        traceInfo('Executing post sign-in callback')
+        try {
+          await onSignInSuccess()
+        } catch (error) {
+          traceInfo(`Error in post sign-in callback: ${error}`)
+        }
+      }
       return
     } else if (selectedFlow.label === 'Device Flow') {
       await authenticationProvider.sign_in_device_flow()
       await authenticationProvider.getSessions()
       await window.showInformationMessage('Sign in success')
+
+      // Execute callback after successful sign-in
+      if (onSignInSuccess) {
+        traceInfo('Executing post sign-in callback')
+        try {
+          await onSignInSuccess()
+        } catch (error) {
+          traceInfo(`Error in post sign-in callback: ${error}`)
+        }
+      }
       return
     } else {
       traceInfo('Invalid flow selected')

vscode/extension/src/extension.ts

@@ -45,12 +45,21 @@ export async function activate(context: vscode.ExtensionContext) {
   traceInfo('Authentication provider registered')
 
   context.subscriptions.push(
-    vscode.commands.registerCommand('sqlmesh.signin', signIn(authProvider)),
+    vscode.commands.registerCommand(
+      'sqlmesh.signin',
+      signIn(authProvider, async () => {
+        traceInfo('Restarting LSP after sign-in')
+        await restart()
+      }),
+    ),
   )
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'sqlmesh.signinSpecifyFlow',
-      signInSpecifyFlow(authProvider),
+      signInSpecifyFlow(authProvider, async () => {
+        traceInfo('Restarting LSP after sign-in')
+        await restart()
+      }),
     ),
   )
   context.subscriptions.push(
@@ -76,13 +85,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  context.subscriptions.push(
-    vscode.commands.registerCommand(
-      'sqlmesh.format',
-      format(authProvider, lspClient),
-    ),
-  )
-
   // Register the webview
   const lineagePanel = new LineagePanel(context.extensionUri, lspClient)
   context.subscriptions.push(
@@ -118,14 +120,35 @@ export async function activate(context: vscode.ExtensionContext) {
       if (isErr(restartResult)) {
         return handleError(
           authProvider,
+          restart,
           restartResult.error,
           'LSP restart failed',
         )
       }
       context.subscriptions.push(lspClient)
+    } else {
+      lspClient = new LSPClient()
+      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(
+    vscode.commands.registerCommand(
+      'sqlmesh.format',
+      format(authProvider, lspClient, restart),
+    ),
+  )
+
   context.subscriptions.push(
     onDidChangePythonInterpreter(async () => {
       await restart()
@@ -140,7 +163,12 @@ export async function activate(context: vscode.ExtensionContext) {
 
   const result = await lspClient.start()
   if (isErr(result)) {
-    return handleError(authProvider, result.error, 'Failed to start LSP')
+    return handleError(
+      authProvider,
+      restart,
+      result.error,
+      'Failed to start LSP',
+    )
   } else {
     context.subscriptions.push(lspClient)
   }

vscode/extension/src/utilities/errors.ts

@@ -69,6 +69,7 @@ interface SqlmeshLspDependenciesMissingError {
 
 export async function handleError(
   authProvider: AuthenticationProviderTobikoCloud,
+  restartLsp: () => Promise<void>,
   error: ErrorType,
   genericErrorPrefix?: string,
 ): Promise<void> {
@@ -85,7 +86,7 @@ export async function handleError(
       )
       return
     case 'not_signed_in':
-      return handleNotSignedInError(authProvider)
+      return handleNotSignedInError(authProvider, restartLsp)
     case 'sqlmesh_not_found':
       return handleSqlmeshNotFoundError()
     case 'sqlmesh_lsp_not_found':
@@ -110,14 +111,15 @@ export async function handleError(
  */
 const handleNotSignedInError = async (
   authProvider: AuthenticationProviderTobikoCloud,
+  restartLsp: () => Promise<void>,
 ): Promise<void> => {
   traceInfo('handleNotSginedInError')
   const result = await window.showInformationMessage(
     'Please sign in to Tobiko Cloud to use SQLMesh',
     'Sign In',
   )
   if (result === 'Sign In') {
-    await signIn(authProvider)()
+    await signIn(authProvider, restartLsp)()
   }
 }