@bian-womp/spark-workbench 0.2.28 → 0.2.30

package/lib/cjs/index.cjs

@@ -1677,7 +1677,7 @@ function useWorkbenchContext() {
     return ctx;
 }
 
-function WorkbenchProvider({ wb, runner, registry, setRegistry, overrides, children, }) {
+function WorkbenchProvider({ wb, runner, registry, setRegistry, overrides, uiVersion, children, }) {
     const [nodeStatus, setNodeStatus] = React.useState({});
     const [edgeStatus, setEdgeStatus] = React.useState({});
     const [events, setEvents] = React.useState([]);
@@ -2219,6 +2219,7 @@ function WorkbenchProvider({ wb, runner, registry, setRegistry, overrides, child
         runAutoLayout,
         updateEdgeType,
         triggerExternal,
+        uiVersion,
     }), [
         wb,
         runner,
@@ -2253,6 +2254,7 @@ function WorkbenchProvider({ wb, runner, registry, setRegistry, overrides, child
         runAutoLayout,
         wb,
         runner,
+        uiVersion,
     ]);
     return (jsxRuntime.jsx(WorkbenchContext.Provider, { value: value, children: children }));
 }
@@ -2947,7 +2949,7 @@ function NodeContextMenu({ open, clientPos, nodeId, onClose, }) {
 }
 
 const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, getDefaultNodeSize }, ref) => {
-    const { wb, registry, inputsMap, outputsMap, valuesTick, nodeStatus, edgeStatus, validationByNode, validationByEdge, } = useWorkbenchContext();
+    const { wb, registry, inputsMap, outputsMap, valuesTick, nodeStatus, edgeStatus, validationByNode, validationByEdge, uiVersion, } = useWorkbenchContext();
     const nodeValidation = validationByNode;
     const edgeValidation = validationByEdge.errors;
     // Keep stable references for nodes/edges to avoid unnecessary updates
@@ -3017,8 +3019,13 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
     const { nodeTypes, resolveNodeType } = React.useMemo(() => {
         // Build nodeTypes map using UI extension registry
         const ui = wb.getUI();
-        const custom = new Map();
-        for (const typeId of Array.from(registry.nodes.keys())) {
+        const custom = new Map(); // Include all types present in registry AND current graph to avoid timing issues
+        const def = wb.export();
+        const ids = new Set([
+            ...Array.from(registry.nodes.keys()),
+            ...def.nodes.map((n) => n.typeId),
+        ]);
+        for (const typeId of ids) {
             const renderer = ui.getNodeRenderer(typeId);
             if (renderer)
                 custom.set(typeId, renderer);
@@ -3032,8 +3039,8 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
         }
         const resolver = (nodeTypeId) => custom.has(nodeTypeId) ? `spark-${nodeTypeId}` : "spark-default";
         return { nodeTypes: types, resolveNodeType: resolver };
-        // registry is stable; ui renderers expected to be set up before mount
-    }, [wb, registry]);
+        // Include uiVersion to recompute when custom renderers are registered
+    }, [wb, registry, uiVersion]);
     const { nodes, edges } = React.useMemo(() => {
         const def = wb.export();
         const sel = wb.getSelection();
@@ -3073,7 +3080,11 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
             })
                 .map((n) => n.id);
             // Detect handle updates (ids/length changes) for targeted debug
-            const toIds = (arr) => Array.isArray(arr) ? arr.map((h) => (h && typeof h === "object" && "id" in h ? String(h.id) : "")).filter(Boolean) : [];
+            const toIds = (arr) => Array.isArray(arr)
+                ? arr
+                    .map((h) => h && typeof h === "object" && "id" in h ? String(h.id) : "")
+                    .filter(Boolean)
+                : [];
             const handlesEqual = (a, b) => {
                 const aIds = toIds(a);
                 const bIds = toIds(b);
@@ -3699,15 +3710,6 @@ function WorkbenchStudio({ engine, onEngineChange, example, onExampleChange, bac
     // Store previous runner for cleanup
     const prevRunnerRef = React.useRef(null);
     const runner = React.useMemo(() => {
-        // Dispose previous runner if it exists
-        if (prevRunnerRef.current) {
-            try {
-                prevRunnerRef.current.dispose();
-            }
-            catch (err) {
-                console.warn("Error disposing previous runner:", err);
-            }
-        }
         let newRunner;
         if (backendKind === "remote-http") {
             const backend = {
@@ -3738,29 +3740,43 @@ function WorkbenchStudio({ engine, onEngineChange, example, onExampleChange, bac
         else {
             newRunner = new LocalGraphRunner(registry);
         }
-        prevRunnerRef.current = newRunner;
         return newRunner;
     }, [registry, backendKind, httpBaseUrl, wsUrl, backendOptions]);
-    // Cleanup runner on unmount
+    // Dispose previous runner after commit; dispose current on unmount
     React.useEffect(() => {
+        const previous = prevRunnerRef.current;
+        prevRunnerRef.current = runner;
+        if (previous && previous !== runner) {
+            try {
+                previous.dispose();
+            }
+            catch (err) {
+                console.warn("Error disposing previous runner:", err);
+            }
+        }
         return () => {
-            if (prevRunnerRef.current) {
+            // Only dispose if this runner is still the current one
+            if (prevRunnerRef.current === runner) {
                 try {
-                    prevRunnerRef.current.dispose();
+                    runner.dispose();
                 }
                 catch (err) {
                     console.warn("Error disposing runner on unmount:", err);
                 }
             }
         };
-    }, []);
+    }, [runner]);
+    // Track UI registration version to trigger nodeTypes recomputation
+    const [uiVersion, setUiVersion] = React.useState(0);
     // Allow external UI registration (e.g., node renderers) with access to wb
     React.useEffect(() => {
         const baseRegisterUI = (_wb) => { };
         overrides?.registerUI?.(baseRegisterUI, { wb, wbRunner: runner });
+        // Increment UI version to trigger nodeTypes recomputation in WorkbenchCanvas
+        setUiVersion((v) => v + 1);
         // eslint-disable-next-line react-hooks/exhaustive-deps
     }, [wb, runner, overrides]);
-    return (jsxRuntime.jsx(WorkbenchProvider, { wb: wb, runner: runner, registry: registry, setRegistry: setRegistry, overrides: overrides, children: jsxRuntime.jsx(WorkbenchStudioCanvas, { setRegistry: setRegistry, autoScroll: autoScroll, onAutoScrollChange: onAutoScrollChange, example: example, onExampleChange: onExampleChange, engine: engine, onEngineChange: onEngineChange, backendKind: backendKind, onBackendKindChange: (v) => {
+    return (jsxRuntime.jsx(WorkbenchProvider, { wb: wb, runner: runner, registry: registry, setRegistry: setRegistry, overrides: overrides, uiVersion: uiVersion, children: jsxRuntime.jsx(WorkbenchStudioCanvas, { setRegistry: setRegistry, autoScroll: autoScroll, onAutoScrollChange: onAutoScrollChange, example: example, onExampleChange: onExampleChange, engine: engine, onEngineChange: onEngineChange, backendKind: backendKind, onBackendKindChange: (v) => {
                 if (runner.isRunning())
                     runner.dispose();
                 onBackendKindChange(v);