From 34e231230f36f29373aa34a71cb79ed1b46ec5a1 Mon Sep 17 00:00:00 2001
From: Aaryan Khandelwal <65252264+aaryan610@users.noreply.github.com>
Date: Tue, 26 Aug 2025 02:13:27 +0530
Subject: [PATCH] [WIKI-498] regression: table bugs #7631

---
 .../plugins/drag-handles/column/plugin.ts     | 34 +++++++++++++----
 .../table/plugins/drag-handles/row/plugin.ts  | 38 ++++++++++++++-----
 .../editor/src/core/plugins/drag-handle.ts    |  4 +-
 3 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/packages/editor/src/core/extensions/table/plugins/drag-handles/column/plugin.ts b/packages/editor/src/core/extensions/table/plugins/drag-handles/column/plugin.ts
index 6cefa24b6..d25591e47 100644
--- a/packages/editor/src/core/extensions/table/plugins/drag-handles/column/plugin.ts
+++ b/packages/editor/src/core/extensions/table/plugins/drag-handles/column/plugin.ts
@@ -14,6 +14,9 @@ import { ColumnDragHandle, ColumnDragHandleProps } from "./drag-handle";
 
 type TableColumnDragHandlePluginState = {
   decorations?: DecorationSet;
+  // track table structure to detect changes
+  tableWidth?: number;
+  tableNodePos?: number;
 };
 
 const TABLE_COLUMN_DRAG_HANDLE_PLUGIN_KEY = new PluginKey("tableColumnHandlerDecorationPlugin");
@@ -31,20 +34,33 @@ export const TableColumnDragHandlePlugin = (editor: Editor): Plugin<TableColumnD
 
         const tableMap = TableMap.get(table.node);
 
-        let isStale = false;
-        const mapped = prev.decorations?.map(tr.mapping, tr.doc);
-        for (let col = 0; col < tableMap.width; col++) {
-          const pos = getTableCellWidgetDecorationPos(table, tableMap, col);
-          if (mapped?.find(pos, pos + 1)?.length !== 1) {
-            isStale = true;
-            break;
+        // Check if table structure changed (width or position)
+        const tableStructureChanged = prev.tableWidth !== tableMap.width || prev.tableNodePos !== table.pos;
+
+        let isStale = tableStructureChanged;
+
+        // Only do position-based stale check if structure hasn't changed
+        if (!isStale) {
+          const mapped = prev.decorations?.map(tr.mapping, tr.doc);
+          for (let col = 0; col < tableMap.width; col++) {
+            const pos = getTableCellWidgetDecorationPos(table, tableMap, col);
+            if (mapped?.find(pos, pos + 1)?.length !== 1) {
+              isStale = true;
+              break;
+            }
           }
         }
 
         if (!isStale) {
-          return { decorations: mapped };
+          const mapped = prev.decorations?.map(tr.mapping, tr.doc);
+          return {
+            decorations: mapped,
+            tableWidth: tableMap.width,
+            tableNodePos: table.pos,
+          };
         }
 
+        // recreate all decorations
         const decorations: Decoration[] = [];
 
         for (let col = 0; col < tableMap.width; col++) {
@@ -63,6 +79,8 @@ export const TableColumnDragHandlePlugin = (editor: Editor): Plugin<TableColumnD
 
         return {
           decorations: DecorationSet.create(newState.doc, decorations),
+          tableWidth: tableMap.width,
+          tableNodePos: table.pos,
         };
       },
     },
diff --git a/packages/editor/src/core/extensions/table/plugins/drag-handles/row/plugin.ts b/packages/editor/src/core/extensions/table/plugins/drag-handles/row/plugin.ts
index d84f942ba..92cf0eea9 100644
--- a/packages/editor/src/core/extensions/table/plugins/drag-handles/row/plugin.ts
+++ b/packages/editor/src/core/extensions/table/plugins/drag-handles/row/plugin.ts
@@ -14,6 +14,9 @@ import { RowDragHandle, RowDragHandleProps } from "./drag-handle";
 
 type TableRowDragHandlePluginState = {
   decorations?: DecorationSet;
+  // track table structure to detect changes
+  tableHeight?: number;
+  tableNodePos?: number;
 };
 
 const TABLE_ROW_DRAG_HANDLE_PLUGIN_KEY = new PluginKey("tableRowDragHandlePlugin");
@@ -31,20 +34,33 @@ export const TableRowDragHandlePlugin = (editor: Editor): Plugin<TableRowDragHan
 
         const tableMap = TableMap.get(table.node);
 
-        let isStale = false;
-        const mapped = prev.decorations?.map(tr.mapping, tr.doc);
-        for (let row = 0; row < tableMap.height; row++) {
-          const pos = getTableCellWidgetDecorationPos(table, tableMap, row * tableMap.width);
-          if (mapped?.find(pos, pos + 1)?.length !== 1) {
-            isStale = true;
-            break;
+        // Check if table structure changed (height or position)
+        const tableStructureChanged = prev.tableHeight !== tableMap.height || prev.tableNodePos !== table.pos;
+
+        let isStale = tableStructureChanged;
+
+        // Only do position-based stale check if structure hasn't changed
+        if (!isStale) {
+          const mapped = prev.decorations?.map(tr.mapping, tr.doc);
+          for (let row = 0; row < tableMap.height; row++) {
+            const pos = getTableCellWidgetDecorationPos(table, tableMap, row * tableMap.width);
+            if (mapped?.find(pos, pos + 1)?.length !== 1) {
+              isStale = true;
+              break;
+            }
           }
         }
 
         if (!isStale) {
-          return { decorations: mapped };
+          const mapped = prev.decorations?.map(tr.mapping, tr.doc);
+          return {
+            decorations: mapped,
+            tableHeight: tableMap.height,
+            tableNodePos: table.pos,
+          };
         }
 
+        // recreate all decorations
         const decorations: Decoration[] = [];
 
         for (let row = 0; row < tableMap.height; row++) {
@@ -61,7 +77,11 @@ export const TableRowDragHandlePlugin = (editor: Editor): Plugin<TableRowDragHan
           decorations.push(Decoration.widget(pos, () => dragHandleComponent.element));
         }
 
-        return { decorations: DecorationSet.create(newState.doc, decorations) };
+        return {
+          decorations: DecorationSet.create(newState.doc, decorations),
+          tableHeight: tableMap.height,
+          tableNodePos: table.pos,
+        };
       },
     },
     props: {
diff --git a/packages/editor/src/core/plugins/drag-handle.ts b/packages/editor/src/core/plugins/drag-handle.ts
index 5d05f80fc..b525759d9 100644
--- a/packages/editor/src/core/plugins/drag-handle.ts
+++ b/packages/editor/src/core/plugins/drag-handle.ts
@@ -379,7 +379,9 @@ const handleNodeSelection = (
   let draggedNodePos = nodePosAtDOM(node, view, options);
   if (draggedNodePos == null || draggedNodePos < 0) return;
 
-  if (node.matches("blockquote")) {
+  if (node.matches("table")) {
+    draggedNodePos = draggedNodePos - 2;
+  } else if (node.matches("blockquote")) {
     draggedNodePos = nodePosAtDOMForBlockQuotes(node, view);
     if (draggedNodePos === null || draggedNodePos === undefined) return;
   } else {