Fix IMC bugs when multiple apps loaded at the same time & prevent same view ID from loading again

Author: Shellishack

npm-packages/react-api/src/hooks/editor/use-snapshot-state.ts

@@ -20,15 +20,20 @@ export default function useSnapShotState<T>(
   );
 
   // Update context whenever state changes
-  const setSnapshotState = (value: T | ((prev: T) => T)) => {
+  const setSnapshotState: React.Dispatch<React.SetStateAction<T>> = (value) => {
     setState((prev) => {
-      const next =
+      const newValue =
         typeof value === "function" ? (value as (prev: T) => T)(prev) : value;
-      setStates((prevStates) => ({
-        ...prevStates,
-        [key]: next,
-      }));
-      return next;
+
+      // Defer the setStates call to next microtask, outside render phase
+      Promise.resolve().then(() => {
+        setStates((prevStates) => ({
+          ...prevStates,
+          [key]: newValue,
+        }));
+      });
+
+      return newValue;
     });
   };
 

npm-packages/react-api/src/lib/use-imc.tsx

@@ -1,6 +1,6 @@
-import { InterModuleCommunication } from "@pulse-editor/shared-utils";
 import {
   IMCMessageTypeEnum,
+  InterModuleCommunication,
   ReceiverHandlerMap,
 } from "@pulse-editor/shared-utils";
 import { useEffect, useState } from "react";
@@ -35,9 +35,8 @@ export default function useIMC(handlerMap: ReceiverHandlerMap) {
       await newImc.initOtherWindow(targetWindow);
       setImc(newImc);
 
-      newImc.sendMessage(IMCMessageTypeEnum.AppReady).then(() => {
-        setIsReady(true);
-      });
+      await newImc.sendMessage(IMCMessageTypeEnum.AppReady);
+      setIsReady(true);
     }
 
     initIMC();

npm-packages/shared-utils/src/imc/poly-imc.ts

@@ -145,18 +145,26 @@ export class ConnectionListener {
 
   private listener: InterModuleCommunication;
 
+  /**
+   *
+   * @param polyIMC The polyIMC instance.
+   * @param newConnectionReceiverHandlerMap Receiver handler map for newly established poly-IMC channel.
+   * @param onConnection Callback function to be called when a new connection is established.
+   * @param expectedOtherWindowId Optional expected other window ID to validate incoming connections.
+   */
   constructor(
     polyIMC: PolyIMC,
     newConnectionReceiverHandlerMap: ReceiverHandlerMap,
-    onConnection?: (senderWindow: Window, message: IMCMessage) => void
+    onConnection?: (senderWindow: Window, message: IMCMessage) => void,
+    expectedOtherWindowId?: string
   ) {
     this.polyIMC = polyIMC;
     this.newConnectionReceiverHandlerMap = newConnectionReceiverHandlerMap;
     this.onConnection = onConnection;
 
     const listener = new InterModuleCommunication();
     this.listener = listener;
-    listener.initThisWindow(window);
+    listener.initThisWindow(window, expectedOtherWindowId);
 
     listener.updateReceiverHandlerMap(
       new Map<IMCMessageTypeEnum, ReceiverHandler>([

web/components/app-loaders/sandbox-app-loader.tsx

@@ -113,7 +113,7 @@ export default function SandboxAppLoader({
 
     if (currentViewId) {
       // Listen for an incoming extension connection
-      console.log("Listening for extension connection...");
+      console.log(`[${currentViewId}]: Listening for app connection...`);
       listenForExtensionConnection();
 
       setIsLookingForExtension(true);
@@ -136,6 +136,8 @@ export default function SandboxAppLoader({
   // When IMC is connected, remove the connection listener
   useEffect(() => {
     if (isConnected && clRef.current) {
+      console.log(`[${currentViewId}]: App connected.`);
+      // Close the connection listener
       clRef.current.close();
       clRef.current = null;
     }
@@ -184,6 +186,7 @@ export default function SandboxAppLoader({
         }: {
           isLoading: boolean;
         } = message.payload;
+        console.log(`[${model.viewId}]: App is loading: `, isLoading);
         setIsLoadingExtension((prev) => isLoading);
         if (onInitialLoaded) {
           onInitialLoaded();
@@ -296,6 +299,7 @@ export default function SandboxAppLoader({
         (senderWindow: Window, message: IMCMessage) => {
           setIsConnected((prev) => true);
         },
+        viewModel.viewId,
       );
       clRef.current = cl;
     }

web/components/marketplace/workflow/workflow-preview-card.tsx

@@ -105,9 +105,7 @@ export default function WorkflowPreviewCard({
               color="primary"
               size="sm"
               onPress={() => {
-                if (isPressable) {
-                  openWorkflow();
-                }
+                openWorkflow();
               }}
             >
               Use

web/components/modals/publish-workflow-modal.tsx

@@ -29,39 +29,50 @@ export default function PublishWorkflowModal({
   const [version, setVersion] = useState("");
 
   async function publishWorkflow() {
-    if (!workflowCanvas) {
-      console.error("Workflow canvas is not available");
-      return;
-    }
+    try {
+      if (!workflowCanvas) {
+        console.error("Workflow canvas is not available");
+        return;
+      }
 
-    setIsOpen(false);
+      setIsOpen(false);
 
-    const res = await captureWorkflowCanvas(workflowCanvas);
-    const dataUrl = res.toDataURL("image/png");
+      const res = await captureWorkflowCanvas(workflowCanvas);
+      const dataUrl = res.toDataURL("image/png");
 
-    const workflow: Workflow = {
-      name: name,
-      thumbnail: dataUrl,
-      content: {
-        nodes: localNodes ?? [],
-        edges: localEdges ?? [],
-        defaultEntryPoint: entryPoint,
-        snapshotStates: await saveAppsSnapshotStates(),
-      },
-      version: version,
-      visibility: "public",
-    };
+      const snapshotStates = await saveAppsSnapshotStates();
 
-    await fetchAPI("/api/workflow/publish", {
-      method: "POST",
-      body: JSON.stringify({ ...workflow }),
-    });
+      const workflow: Workflow = {
+        name: name,
+        thumbnail: dataUrl,
+        content: {
+          nodes: localNodes ?? [],
+          edges: localEdges ?? [],
+          defaultEntryPoint: entryPoint,
+          snapshotStates: snapshotStates,
+        },
+        version: version,
+        visibility: "public",
+      };
 
-    addToast({
-      title: "Workflow Published",
-      description: "Your workflow has been published successfully.",
-      color: "success",
-    });
+      await fetchAPI("/api/workflow/publish", {
+        method: "POST",
+        body: JSON.stringify({ ...workflow }),
+      });
+
+      addToast({
+        title: "Workflow Published",
+        description: "Your workflow has been published successfully.",
+        color: "success",
+      });
+    } catch (error) {
+      console.error("Error publishing workflow:", error);
+      addToast({
+        title: "Error",
+        description: "There was an error publishing your workflow.",
+        color: "danger",
+      });
+    }
   }
 
   async function handlePress() {

web/components/providers/imc-provider.tsx

@@ -27,7 +27,7 @@ import {
   STTConfig,
   TTSConfig,
 } from "@pulse-editor/shared-utils";
-import { createContext, useContext, useEffect, useRef, useState } from "react";
+import { createContext, useContext, useEffect, useRef } from "react";
 import { EditorContext } from "./editor-context-provider";
 
 export const IMCContext = createContext<IMCContextType | undefined>(undefined);
@@ -39,7 +39,7 @@ export default function InterModuleCommunicationProvider({
 }) {
   const editorContext = useContext(EditorContext);
 
-  const [polyIMC, setPolyIMC] = useState<PolyIMC | undefined>(undefined);
+  const polyIMCRef = useRef<PolyIMC | undefined>(undefined);
   const imcInitializedMapRef = useRef<Map<string, boolean>>(new Map());
   const imcInitializedResolvePromisesRef = useRef<{
     [key: string]: () => void;
@@ -53,29 +53,23 @@ export default function InterModuleCommunicationProvider({
   useEffect(() => {
     // @ts-expect-error set window viewId
     window.viewId = "Pulse Editor Main";
+    polyIMCRef.current = new PolyIMC(getHandlerMap());
 
     return () => {
       // Cleanup the polyIMC instance when the component unmounts
-      if (polyIMC) {
-        polyIMC.close();
-        setPolyIMC(undefined);
+      if (polyIMCRef) {
+        polyIMCRef.current?.close();
+        polyIMCRef.current = undefined;
       }
     };
   }, []);
 
-  useEffect(() => {
-    if (!polyIMC) {
-      const newPolyIMC = new PolyIMC(getHandlerMap());
-      setPolyIMC(newPolyIMC);
-    }
-  }, [polyIMC, setPolyIMC]);
-
   // Update the base handler map as editor context changes
   useEffect(() => {
-    if (polyIMC) {
-      polyIMC.updateBaseReceiverHandlerMap(getHandlerMap());
+    if (polyIMCRef.current) {
+      polyIMCRef.current?.updateBaseReceiverHandlerMap(getHandlerMap());
     }
-  }, [polyIMC, editorContext]);
+  }, [editorContext]);
 
   function markIMCInitialized(viewId: string) {
     imcInitializedMapRef.current.set(viewId, true);
@@ -489,7 +483,7 @@ export default function InterModuleCommunicationProvider({
   return (
     <IMCContext.Provider
       value={{
-        polyIMC,
+        polyIMC: polyIMCRef.current,
         resolveWhenViewInitialized,
         markIMCInitialized,
         resolveWhenActionRegistered,

web/lib/hooks/use-canvas-workflow.ts

@@ -385,12 +385,18 @@ export default function useCanvasWorkflow(
     const appStates = await Promise.all(
       apps.map(async (app) => {
         if (!app.viewId) return null;
-        const { states } = await imcContext?.polyIMC?.sendMessage(
-          app.viewId,
-          IMCMessageTypeEnum.EditorAppStateSnapshotSave,
-        );
 
-        return { appId: app.viewId, states: states };
+        // Do a time out because the app may not use snapshot feature
+        return await Promise.race([
+          new Promise<any>((resolve) => setTimeout(() => resolve(null), 2000)),
+          (async () => {
+            const { states } = await imcContext?.polyIMC?.sendMessage(
+              app.viewId,
+              IMCMessageTypeEnum.EditorAppStateSnapshotSave,
+            );
+            return { appId: app.viewId, states: states };
+          })(),
+        ]);
       }),
     );
 

web/lib/hooks/use-tab-view-manager.ts

@@ -1,5 +1,6 @@
 import { EditorContext } from "@/components/providers/editor-context-provider";
 import { IMCContext } from "@/components/providers/imc-provider";
+import { addToast } from "@heroui/react";
 import { ViewModeEnum } from "@pulse-editor/shared-utils";
 import { useContext, useEffect, useState } from "react";
 import { v4 } from "uuid";
@@ -206,6 +207,15 @@ export function useTabViewManager() {
         tabIndex: newIdx,
       };
     });
+
+    if (view.type === ViewModeEnum.Canvas) {
+      // Remove the app nodes' view IDs from IMC context
+      (view.config as CanvasViewConfig).appConfigs?.forEach((appConfig) => {
+        imcContext?.polyIMC?.removeChannel(appConfig.viewId);
+      });
+    } else if (view.type === ViewModeEnum.App) {
+      imcContext?.polyIMC?.removeChannel((view.config as AppViewConfig).viewId);
+    }
   }
 
   /**
@@ -222,6 +232,19 @@ export function useTabViewManager() {
         tabIndex: -1,
       };
     });
+
+    // For all tab views, remove their view IDs from IMC context
+    tabViews.forEach((view) => {
+      if (view.type === ViewModeEnum.Canvas) {
+        (view.config as CanvasViewConfig).appConfigs?.forEach((appConfig) => {
+          imcContext?.polyIMC?.removeChannel(appConfig.viewId);
+        });
+      } else if (view.type === ViewModeEnum.App) {
+        imcContext?.polyIMC?.removeChannel(
+          (view.config as AppViewConfig).viewId,
+        );
+      }
+    });
   }
 
   function viewCount(): number {
@@ -265,6 +288,20 @@ export function useTabViewManager() {
       throw new Error("IMC context is not available");
     }
 
+    // Prohibit creating canvas if any app's view ID in the canvas already exists
+    const existViewId = canvasConfig.appConfigs?.find((appConfig) =>
+      imcContext?.polyIMC?.hasChannel(appConfig.viewId),
+    );
+
+    if (existViewId) {
+      addToast({
+        title: "Error creating canvas",
+        description: `Same app nodes already exist. Your workflow might already be opened in another tab.`,
+        color: "danger",
+      });
+      return undefined;
+    }
+
     const newTabView: TabView = {
       type: ViewModeEnum.Canvas,
       config: canvasConfig,
@@ -297,6 +334,10 @@ export function useTabViewManager() {
       currentTab = await createCanvasTabView({
         viewId: `canvas-${v4()}`,
       } as CanvasViewConfig);
+      if (!currentTab) {
+        console.error("Failed to create a new canvas tab");
+        return;
+      }
     }
 
     const newCanvasConfig: CanvasViewConfig = {