feat(vscode): refresh lineage on file changes (#4415)

Author: benfdking

vscode/bus/src/callbacks.ts

@@ -34,6 +34,9 @@ export type VSCodeCallback = {
     requestId: string
     response: any
   }
+  savedFile: {
+    fileUri: string
+  }
 }
 
 export type VSCodeEvent = {

vscode/extension/src/webviews/lineagePanel.ts

@@ -26,6 +26,16 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
       this.panel.webview.html = this.getHTML(this.panel.webview)
     }
 
+    workspace.onDidSaveTextDocument(document => {
+      this.panel?.webview.postMessage({
+        key: 'vscode_send',
+        payload: {
+          key: 'savedFile',
+          payload: { fileUri: document.uri.fsPath },
+        },
+      })
+    })
+
     window.onDidChangeActiveTextEditor(editor => {
       if (editor) {
         this.panel?.webview.postMessage({

vscode/react/src/hooks/eventBus.tsx

@@ -15,6 +15,9 @@ export type EventMap = {
   changeFocusedFile: {
     fileUri: string
   }
+  savedFile: {
+    fileUri: string
+  }
 }
 
 /**

vscode/react/src/pages/lineage.tsx

@@ -3,6 +3,7 @@ import {
   QueryCache,
   QueryClient,
   QueryClientProvider,
+  useQueryClient,
 } from '@tanstack/react-query'
 import { useApiModels } from '@/api'
 import LineageFlowProvider from '@/components/graph/context'
@@ -24,9 +25,11 @@ export function LineagePage() {
         const payload: VSCodeEvent = event.data.payload
         switch (payload.key) {
           case 'changeFocusOnFile':
-            console.log('emitted focus to file:', payload.payload.path)
             emit('changeFocusedFile', { fileUri: payload.payload.path })
             break
+          case 'savedFile':
+            emit('savedFile', { fileUri: payload.payload.fileUri })
+            break
           default:
             console.error(
               'Unhandled message type in lineage page:',
@@ -65,6 +68,7 @@ function Lineage() {
     undefined,
   )
   const { on } = useEventBus()
+  const queryClient = useQueryClient()
 
   const { data: models, isLoading: isLoadingModels } = useApiModels()
   React.useEffect(() => {
@@ -104,6 +108,10 @@ function Lineage() {
       setSelectedModelSet(model.name)
     }
   })
+  on('savedFile', () => {
+    queryClient.invalidateQueries()
+  })
+
   return (
     <LineageComponentFromWeb
       selectedModel={selectedModel}