fix(vscode): on startup pick the right model (#4556)

Author: benfdking

vscode/bus/src/callbacks.ts

@@ -1,15 +1,13 @@
+import type { Result } from './result'
+
+export type CallbackShape = Record<string, any>
+
 export type Callback = {
   openFile: {
     uri: string
   }
-  queryRequest: {
-    requestId: string
-    url: string
-    method?: string
-    params?: Record<string, string>
-    body?: Record<string, any>
-  }
-}
+  rpcResponse: RPCResponse
+} & CallbackShape
 
 /**
  * A tuple type representing a callback event with its associated payload.
@@ -23,22 +21,52 @@ export type CallbackEvent = {
   [K in keyof Callback]: { key: K; payload: Callback[K] }
 }[keyof Callback]
 
+export type VSCodeCallbackShape = Record<string, any>
+
 /**
  * A tuple type representing a VSCode event with its associated payload.
  */
 export type VSCodeCallback = {
   changeFocusOnFile: {
     path: string
   }
-  queryResponse: {
-    requestId: string
-    response: any
-  }
   savedFile: {
     fileUri: string
   }
-}
+  rpcRequest: RPCRequest
+} & VSCodeCallbackShape
 
 export type VSCodeEvent = {
   [K in keyof VSCodeCallback]: { key: K; payload: VSCodeCallback[K] }
 }[keyof VSCodeCallback]
+
+type RPCMethodsShape = Record<string, { params: any; result: any }>
+
+export type RPCMethods = {
+  get_active_file: {
+    params: {}
+    result: {
+      fileUri?: string
+    }
+  }
+  api_query: {
+    params: {
+      url: string
+      method: string
+      params: any
+      body: any
+    }
+    result: any
+  }
+} & RPCMethodsShape
+
+export type RPCRequest = {
+  requestId: string
+  method: keyof RPCMethods
+  params: RPCMethods[keyof RPCMethods]['params']
+}
+
+export type RPCResponse = {
+  requestId: string
+  result: Result<RPCMethods[keyof RPCMethods]['result'], string>
+}

vscode/extension/src/webviews/lineagePanel.ts

@@ -1,4 +1,4 @@
-import { CallbackEvent } from '@bus/callbacks'
+import { CallbackEvent, RPCRequest } from '@bus/callbacks'
 import {
   Disposable,
   Uri,
@@ -83,18 +83,43 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
             await window.showTextDocument(document)
             break
           }
-          case 'queryRequest': {
-            const payload = message.payload
-            const requestId = message.payload.requestId
-            const response = await this.lsp.call_custom_method(
-              'sqlmesh/api',
-              payload as any,
-            )
-            webviewView.webview.postMessage({
-              key: 'query_response',
-              payload: response,
-              requestId,
-            })
+          case 'rpcRequest': {
+            const payload: RPCRequest = message.payload
+            const requestId = payload.requestId
+            switch (payload.method) {
+              case 'api_query': {
+                const response = await this.lsp.call_custom_method(
+                  'sqlmesh/api',
+                  payload.params,
+                )
+                const responseCallback: CallbackEvent = {
+                  key: 'rpcResponse',
+                  payload: {
+                    requestId,
+                    result: response,
+                  },
+                }
+                webviewView.webview.postMessage(responseCallback)
+                break
+              }
+              case 'get_active_file': {
+                const active_file = window.activeTextEditor?.document.uri.fsPath
+                const responseCallback: CallbackEvent = {
+                  key: 'rpcResponse',
+                  payload: {
+                    requestId,
+                    result: {
+                      fileUri: active_file,
+                    },
+                  },
+                }
+                webviewView.webview.postMessage(responseCallback)
+                break
+              }
+              default: {
+                throw new Error(`Unhandled RPC method: ${payload.method}`)
+              }
+            }
             break
           }
           default:

vscode/react/src/api/index.ts

@@ -70,7 +70,15 @@ export function useApiModelLineage(
   return useQuery({
     queryKey: ['/api/lineage', modelName],
     queryFn: async ({ signal }) => {
-      return await modelLineageApiLineageModelNameGet(modelName, { signal })
+      try {
+        const response = await modelLineageApiLineageModelNameGet(modelName, {
+          signal,
+        })
+        return response
+      } catch (error) {
+        console.error('error fetching lineage', error)
+        throw error
+      }
     },
   })
 }

vscode/react/src/api/instance.ts

@@ -1,4 +1,4 @@
-import { sendVSCodeMessage } from '@/utils/vscodeapi'
+import { callRpc } from '@/utils/rpc'
 import { isErr } from '@bus/result'
 
 declare global {
@@ -38,43 +38,16 @@ export async function fetchAPI<T = any, B extends object = any>(
   config: FetchOptions<B>,
   _options?: Partial<FetchOptionsWithSignal>,
 ): Promise<T & ResponseWithDetail> {
-  // Generate a unique ID for this request
-  // Create a promise that will resolve when we get a response with matching ID
-  return new Promise((resolve, reject) => {
-    const requestId = `query_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`
-    const messageHandler = (event: MessageEvent) => {
-      if (
-        event.data &&
-        event.data.key === 'query_response' &&
-        event.data.requestId === requestId
-      ) {
-        // Remove the listener once we get our response
-        window.removeEventListener('message', messageHandler)
-
-        const payload = event.data.payload
-        if (isErr(payload)) {
-          reject(new Error(payload.error as string))
-        } else {
-          resolve(payload.value.data)
-        }
-      }
-    }
-
-    // Add the listener
-    window.addEventListener('message', messageHandler)
-
-    sendVSCodeMessage('queryRequest', {
-      requestId,
-      url: config.url,
-      params: config.params as any,
-      body: config.data,
-      method: config.method,
-    })
-
-    // Set a timeout to prevent hanging promises
-    setTimeout(() => {
-      window.removeEventListener('message', messageHandler)
-      reject(new Error('Query request timed out'))
-    }, 30000) // 30 second timeout
-  })
+  const request = {
+    url: config.url,
+    method: config.method,
+    params: config.params,
+    body: config.data,
+  }
+  const result = await callRpc('api_query', request)
+  if (isErr(result)) {
+    throw new Error(result.error)
+  }
+  const response = result.value.data
+  return response
 }

vscode/react/src/components/graph/ModelLineage.tsx

@@ -78,7 +78,7 @@ export function ModelLineage({
     useApiModelLineage(model.name)
   const { isFetching: isFetchingModels } = useApiModels()
 
-  const [isMegringModels, setIsMergingModels] = useState(false)
+  const [isMergingModels, setIsMergingModels] = useState(false)
   const [modelLineage, setModelLineage] = useState<
     ModelLineageApiLineageModelNameGet200 | undefined
   >(undefined)
@@ -91,7 +91,6 @@ export function ModelLineage({
     getModelLineage()
       .then(({ data }) => {
         setModelLineage(data)
-
         if (isNil(data)) return
 
         setIsMergingModels(true)
@@ -164,7 +163,7 @@ export function ModelLineage({
   }
 
   const isFetching =
-    isFetchingModelLineage || isFetchingModels || isMegringModels
+    isFetchingModelLineage || isFetchingModels || isMergingModels
 
   return (
     <div className="relative h-full w-full overflow-hidden">

vscode/react/src/components/graph/context.tsx

@@ -104,7 +104,7 @@ export const LineageFlowContext = createContext<LineageFlow>({
   setActiveEdges: () => {},
   handleClickModel: () => {},
   setManuallySelectedColumn: () => {},
-  handleError: () => {},
+  handleError: error => console.error(error),
   setLineage: () => {},
   setLineageCache: () => {},
   isActiveColumn: () => false,

vscode/react/src/components/graph/help.ts

@@ -176,7 +176,6 @@ function getNodeMap({
       ? decodeURI(modelName)
       : modelName
     const model = Object.values(models).find(m => m.fqn === decodedModelName)
-    console.log('model', model)
     const nodeType: LineageNodeModelType = isNotNil(model)
       ? (model.type as LineageNodeModelType)
       : // If model name present in lineage but not in global models