@bian-womp/spark-workbench 0.3.1 → 0.3.3

package/lib/esm/index.js

@@ -1,4 +1,4 @@
-import { generateId, createSimpleGraphRegistry, GraphBuilder, getTypedOutputValue, isTypedOutput, getInputTypeId, UnifiedEngine, getTypedOutputTypeId, isInputPrivate, offsetImportedPositions, createSimpleGraphDef, createAsyncGraphDef, createAsyncGraphRegistry, createProgressGraphDef, createProgressGraphRegistry, createValidationGraphDef, createValidationGraphRegistry } from '@bian-womp/spark-graph';
+import { generateId, createSimpleGraphRegistry, GraphBuilder, getTypedOutputValue, isTypedOutput, getInputTypeId, LocalEngine, parseJsonPath, setValueAtPathWithCreation, getTypedOutputTypeId, isInputPrivate, offsetImportedPositions, createSimpleGraphDef, createAsyncGraphDef, createAsyncGraphRegistry, createProgressGraphDef, createProgressGraphRegistry, createValidationGraphDef, createValidationGraphRegistry } from '@bian-womp/spark-graph';
 import lod from 'lodash';
 import { RemoteRuntimeClient } from '@bian-womp/spark-remote';
 import { Position, Handle, NodeResizer, getBezierPath, BaseEdge, useReactFlow, ReactFlowProvider, ReactFlow, Background, BackgroundVariant, MiniMap, Controls } from '@xyflow/react';
@@ -144,6 +144,7 @@ class InMemoryWorkbench extends AbstractWorkbench {
             edges: [],
         };
         this.nodeNames = {};
+        this.customData = {};
         this.runtimeState = null;
         this.viewport = null;
         this.historyState = undefined;
@@ -172,6 +173,17 @@ class InMemoryWorkbench extends AbstractWorkbench {
         const filteredSizes = Object.fromEntries(Object.entries(this.sizes).filter(([id]) => defNodeIds.has(id)));
         const filteredNodes = this.selection.nodes.filter((id) => defNodeIds.has(id));
         const filteredEdges = this.selection.edges.filter((id) => defEdgeIds.has(id));
+        // Clean up extData for removed nodes/edges
+        if (this.customData.nodes) {
+            const filteredExtNodes = Object.fromEntries(Object.entries(this.customData.nodes).filter(([id]) => defNodeIds.has(id)));
+            this.customData.nodes =
+                Object.keys(filteredExtNodes).length > 0 ? filteredExtNodes : undefined;
+        }
+        if (this.customData.edges) {
+            const filteredExtEdges = Object.fromEntries(Object.entries(this.customData.edges).filter(([id]) => defEdgeIds.has(id)));
+            this.customData.edges =
+                Object.keys(filteredExtEdges).length > 0 ? filteredExtEdges : undefined;
+        }
         this.positions = filteredPositions;
         this.sizes = filteredSizes;
         this.selection = { nodes: filteredNodes, edges: filteredEdges };
@@ -849,6 +861,107 @@ class InMemoryWorkbench extends AbstractWorkbench {
             ...options,
         });
     }
+    /**
+     * Get custom data for a specific node.
+     */
+    getCustomNodeData(nodeId) {
+        return this.customData.nodes?.[nodeId];
+    }
+    /**
+     * Set custom data for a specific node.
+     */
+    setCustomNodeData(nodeId, data, options) {
+        if (!this.customData.nodes) {
+            this.customData.nodes = {};
+        }
+        if (data === undefined) {
+            delete this.customData.nodes[nodeId];
+            if (Object.keys(this.customData.nodes).length === 0) {
+                delete this.customData.nodes;
+            }
+        }
+        else {
+            this.customData.nodes[nodeId] = data;
+        }
+        this.emit("graphUiChanged", {
+            change: { type: "customNodeData", nodeId, data },
+            ...options,
+        });
+    }
+    /**
+     * Get custom data for a specific edge.
+     */
+    getCustomEdgeData(edgeId) {
+        return this.customData.edges?.[edgeId];
+    }
+    /**
+     * Set custom data for a specific edge.
+     */
+    setCustomEdgeData(edgeId, data, options) {
+        if (!this.customData.edges) {
+            this.customData.edges = {};
+        }
+        if (data === undefined) {
+            delete this.customData.edges[edgeId];
+            if (Object.keys(this.customData.edges).length === 0) {
+                delete this.customData.edges;
+            }
+        }
+        else {
+            this.customData.edges[edgeId] = data;
+        }
+        this.emit("graphUiChanged", {
+            change: { type: "customEdgeData", edgeId, data },
+            ...options,
+        });
+    }
+    /**
+     * Get custom metadata.
+     */
+    getCustomMetaData() {
+        return this.customData.meta;
+    }
+    /**
+     * Set custom metadata.
+     */
+    setCustomMetaData(meta, options) {
+        if (meta === undefined) {
+            delete this.customData.meta;
+        }
+        else {
+            this.customData.meta = meta;
+        }
+        this.emit("graphUiChanged", {
+            change: { type: "customMetaData", meta },
+            ...options,
+        });
+    }
+    /**
+     * Get all custom data.
+     */
+    getCustomData() {
+        return { ...this.customData };
+    }
+    /**
+     * Set all custom data.
+     */
+    setCustomData(custom, options) {
+        if (custom === undefined) {
+            this.customData = {};
+        }
+        else {
+            this.customData = lod.pick(custom, ["nodes", "edges", "meta"]);
+        }
+        this.emit("graphUiChanged", {
+            change: {
+                type: "customData",
+                nodes: custom?.nodes,
+                edges: custom?.edges,
+                meta: custom?.meta,
+            },
+            ...options,
+        });
+    }
 }
 
 class CLIWorkbench {
@@ -909,36 +1022,14 @@ class AbstractGraphRunner {
         this.stagedInputs = {};
         this.runnerId = "";
     }
-    launch(def, opts) {
-        // Auto-stop if engine is already running
-        if (this.engine) {
-            this.stop();
-        }
-    }
     async whenIdle() {
         await this.engine?.whenIdle();
     }
-    stop() {
-        if (!this.engine)
-            return;
-        // Dispose engine (cleans up timers, listeners, etc.)
-        this.engine.dispose();
-        this.engine = undefined;
-        // Emit status but keep runtime alive
-        if (this.runMode) {
-            this.runMode = undefined;
-            this.emit("status", { running: false, runMode: undefined });
-        }
-    }
     setRunMode(runMode) {
-        if (!this.engine) {
-            throw new Error("Cannot set run mode: engine not running");
+        if (this.engine) {
+            this.engine.setRunMode(runMode);
+            this.emit("status", { running: true, runMode });
         }
-        // Update engine run mode (this will update pause/resume state)
-        this.engine.setRunMode(runMode);
-        // Update local state and emit status event
-        this.runMode = runMode;
-        this.emit("status", { running: true, runMode: this.runMode });
     }
     getInputDefaults(def) {
         const out = {};
@@ -968,17 +1059,10 @@ class AbstractGraphRunner {
         this.engine = undefined;
         this.runtime?.dispose();
         this.runtime = undefined;
-        if (this.runMode) {
-            this.runMode = undefined;
-            this.emit("status", { running: false, runMode: undefined });
-        }
     }
     isRunning() {
         return !!this.engine;
     }
-    getRunMode() {
-        return this.runMode;
-    }
     // Optional undo/redo support
     async undo() {
         return false;
@@ -997,12 +1081,12 @@ let localRunnerCounter = 0;
 class LocalGraphRunner extends AbstractGraphRunner {
     constructor(registry) {
         super(registry, { kind: "local" });
+        this.extData = {};
         this.setEnvironment = (env, opts) => {
             if (!this.runtime)
                 return;
-            const wasPaused = this.runtime.isPaused();
-            if (opts?.dry && !wasPaused)
-                this.runtime.pause();
+            // Use requestPause for dry mode to temporarily pause without affecting base run mode
+            const releasePause = opts?.dry ? this.runtime.requestPause() : null;
             try {
                 if (opts?.merge) {
                     const current = this.runtime.getEnvironment();
@@ -1014,8 +1098,7 @@ class LocalGraphRunner extends AbstractGraphRunner {
                 }
             }
             finally {
-                if (opts?.dry && !wasPaused)
-                    this.runtime.resume();
+                releasePause?.();
             }
         };
         this.getEnvironment = () => {
@@ -1036,24 +1119,21 @@ class LocalGraphRunner extends AbstractGraphRunner {
     update(def, options) {
         if (!this.runtime)
             return;
-        const wasPaused = this.runtime.isPaused();
-        // Pause runtime if dry option is set (to prevent execution) or if not paused already
-        if (options?.dry && !wasPaused) {
-            this.runtime.pause();
-        }
+        // Use requestPause for dry mode to temporarily pause without affecting base run mode
+        const releasePause = options?.dry ? this.runtime.requestPause() : null;
         try {
             this.runtime.update(def, this.registry);
             this.emit("invalidate", { reason: "graph-updated" });
         }
         finally {
-            // Resume only if we paused it due to dry option
-            if (options?.dry && !wasPaused) {
-                this.runtime.resume();
-            }
+            releasePause?.();
         }
     }
     launch(def, opts) {
-        super.launch(def, opts);
+        if (this.engine) {
+            this.engine.dispose();
+            this.engine = undefined;
+        }
         this.build(def);
         if (!this.runtime)
             throw new Error("Runtime not built");
@@ -1072,7 +1152,7 @@ class LocalGraphRunner extends AbstractGraphRunner {
         if (!this.runtime)
             throw new Error("Runtime not built");
         // Use shared engine factory
-        this.engine = new UnifiedEngine(this.runtime, opts?.runMode);
+        this.engine = new LocalEngine(this.runtime, opts?.runMode);
         if (!this.engine)
             throw new Error("Failed to create engine");
         this.engine.on("value", (e) => this.emit("value", e));
@@ -1080,8 +1160,8 @@ class LocalGraphRunner extends AbstractGraphRunner {
         this.engine.on("invalidate", (e) => this.emit("invalidate", e));
         this.engine.on("stats", (e) => this.emit("stats", e));
         this.engine.launch(opts?.invalidate);
-        this.runMode = opts?.runMode ?? "manual";
-        this.emit("status", { running: true, runMode: this.runMode });
+        const runMode = opts?.runMode ?? "manual";
+        this.emit("status", { running: true, runMode });
         for (const [nodeId, map] of Object.entries(this.stagedInputs)) {
             this.engine.setInputs(nodeId, map);
         }
@@ -1175,6 +1255,24 @@ class LocalGraphRunner extends AbstractGraphRunner {
             this.engine.copyOutputs(fromNodeId, toNodeId, options);
         }
     }
+    async setExtData(data) {
+        if (!data || typeof data !== "object") {
+            this.extData = {};
+            return;
+        }
+        this.extData = { ...this.extData, ...data };
+    }
+    async updateExtData(updates) {
+        if (!this.extData ||
+            typeof this.extData !== "object" ||
+            Array.isArray(this.extData)) {
+            this.extData = {};
+        }
+        for (const { path, value } of updates) {
+            const pathSegments = parseJsonPath(path);
+            setValueAtPathWithCreation(this.extData, pathSegments, value);
+        }
+    }
     async snapshotFull() {
         const def = undefined; // UI will supply def/positions on download for local
         const inputs = this.getInputs(this.runtime
@@ -1196,7 +1294,8 @@ class LocalGraphRunner extends AbstractGraphRunner {
             }
             : { nodes: [], edges: [] });
         const environment = this.getEnvironment() || {};
-        return { def, environment, inputs, outputs };
+        const extData = this.extData;
+        return { def, environment, inputs, outputs, extData };
     }
     async applySnapshotFull(payload, options) {
         if (payload.def && !options?.skipBuild) {
@@ -1204,6 +1303,9 @@ class LocalGraphRunner extends AbstractGraphRunner {
         }
         this.setEnvironment?.(payload.environment || {}, { merge: false });
         this.hydrate(payload, { dry: options?.dry });
+        if (payload.extData) {
+            await this.setExtData(payload.extData);
+        }
     }
     hydrate(snapshot, opts) {
         // Hydrate via runtime for exact restore (this emits events on runtime emitter)
@@ -1573,7 +1675,10 @@ class RemoteGraphRunner extends AbstractGraphRunner {
         }
     }
     launch(def, opts) {
-        super.launch(def, opts);
+        if (this.engine) {
+            this.engine.dispose();
+            this.engine = undefined;
+        }
         // Remote: build remotely then launch
         this.ensureClient().then(async (client) => {
             await client.api.build(def);
@@ -1629,8 +1734,8 @@ class RemoteGraphRunner extends AbstractGraphRunner {
             this.listenersBound = true;
         }
         this.engine = eng;
-        this.runMode = opts?.runMode ?? "manual";
-        this.emit("status", { running: true, runMode: this.runMode });
+        const runMode = opts?.runMode ?? "manual";
+        this.emit("status", { running: true, runMode });
         // Re-apply staged inputs using client.setInputs for consistency
         for (const [nodeId, map] of Object.entries(this.stagedInputs)) {
             await eng.setInputs(nodeId, map, undefined).catch(() => {
@@ -1646,7 +1751,10 @@ class RemoteGraphRunner extends AbstractGraphRunner {
      * the runtime state that was just restored.
      */
     launchExisting(def, opts) {
-        super.launch(def, opts);
+        if (this.engine) {
+            this.engine.dispose();
+            this.engine = undefined;
+        }
         // Remote: attach to existing runtime and launch (do NOT rebuild)
         this.ensureClient().then(async (client) => {
             // NOTE: We do NOT call client.build(def) here because the backend runtime
@@ -1657,14 +1765,10 @@ class RemoteGraphRunner extends AbstractGraphRunner {
         });
     }
     setRunMode(runMode) {
-        if (!this.engine) {
-            throw new Error("Cannot set run mode: engine not running");
+        if (this.engine) {
+            this.engine.setRunMode(runMode);
+            this.emit("status", { running: true, runMode });
         }
-        // Update engine run mode (sends SetRunMode command to backend)
-        this.engine.setRunMode(runMode);
-        // Update local state and emit status event
-        this.runMode = runMode;
-        this.emit("status", { running: true, runMode: this.runMode });
     }
     async computeNode(nodeId, options) {
         const client = await this.ensureClient();
@@ -1727,6 +1831,10 @@ class RemoteGraphRunner extends AbstractGraphRunner {
         const client = await this.ensureClient();
         await client.api.setExtData(data);
     }
+    async updateExtData(updates) {
+        const client = await this.ensureClient();
+        await client.api.updateExtData(updates);
+    }
     async commit(reason) {
         const client = await this.ensureClient();
         try {
@@ -2598,35 +2706,40 @@ function toReactFlow(def, positions, sizes, registry, opts) {
         }));
         const handleLayout = geom.handleLayout;
         const handles = geom.handles;
+        const baseData = {
+            typeId: n.typeId,
+            params: n.params,
+            inputHandles,
+            outputHandles,
+            inputConnected: Object.fromEntries(inputHandles.map((h) => [
+                h.id,
+                !!connectedInputs[n.nodeId]?.has(h.id),
+            ])),
+            handleLayout,
+            showValues: opts.showValues,
+            renderWidth,
+            renderHeight,
+            initialWidth: initialGeom.width,
+            initialHeight: initialGeom.height,
+            inputValues: opts.inputs?.[n.nodeId],
+            inputDefaults: opts.inputDefaults?.[n.nodeId],
+            outputValues: opts.outputs?.[n.nodeId],
+            status: opts.nodeStatus?.[n.nodeId],
+            validation: {
+                inputs: opts.nodeValidation?.inputs?.[n.nodeId] ?? [],
+                outputs: opts.nodeValidation?.outputs?.[n.nodeId] ?? [],
+                issues: opts.nodeValidation?.issues?.[n.nodeId] ?? [],
+            },
+            toString: opts.toString,
+            toElement: opts.toElement,
+        };
+        const customNodeData = opts.customData?.nodes;
+        const mergedData = customNodeData?.[n.nodeId]
+            ? { ...baseData, custom: customNodeData[n.nodeId] }
+            : baseData;
         return {
             id: n.nodeId,
-            data: {
-                typeId: n.typeId,
-                params: n.params,
-                inputHandles,
-                outputHandles,
-                inputConnected: Object.fromEntries(inputHandles.map((h) => [
-                    h.id,
-                    !!connectedInputs[n.nodeId]?.has(h.id),
-                ])),
-                handleLayout,
-                showValues: opts.showValues,
-                renderWidth,
-                renderHeight,
-                initialWidth: initialGeom.width,
-                initialHeight: initialGeom.height,
-                inputValues: opts.inputs?.[n.nodeId],
-                inputDefaults: opts.inputDefaults?.[n.nodeId],
-                outputValues: opts.outputs?.[n.nodeId],
-                status: opts.nodeStatus?.[n.nodeId],
-                validation: {
-                    inputs: opts.nodeValidation?.inputs?.[n.nodeId] ?? [],
-                    outputs: opts.nodeValidation?.outputs?.[n.nodeId] ?? [],
-                    issues: opts.nodeValidation?.issues?.[n.nodeId] ?? [],
-                },
-                toString: opts.toString,
-                toElement: opts.toElement,
-            },
+            data: mergedData,
             position: positions[n.nodeId] ?? { x: 0, y: 0 },
             type: opts.resolveNodeType?.(n.typeId) ?? "spark-default",
             selected: opts.selectedNodeIds
@@ -2669,6 +2782,25 @@ function toReactFlow(def, positions, sizes, registry, opts) {
         const targetHandleTypeId = targetHandles?.inputs[e.target.handle]
             ? getInputTypeId(targetHandles.inputs, e.target.handle) ?? "unknown"
             : "unknown";
+        const baseEdgeData = {
+            sourceNodeId: e.source.nodeId,
+            sourceNodeTypeId: sourceNode?.typeId || "unknown",
+            sourceHandle: e.source.handle,
+            sourceHandleTypeId,
+            targetNodeId: e.target.nodeId,
+            targetNodeTypeId: targetNode?.typeId || "unknown",
+            targetHandle: e.target.handle,
+            targetHandleTypeId,
+            edgeTypeId,
+            isRunning,
+            hasError,
+            isInvalid: isInvalidEdge,
+            isMissing,
+        };
+        const customEdgeData = opts.customData?.edges;
+        const mergedEdgeData = customEdgeData?.[e.id]
+            ? { ...baseEdgeData, custom: customEdgeData[e.id] }
+            : baseEdgeData;
         return {
             id: e.id,
             source: e.source.nodeId,
@@ -2682,21 +2814,7 @@ function toReactFlow(def, positions, sizes, registry, opts) {
             style,
             label: edgeTypeId,
             type: "default",
-            data: {
-                sourceNodeId: e.source.nodeId,
-                sourceNodeTypeId: sourceNode?.typeId || "unknown",
-                sourceHandle: e.source.handle,
-                sourceHandleTypeId,
-                targetNodeId: e.target.nodeId,
-                targetNodeTypeId: targetNode?.typeId || "unknown",
-                targetHandle: e.target.handle,
-                targetHandleTypeId,
-                edgeTypeId,
-                isRunning,
-                hasError,
-                isInvalid: isInvalidEdge,
-                isMissing,
-            },
+            data: mergedEdgeData,
         };
     });
     return { nodes, edges };
@@ -2852,6 +2970,9 @@ async function upload(parsed, wb, runner) {
     if (extData.runtime && typeof extData.runtime === "object") {
         wb.setRuntimeState(extData.runtime);
     }
+    if (extData.custom && typeof extData.custom === "object") {
+        wb.setCustomData(extData.custom);
+    }
     if (runner.isRunning()) {
         await runner.applySnapshotFull({
             def: wb.def,
@@ -2868,6 +2989,9 @@ async function upload(parsed, wb, runner) {
                 runner.setInputs(nodeId, map, { dry: true });
             }
         }
+        if (extData) {
+            await runner.setExtData(extData);
+        }
     }
 }
 
@@ -3664,15 +3788,7 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
     const graphUiTick = useWorkbenchGraphUiTick(wb);
     const versionTick = useWorkbenchVersionTick(runner);
     const valuesTick = versionTick + graphTick + graphUiTick;
-    // Keep local runMode state loosely in sync with runner status.
-    // - Seed from runner.getRunMode() on mount if available.
-    // - On status events, update only when a non-undefined runMode is reported,
-    //   so the UI preserves the last selected mode after stop().
     useEffect(() => {
-        const initialMode = runner.getRunMode();
-        if (initialMode) {
-            setRunModeState(initialMode);
-        }
         const offRunnerStatus = runner.on("status", (status) => {
             if (status.runMode) {
                 setRunModeState(status.runMode);
@@ -3832,12 +3948,13 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
             workbench.setRuntimeState(metadata);
             const fullUiState = workbench.getUIState();
             const uiWithoutViewport = excludeViewportFromUIState(fullUiState);
-            await graphRunner.setExtData?.({
-                ...(Object.keys(uiWithoutViewport || {}).length > 0
-                    ? { ui: uiWithoutViewport }
-                    : {}),
-                runtime: metadata,
-            });
+            // Use updateExtData for efficient batched partial updates
+            const updates = [];
+            if (Object.keys(uiWithoutViewport || {}).length > 0) {
+                updates.push({ path: "ui", value: uiWithoutViewport });
+            }
+            updates.push({ path: "runtime", value: metadata });
+            await graphRunner.updateExtData(updates);
         }
         catch (err) {
             console.warn("[WorkbenchContext] Failed to save runtime metadata:", err);
@@ -4270,6 +4387,18 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
                     else if (changeType === "selection") {
                         reason = "selection";
                     }
+                    else if (changeType === "customNodeData") {
+                        reason = "custom-node-data";
+                    }
+                    else if (changeType === "customEdgeData") {
+                        reason = "custom-edge-data";
+                    }
+                    else if (changeType === "customMetaData") {
+                        reason = "custom-meta";
+                    }
+                    else if (changeType === "customData") {
+                        reason = "custom";
+                    }
                 }
                 await saveUiRuntimeMetadata(wb, runner);
                 const history = await runner
@@ -4375,27 +4504,11 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
         };
     }, [runner, wb]);
     const isRunning = useCallback(() => runner.isRunning(), [runner]);
-    const getRunMode = useCallback(() => runner.getRunMode(), [runner]);
-    const stop = useCallback(() => runner.stop(), [runner]);
-    // Run mode actions
     const setRunMode = useCallback((mode) => {
         if (mode === runMode)
             return;
-        const wasRunning = runner.isRunning();
-        if (wasRunning) {
-            // Use setRunMode to change run mode without rebuilding
-            try {
-                runner.setRunMode(mode);
-                setRunModeState(mode);
-            }
-            catch (err) {
-                console.error("Failed to set run mode:", err);
-            }
-        }
-        else {
-            // Just update state if not running (will be applied on next launch)
-            setRunModeState(mode);
-        }
+        runner.setRunMode(mode);
+        setRunModeState(mode);
     }, [runMode, runner]);
     const runNodeAction = useCallback(async (nodeId) => {
         await runner.computeNode(nodeId);
@@ -4497,8 +4610,6 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
         removeRegistryError,
         removeInputValidationError,
         isRunning,
-        getRunMode,
-        stop,
         runMode,
         setRunMode,
         runNode: runNodeAction,
@@ -4539,8 +4650,6 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
         events,
         clearEvents,
         isRunning,
-        getRunMode,
-        stop,
         runMode,
         setRunMode,
         runNodeAction,
@@ -5574,6 +5683,7 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
                 status: n.data.status,
                 validation: n.data.validation,
                 inputConnected: n.data.inputConnected,
+                custom: n.data.custom,
             },
         });
         return lod.isEqual(pick(a), pick(b));
@@ -5589,6 +5699,14 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
             style: e.style,
             label: e.label,
             type: e.type,
+            data: e.data && {
+                edgeTypeId: e.data.edgeTypeId,
+                isRunning: e.data.isRunning,
+                hasError: e.data.hasError,
+                isInvalid: e.data.isInvalid,
+                isMissing: e.data.isMissing,
+                custom: e.data.custom,
+            },
         });
         return lod.isEqual(pick(a), pick(b));
     };
@@ -5676,6 +5794,7 @@ const WorkbenchCanvas = React.forwardRef(({ showValues, toString, toElement, get
             selectedEdgeIds: new Set(sel.edges),
             getDefaultNodeSize,
             ui,
+            customData: wb.getCustomData(),
         });
         // Retain references for unchanged items
         const stableNodes = retainStabilityById(prevNodesRef.current, out.nodes, isSameNode);
@@ -6223,14 +6342,14 @@ function WorkbenchStudioCanvas({ autoScroll, onAutoScrollChange, example, onExam
         const isConnecting = transportStatus.state === "connecting" ||
             transportStatus.state === "retrying";
         // Only allow Start/Stop when transport is connected or local
+        // For local backend, always allow control (transport state is "local")
+        // For remote backends, require connection
         const canControl = transportStatus.state === "connected" ||
-            transportStatus.state === "local";
+            transportStatus.state === "local" ||
+            backendKind === "local"; // Always allow control for local backend
         if (isConnecting) {
             return (jsxs("button", { className: "border rounded px-2 py-1.5 text-gray-500 border-gray-400 flex items-center gap-1 disabled:opacity-50", disabled: true, title: "Connecting to backend...", children: [jsx(ClockClockwiseIcon, { size: 16, className: "animate-spin" }), jsx("span", { className: "font-medium ml-1", children: "Connecting..." })] }));
         }
-        if (isGraphRunning) {
-            return (jsxs("button", { className: "border rounded px-2 py-1.5 text-red-700 border-red-600 flex items-center gap-1 disabled:opacity-50 disabled:text-gray-400 disabled:border-gray-300", onClick: () => runner.stop(), disabled: !canControl, title: canControl ? "Stop engine" : "Waiting for connection", children: [jsx(StopIcon, { size: 16, weight: "fill" }), jsx("span", { className: "font-medium ml-1", children: "Stop" })] }));
-        }
         return (jsxs("button", { className: "border rounded px-2 py-1.5 text-green-700 border-green-600 flex items-center gap-1 disabled:text-gray-400 disabled:border-gray-300 disabled:opacity-50", onClick: (evt) => {
                 if (evt.shiftKey && !confirm("Invalidate and re-run graph?"))
                     return;
@@ -6247,7 +6366,7 @@ function WorkbenchStudioCanvas({ autoScroll, onAutoScrollChange, example, onExam
             }, disabled: !canControl, title: !canControl
                 ? "Waiting for connection"
                 : `Start ${runMode === "manual" ? "manual" : "auto"} mode`, children: [jsx(PlayIcon, { size: 16, weight: "fill" }), jsx("span", { className: "font-medium ml-1", children: "Start" })] }));
-    }, [transportStatus, isGraphRunning, runner, runMode, wb]);
+    }, [transportStatus, isGraphRunning, runner, runMode, wb, backendKind]);
     const defaultExamples = useMemo(() => [
         {
             id: "simple",
@@ -6617,9 +6736,7 @@ function WorkbenchStudioCanvas({ autoScroll, onAutoScrollChange, example, onExam
                             if (mode !== runMode) {
                                 await setRunMode(mode);
                             }
-                        }, disabled: isGraphRunning, title: isGraphRunning
-                            ? "Stop before switching run mode"
-                            : "Select run mode", children: [jsx("option", { value: "manual", children: "Manual" }), jsx("option", { value: "auto", children: "Auto" })] }), renderStartStopButton(), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: runAutoLayout, children: jsx(TreeStructureIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: () => canvasRef.current?.fitView?.(), title: "Fit View", children: jsx(CornersOutIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: download$1, children: jsx(DownloadIcon, { size: 24 }) }), jsx("input", { ref: uploadInputRef, type: "file", accept: "application/json,.json", className: "hidden", onChange: onUploadPicked }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: triggerUpload, children: jsx(UploadIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: async () => {
+                        }, title: "Select run mode", children: [jsx("option", { value: "manual", children: "Manual" }), jsx("option", { value: "auto", children: "Auto" })] }), renderStartStopButton(), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: runAutoLayout, children: jsx(TreeStructureIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: () => canvasRef.current?.fitView?.(), title: "Fit View", children: jsx(CornersOutIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: download$1, children: jsx(DownloadIcon, { size: 24 }) }), jsx("input", { ref: uploadInputRef, type: "file", accept: "application/json,.json", className: "hidden", onChange: onUploadPicked }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: triggerUpload, children: jsx(UploadIcon, { size: 24 }) }), jsx("button", { className: "border border-gray-300 rounded p-1", onClick: async () => {
                             await downloadCanvasThumbnail(canvasContainerRef.current);
                         }, title: "Download Flow Thumbnail (SVG)", children: jsx(ImageIcon, { size: 24 }) }), jsxs("label", { className: "flex items-center gap-1", children: [jsx("input", { type: "checkbox", checked: debug, onChange: (e) => onDebugChange(e.target.checked) }), jsx(BugBeetleIcon, { size: 24, weight: debug ? "fill" : undefined })] }), jsxs("label", { className: "flex items-center gap-1", children: [jsx("input", { type: "checkbox", checked: showValues, onChange: (e) => onShowValuesChange(e.target.checked) }), jsx(ListBulletsIcon, { size: 24, weight: showValues ? "fill" : undefined })] })] }), jsxs("div", { className: "flex flex-1 min-h-0", children: [jsx("div", { className: "flex-1 min-w-0", ref: canvasContainerRef, children: jsx(WorkbenchCanvas, { ref: canvasRef, showValues: showValues, toString: toString, toElement: toElement, getDefaultNodeSize: overrides?.getDefaultNodeSize }) }), jsx(Inspector, { setInput: setInput, debug: debug, autoScroll: autoScroll, hideWorkbench: hideWorkbench, onAutoScrollChange: onAutoScrollChange, onHideWorkbenchChange: onHideWorkbenchChange, toString: toString, contextPanel: overrides?.contextPanel })] })] }));
 }