fix(socket): keep remote workflow operation side effects in sync

apps/sim/hooks/use-collaborative-workflow.ts

@@ -10,6 +10,7 @@ import { useUndoRedo } from '@/hooks/use-undo-redo'
 import {
   BLOCK_OPERATIONS,
   BLOCKS_OPERATIONS,
+  EDGE_OPERATIONS,
   EDGES_OPERATIONS,
   OPERATION_TARGETS,
   SUBBLOCK_OPERATIONS,
@@ -31,6 +32,38 @@ import { findAllDescendantNodes, isBlockProtected } from '@/stores/workflows/wor
 
 const logger = createLogger('CollaborativeWorkflow')
 
+function parseUndoRedoStackKey(key: string): { workflowId: string; userId: string } | null {
+  const separatorIndex = key.indexOf(':')
+
+  if (separatorIndex <= 0 || separatorIndex === key.length - 1) {
+    return null
+  }
+
+  return {
+    workflowId: key.slice(0, separatorIndex),
+    userId: key.slice(separatorIndex + 1),
+  }
+}
+
+function pruneUndoRedoStacksForWorkflow(
+  workflowId: string,
+  graph?: { blocksById: Record<string, BlockState>; edgesById: Record<string, Edge> }
+) {
+  const resolvedGraph = graph ?? {
+    blocksById: useWorkflowStore.getState().blocks,
+    edgesById: Object.fromEntries(useWorkflowStore.getState().edges.map((edge) => [edge.id, edge])),
+  }
+
+  const undoRedoStore = useUndoRedoStore.getState()
+  const stackKeys = Object.keys(undoRedoStore.stacks)
+  stackKeys.forEach((key) => {
+    const parsedKey = parseUndoRedoStackKey(key)
+    if (parsedKey?.workflowId === workflowId) {
+      undoRedoStore.pruneInvalidEntries(parsedKey.workflowId, parsedKey.userId, resolvedGraph)
+    }
+  })
+}
+
 export function useCollaborativeWorkflow() {
   const undoRedo = useUndoRedo()
   const isUndoRedoInProgress = useRef(false)
@@ -191,14 +224,110 @@ export function useCollaborativeWorkflow() {
                 .getState()
                 .setBlockCanonicalMode(payload.id, payload.canonicalId, payload.canonicalMode)
               break
+            case BLOCK_OPERATIONS.UPDATE_POSITION: {
+              if (payload.id && payload.position) {
+                // Check if this node is currently being dragged locally
+                const dragStart = useWorkflowStore.getState().getDragStartPosition()
+                if (dragStart?.id === payload.id) {
+                  logger.debug('Ignoring remote position update for locally dragged node', {
+                    nodeId: payload.id,
+                  })
+                  break
+                }
+
+                useWorkflowStore
+                  .getState()
+                  .batchUpdatePositions([{ id: payload.id, position: payload.position }])
+              }
+              break
+            }
+            case BLOCK_OPERATIONS.UPDATE_PARENT: {
+              const store = useWorkflowStore.getState()
+              const block = store.blocks[payload.id]
+              if (block) {
+                // Reconcile with local dragging
+                const dragStart = store.getDragStartPosition()
+                if (dragStart?.id === payload.id) {
+                  logger.info('Remote parent update for locally dragged node - reconciling', {
+                    nodeId: payload.id,
+                    newParentId: payload.parentId,
+                  })
+
+                  // When parent changes remotely while dragging, we update the store immediately.
+                  // React Flow's onNodeDragStop uses store data to calculate final positions.
+                  // By updating parentId and extent now, we ensure the next drag event
+                  // and the final drop use the correct coordinate space.
+                }
+
+                store.batchUpdateBlocksWithParent([
+                  {
+                    id: payload.id,
+                    position: block.position,
+                    parentId: payload.parentId || undefined,
+                  },
+                ])
+
+                if (activeWorkflowId) {
+                  pruneUndoRedoStacksForWorkflow(activeWorkflowId)
+                }
+              }
+              break
+            }
           }
         } else if (target === OPERATION_TARGETS.BLOCKS) {
           switch (operation) {
             case BLOCKS_OPERATIONS.BATCH_UPDATE_POSITIONS: {
               const { updates } = payload
               if (Array.isArray(updates)) {
-                useWorkflowStore.getState().batchUpdatePositions(updates)
+                const store = useWorkflowStore.getState()
+                const dragStart = store.getDragStartPosition()
+                const localDraggedId = dragStart?.id
+
+                // Filter out updates for blocks currently being dragged locally
+                const filteredUpdates = localDraggedId
+                  ? updates.filter((u) => u.id !== localDraggedId)
+                  : updates
+
+                if (filteredUpdates.length > 0) {
+                  store.batchUpdatePositions(filteredUpdates)
+                }
+              }
+              break
+            }
+            case BLOCKS_OPERATIONS.BATCH_UPDATE_PARENT: {
+              const { updates } = payload
+              logger.info('Received batch-update-parent from remote user', {
+                userId,
+                count: (updates || []).length,
+              })
+
+              if (updates && updates.length > 0) {
+                const store = useWorkflowStore.getState()
+                const dragStart = store.getDragStartPosition()
+                const localDraggedId = dragStart?.id
+
+                if (localDraggedId && updates.some((u: any) => u.id === localDraggedId)) {
+                  logger.info('Remote batch parent update contains locally dragged node', {
+                    localDraggedId,
+                  })
+                }
+
+                store.batchUpdateBlocksWithParent(
+                  updates.map(
+                    (u: { id: string; parentId: string; position: { x: number; y: number } }) => ({
+                      id: u.id,
+                      position: u.position,
+                      parentId: u.parentId || undefined,
+                    })
+                  )
+                )
+
+                if (activeWorkflowId) {
+                  pruneUndoRedoStacksForWorkflow(activeWorkflowId)
+                }
               }
+
+              logger.info('Successfully applied batch-update-parent from remote user')
               break
             }
           }
@@ -209,21 +338,9 @@ export function useCollaborativeWorkflow() {
               if (Array.isArray(ids) && ids.length > 0) {
                 useWorkflowStore.getState().batchRemoveEdges(ids)
 
-                const updatedBlocks = useWorkflowStore.getState().blocks
-                const updatedEdges = useWorkflowStore.getState().edges
-                const graph = {
-                  blocksById: updatedBlocks,
-                  edgesById: Object.fromEntries(updatedEdges.map((e) => [e.id, e])),
+                if (activeWorkflowId) {
+                  pruneUndoRedoStacksForWorkflow(activeWorkflowId)
                 }
-
-                const undoRedoStore = useUndoRedoStore.getState()
-                const stackKeys = Object.keys(undoRedoStore.stacks)
-                stackKeys.forEach((key) => {
-                  const [wfId, uId] = key.split(':')
-                  if (wfId === activeWorkflowId) {
-                    undoRedoStore.pruneInvalidEntries(wfId, uId, graph)
-                  }
-                })
               }
               break
             }
@@ -241,6 +358,25 @@ export function useCollaborativeWorkflow() {
               break
             }
           }
+        } else if (target === OPERATION_TARGETS.EDGE) {
+          switch (operation) {
+            case EDGE_OPERATIONS.REMOVE: {
+              if (payload.id) {
+                logger.info('Received remove-edge from remote user', {
+                  userId,
+                  edgeId: payload.id,
+                })
+                useWorkflowStore.getState().batchRemoveEdges([payload.id])
+
+                if (activeWorkflowId) {
+                  pruneUndoRedoStacksForWorkflow(activeWorkflowId)
+                }
+
+                logger.info('Successfully applied remove-edge from remote user')
+              }
+              break
+            }
+          }
         } else if (target === OPERATION_TARGETS.SUBFLOW) {
           switch (operation) {
             case SUBFLOW_OPERATIONS.UPDATE:
@@ -376,7 +512,12 @@ export function useCollaborativeWorkflow() {
               })
 
               if (ids && ids.length > 0) {
-                useWorkflowStore.getState().batchRemoveBlocks(ids)
+                const store = useWorkflowStore.getState()
+                store.batchRemoveBlocks(ids)
+
+                if (activeWorkflowId) {
+                  pruneUndoRedoStacksForWorkflow(activeWorkflowId)
+                }
               }
 
               logger.info('Successfully applied batch-remove-blocks from remote user')
@@ -424,28 +565,6 @@ export function useCollaborativeWorkflow() {
               logger.info('Successfully applied batch-toggle-locked from remote user')
               break
             }
-            case BLOCKS_OPERATIONS.BATCH_UPDATE_PARENT: {
-              const { updates } = payload
-              logger.info('Received batch-update-parent from remote user', {
-                userId,
-                count: (updates || []).length,
-              })
-
-              if (updates && updates.length > 0) {
-                useWorkflowStore.getState().batchUpdateBlocksWithParent(
-                  updates.map(
-                    (u: { id: string; parentId: string; position: { x: number; y: number } }) => ({
-                      id: u.id,
-                      position: u.position,
-                      parentId: u.parentId || undefined,
-                    })
-                  )
-                )
-              }
-
-              logger.info('Successfully applied batch-update-parent from remote user')
-              break
-            }
           }
         }
       } catch (error) {
@@ -585,20 +704,11 @@ export function useCollaborativeWorkflow() {
 
                 logger.info(`Successfully loaded reverted workflow state for ${workflowId}`)
 
-                const graph = {
+                pruneUndoRedoStacksForWorkflow(workflowId, {
                   blocksById: workflowData.state.blocks || {},
                   edgesById: Object.fromEntries(
-                    (workflowData.state.edges || []).map((e: any) => [e.id, e])
+                    (workflowData.state.edges || []).map((edge: any) => [edge.id, edge])
                   ),
-                }
-
-                const undoRedoStore = useUndoRedoStore.getState()
-                const stackKeys = Object.keys(undoRedoStore.stacks)
-                stackKeys.forEach((key) => {
-                  const [wfId, userId] = key.split(':')
-                  if (wfId === workflowId) {
-                    undoRedoStore.pruneInvalidEntries(wfId, userId, graph)
-                  }
                 })
               } finally {
                 isApplyingRemoteChange.current = false