Merge pull request #365 from stone-lyl/test-file

fix: resolve "webView is disposed" issue when the webView is still being used
ajthinking · Jan 11, 2025 · d534cad · d534cad
2 parents df2480b + 58d3100
commit d534cad
Show file tree

Hide file tree

Showing 14 changed files with 78 additions and 37 deletions.
diff --git a/packages/core/src/InputObserverController.ts b/packages/core/src/InputObserverController.ts
@@ -63,7 +63,7 @@ export class InputObserverController {
     return items;
   }
 
-  observeLinkUpdate(observer: ObserveLinkUpdate): void {
+  observeLinkUpdate(observer: ObserveLinkUpdate): () => void {
     const subscription = this.items$.pipe(
       filter(payload => observer.linkIds.includes(payload.linkId)),
       map(payload => payload.items),
@@ -76,9 +76,16 @@ export class InputObserverController {
     ).subscribe();
 
     if (observer?.observerId && subscription) this.observerMap.set(observer.observerId, subscription);
+
+    return () => {
+      this.deleteExecutionObserver({
+        type: RequestObserverType.cancelObservation,
+        observerId: observer.observerId
+      });
+    };
   }
 
-  addLinkItemsObserver(observer: ObserveLinkItems): void {
+  addLinkItemsObserver(observer: ObserveLinkItems): () => void {
     const subscription = this.items$.pipe(
       filter(payload => observer.linkIds.includes(payload.linkId)),
       bufferTime(observer.throttleMs ?? ThrottleMS),
@@ -106,9 +113,16 @@ export class InputObserverController {
     ).subscribe();
 
     if (observer?.observerId && subscription) this.observerMap.set(observer.observerId, subscription);
+
+    return () => {
+      this.deleteExecutionObserver({
+        type: RequestObserverType.cancelObservation,
+        observerId: observer.observerId
+      });
+    };
   }
 
-  addLinkCountsObserver(observer: ObserveLinkCounts): void {
+  addLinkCountsObserver(observer: ObserveLinkCounts): () => void {
     const subscription = this.links$.pipe(
       filter(payload => observer.linkIds.includes(payload.linkId)),
       map(payload => payload),
@@ -123,9 +137,16 @@ export class InputObserverController {
     ).subscribe();
 
     if (observer?.observerId && subscription) this.observerMap.set(observer.observerId, subscription);
+
+    return () => {
+      this.deleteExecutionObserver({
+        type: RequestObserverType.cancelObservation,
+        observerId: observer.observerId
+      });
+    };
   }
 
-  addNodeStatusObserver(observer: ObserveNodeStatus): void {
+  addNodeStatusObserver(observer: ObserveNodeStatus): () => void {
     const subscription = this.nodeStatus$.pipe(
       filter(payload => observer.nodeIds.includes(payload.nodeId)),
       bufferTime(observer.throttleMs ?? ThrottleMS),
@@ -143,7 +164,14 @@ export class InputObserverController {
       })
     ).subscribe();
 
-    if (observer?.observerId && subscription) this.observerMap.set(observer.observerId, subscription);
+    if (observer.observerId && subscription) this.observerMap.set(observer.observerId, subscription);
+
+    return () => {
+      this.deleteExecutionObserver({
+        type: RequestObserverType.cancelObservation,
+        observerId: observer.observerId
+      });
+    };
   }
 
   deleteExecutionObserver(observer: CancelObservation): void {

diff --git a/packages/core/src/types/NodeDescription.ts b/packages/core/src/types/NodeDescription.ts
@@ -27,9 +27,8 @@ export const NodeDescriptionResponseSchema = z.object({
     invalid_type_error: 'path must be a string'
   }).optional(),
   status: z.literal('server-post', {
-    required_error: 'status is required',
     invalid_type_error: 'status must be server-post'
-  }),
+  }).optional(),
   type: z.literal('getNodeDescriptions', {
     required_error: 'type is required',
     invalid_type_error: 'type must be getNodeDescriptions'

diff --git a/packages/ds-ext/src/DiagramEditorProvider.ts b/packages/ds-ext/src/DiagramEditorProvider.ts
@@ -92,8 +92,15 @@ export class DiagramEditorProvider implements vscode.CustomEditorProvider<Diagra
     // Set the webview's HTML content
     webviewPanel.webview.html = this.getWebviewContent(webviewPanel.webview, document);
 
-    // Handle messages from the webview
-    webviewPanel.webview.onDidReceiveMessage(event => {
+    const postMessage = (msg: any) => {
+      // Webview is disposed, but webviewPanel.postMessage is still being referenced. To prevent errors, we need to mock webviewPanel.postMessage.
+      if (webviewPanel?.webview?.postMessage) {
+        return webviewPanel?.webview?.postMessage(msg);
+      }
+    };
+
+    const disposables: (Promise<() => void>)[] = [];
+    const unsubscribe = webviewPanel.webview.onDidReceiveMessage(event => {
       const handlers: Record<string, MessageHandler> = {
         run: onRun,
         getNodeDescriptions: onGetNodeDescriptions,
@@ -114,7 +121,15 @@ export class DiagramEditorProvider implements vscode.CustomEditorProvider<Diagra
         return;
       }
 
-      handler({ webviewPanel, event, document, inputObserverController: this.inputObserverController });
+      const disposable = handler({ postMessage, event, document, inputObserverController: this.inputObserverController });
+      if(typeof disposable === 'function' && disposable){
+        disposables.push(disposable);
+      }
+    });
+
+    webviewPanel.onDidDispose(() => {
+      unsubscribe.dispose();
+      disposables.forEach(async (disposable) => (await disposable)());
     });
   }
 

diff --git a/packages/ds-ext/src/MessageHandler.ts b/packages/ds-ext/src/MessageHandler.ts
@@ -1,12 +1,11 @@
-import * as vscode from 'vscode';
 import { DiagramDocument } from './DiagramDocument';
 import { InputObserverController } from '@data-story/core';
 
 export type MessageHandlerArgs = {
-  webviewPanel: vscode.WebviewPanel;
+  postMessage: (msg: any) => Thenable<boolean> | undefined;
   event: any;
   document: DiagramDocument,
   inputObserverController: InputObserverController
 };
 
-export type MessageHandler = (args: MessageHandlerArgs) => Promise<void>;
+export type MessageHandler = (args: MessageHandlerArgs) => Promise<() => void> | Promise<void>;
diff --git a/packages/ds-ext/src/commands/makeDensityDatasets.ts b/packages/ds-ext/src/commands/makeDensityDatasets.ts
@@ -26,7 +26,7 @@ export const makeDensityDatasets = (directory: string) => {
 
   // Generate datasets based on density vectors
   const createDatasets = (fileDensity: any, itemsDensity: any, schemaDensity: any) => {
-    const datasetDir = path.join(datasetsDir, `${fileDensity}:${itemsDensity}:${schemaDensity}`);
+    const datasetDir = path.join(datasetsDir, `${fileDensity}_${itemsDensity}_${schemaDensity}`);
     if (!fs.existsSync(datasetDir)) fs.mkdirSync(datasetDir);
 
     for (let i = 0; i < fileDensity; i++) {

diff --git a/packages/ds-ext/src/messageHandlers/cancelObservation.ts b/packages/ds-ext/src/messageHandlers/cancelObservation.ts
@@ -1,9 +1,9 @@
 import { MessageHandler } from '../MessageHandler';
 import { CancelObservation } from '@data-story/core';
 
-export const cancelObservation: MessageHandler = async({ event, webviewPanel, inputObserverController }) => {
+export const cancelObservation: MessageHandler = async({ event, postMessage, inputObserverController }) => {
   inputObserverController.deleteExecutionObserver(event as CancelObservation);
-  webviewPanel.webview.postMessage({
+  postMessage?.({
     msgId: event!.msgId,
   });
 };
diff --git a/packages/ds-ext/src/messageHandlers/getDataFromStorage.ts b/packages/ds-ext/src/messageHandlers/getDataFromStorage.ts
@@ -5,9 +5,9 @@ import {
   LinkId
 } from '@data-story/core';
 
-export const getDataFromStorage: MessageHandler = async({ event, webviewPanel, inputObserverController }) => {
+export const getDataFromStorage: MessageHandler = async({ event, postMessage, inputObserverController }) => {
   const result: Record<LinkId, ItemValue[]> = await inputObserverController.getDataFromStorage(event as GetDataFromStorageParams);
-  webviewPanel.webview.postMessage({
+  postMessage?.({
     ...event,
     data: result,
   });

diff --git a/packages/ds-ext/src/messageHandlers/getDiagram.ts b/packages/ds-ext/src/messageHandlers/getDiagram.ts
@@ -1,14 +1,14 @@
 import { Diagram } from '@data-story/core';
 import { MessageHandler } from '../MessageHandler';
 
-export const getDiagram: MessageHandler = async ({ event, webviewPanel, document }) => {
+export const getDiagram: MessageHandler = async ({ event, postMessage, document }) => {
   const diagramData = new TextDecoder().decode(document.data);
   let diagram = new Diagram();
   if(diagramData) {
     diagram = JSON.parse(diagramData);
   }
 
-  webviewPanel.webview.postMessage({
+  postMessage?.({
     ...event,
     type: 'getDiagram',
     diagram: diagram,

diff --git a/packages/ds-ext/src/messageHandlers/observeLinkCounts.ts b/packages/ds-ext/src/messageHandlers/observeLinkCounts.ts
@@ -4,11 +4,11 @@ import {
   RequestObserverType
 } from '@data-story/core';
 
-export const observeLinkCounts: MessageHandler = async({ event, webviewPanel, inputObserverController }) => {
-  inputObserverController.addLinkCountsObserver({
+export const observeLinkCounts: MessageHandler = async({ event, postMessage, inputObserverController }) => {
+  return inputObserverController.addLinkCountsObserver({
     ...event,
     onReceive: ({ links }) => {
-      webviewPanel.webview.postMessage({
+      postMessage?.({
         links: links,
         type: RequestObserverType.observeLinkCounts,
         msgId: event!.msgId,

diff --git a/packages/ds-ext/src/messageHandlers/observeLinkItems.ts b/packages/ds-ext/src/messageHandlers/observeLinkItems.ts
@@ -1,11 +1,11 @@
 import { MessageHandler } from '../MessageHandler';
 import { InputObserver, type ObserveLinkItems, ItemValue, RequestObserverType } from '@data-story/core';
 
-export const observeLinkItems: MessageHandler = async ({ event, webviewPanel, inputObserverController }) => {
-  inputObserverController.addLinkItemsObserver({
+export const observeLinkItems: MessageHandler = async ({ event, postMessage, inputObserverController }) => {
+  return inputObserverController.addLinkItemsObserver({
     ...event,
     onReceive: (items: ItemValue[], inputObserver: InputObserver) => {
-      webviewPanel.webview.postMessage({
+      postMessage?.({
         items,
         inputObserver,
         type: RequestObserverType.observeLinkItems,

diff --git a/packages/ds-ext/src/messageHandlers/observeLinkUpdate.ts b/packages/ds-ext/src/messageHandlers/observeLinkUpdate.ts
@@ -4,11 +4,11 @@ import {
   RequestObserverType
 } from '@data-story/core';
 
-export const observeLinkUpdate: MessageHandler = async({ event, webviewPanel, inputObserverController }) => {
-  inputObserverController.observeLinkUpdate({
+export const observeLinkUpdate: MessageHandler = async({ event, postMessage, inputObserverController }) => {
+  return inputObserverController.observeLinkUpdate({
     ...event,
     onReceive: () => {
-      webviewPanel.webview.postMessage({
+      postMessage?.({
         linkIds: event.linkIds,
         type: RequestObserverType.observeLinkUpdate,
         msgId: event!.msgId,

diff --git a/packages/ds-ext/src/messageHandlers/observeNodeStatus.ts b/packages/ds-ext/src/messageHandlers/observeNodeStatus.ts
@@ -4,11 +4,11 @@ import {
   RequestObserverType
 } from '@data-story/core';
 
-export const observeNodeStatus: MessageHandler = async({ event, webviewPanel, inputObserverController }) => {
-  inputObserverController.addNodeStatusObserver({
+export const observeNodeStatus: MessageHandler = async({ event, postMessage, inputObserverController }) => {
+  return inputObserverController.addNodeStatusObserver({
     ...event,
     onReceive: ({ nodes }) => {
-      webviewPanel.webview.postMessage({
+      postMessage?.({
         nodes: nodes,
         type: RequestObserverType.observeNodeStatus,
         msgId: event!.msgId,

diff --git a/packages/ds-ext/src/messageHandlers/onGetNodeDescriptions.ts b/packages/ds-ext/src/messageHandlers/onGetNodeDescriptions.ts
@@ -1,10 +1,10 @@
 import { MessageHandler } from '../MessageHandler';
 import { createAndBootApp } from '../app/createAndBootApp';
 
-export const onGetNodeDescriptions: MessageHandler = async ({ event, webviewPanel }) => {
+export const onGetNodeDescriptions: MessageHandler = async ({ event, postMessage }) => {
   const app = await createAndBootApp();
 
-  webviewPanel.webview.postMessage({
+  postMessage?.({
     ...event,
     awaited: true,
     availableNodes: app.descriptions(),

diff --git a/packages/ds-ext/src/messageHandlers/onRun.ts b/packages/ds-ext/src/messageHandlers/onRun.ts
@@ -21,7 +21,7 @@ function setWorkspaceFolderPath() {
   }
 }
 
-export const onRun: MessageHandler = async ({ event, webviewPanel, inputObserverController }) => {
+export const onRun: MessageHandler = async ({ event, postMessage, inputObserverController }) => {
   const app = await createAndBootApp();
 
   const diagram = new Diagram({
@@ -45,13 +45,13 @@ export const onRun: MessageHandler = async ({ event, webviewPanel, inputObserver
     for await(const update of execution) {}
 
     const endTime = Date.now();
-    webviewPanel.webview.postMessage({
+    postMessage?.({
       msgId,
       type: 'ExecutionResult',
       time: endTime - startTime
     });
   } catch(error: any) {
-    webviewPanel.webview.postMessage({
+    postMessage?.({
       msgId,
       type: 'ExecutionFailure',
       error: error.message