fix: Grips on merged table cells (#11003)

shared/editor/commands/table.ts

@@ -588,8 +588,8 @@ export function selectRow(index: number, expand = false): Command {
       const $pos = state.doc.resolve(rect.tableStart + pos);
       const rowSelection =
         expand && state.selection instanceof CellSelection
-          ? RowSelection.rowSelection(state.selection.$anchorCell, $pos)
-          : RowSelection.rowSelection($pos);
+          ? RowSelection.rowSelection(state.selection.$anchorCell, $pos, index)
+          : RowSelection.rowSelection($pos, $pos, index);
       dispatch(state.tr.setSelection(rowSelection));
       return true;
     }

shared/editor/components/Styles.ts

@@ -1704,14 +1704,22 @@ table {
   tr:first-child td {
     border-top: 0;
   }
-  tr:first-child th:first-child,
-  tr:first-child td:first-child {
+  tr:first-child th[data-first-column],
+  tr:first-child td[data-first-column] {
     border-radius: ${EditorStyleHelper.blockRadius} 0 0 0;
   }
-  tr:last-child th:first-child,
-  tr:last-child td:first-child {
+  th[data-first-column][data-last-row],
+  td[data-first-column][data-last-row] {
     border-radius: 0 0 0 ${EditorStyleHelper.blockRadius};
   }
+  tr:first-child th[data-last-column],
+  tr:first-child td[data-last-column] {
+    border-radius: 0 ${EditorStyleHelper.blockRadius} 0 0;
+  }
+  th[data-last-column][data-last-row],
+  td[data-last-column][data-last-row] {
+    border-radius: 0 0 ${EditorStyleHelper.blockRadius} 0;
+  }
 
   td .component-embed {
     padding: 4px 0;
@@ -1881,15 +1889,16 @@ table {
       border-top-left-radius: 3px;
       border-bottom-left-radius: 3px;
     }
-    &.last::after {
-      border-top-right-radius: 3px;
-      border-bottom-right-radius: 3px;
-    }
     &.selected::after {
       background: ${props.theme.tableSelected};
     }
   }
 
+  [data-last-column] .${EditorStyleHelper.tableGripColumn}::after {
+    border-top-right-radius: 3px;
+    border-bottom-right-radius: 3px;
+  }
+
   .${EditorStyleHelper.tableGripRow} {
     &::after {
       content: "";
@@ -1911,15 +1920,16 @@ table {
       border-top-left-radius: 3px;
       border-top-right-radius: 3px;
     }
-    &.last::after {
-      border-bottom-left-radius: 3px;
-      border-bottom-right-radius: 3px;
-    }
     &.selected::after {
       background: ${props.theme.tableSelected};
     }
   }
 
+  [data-last-row] .${EditorStyleHelper.tableGripRow}::after {
+    border-bottom-left-radius: 3px;
+    border-bottom-right-radius: 3px;
+  }
+
   .${EditorStyleHelper.tableGrip} {
     &::after {
       content: "";

shared/editor/nodes/TableCell.ts

@@ -1,19 +1,14 @@
 import type { Token } from "markdown-it";
-import type { NodeSpec } from "prosemirror-model";
-import { Slice } from "prosemirror-model";
+import {
+  type Node as ProsemirrorNode,
+  type NodeSpec,
+  Slice,
+} from "prosemirror-model";
 import type { EditorState } from "prosemirror-state";
 import { Plugin, PluginKey } from "prosemirror-state";
-import { DecorationSet, Decoration } from "prosemirror-view";
-import { addRowBefore, selectRow, selectTable } from "../commands/table";
+import { Decoration, DecorationSet } from "prosemirror-view";
+import { TableMap } from "prosemirror-tables";
 import { getCellAttrs, setCellAttrs } from "../lib/table";
-import {
-  getCellsInColumn,
-  getRowIndexInMap,
-  isRowSelected,
-  isTableSelected,
-} from "../queries/table";
-import { EditorStyleHelper } from "../styles/EditorStyleHelper";
-import { cn } from "../styles/utils";
 import Node from "./Node";
 
 export default class TableCell extends Node {
@@ -51,108 +46,85 @@ export default class TableCell extends Node {
   }
 
   get plugins() {
-    function buildAddRowDecoration(pos: number, index: number) {
-      const className = cn(EditorStyleHelper.tableAddRow, {
-        first: index === 0,
-      });
-
-      return Decoration.widget(
-        pos + 1,
-        () => {
-          const plus = document.createElement("a");
-          plus.role = "button";
-          plus.className = className;
-          plus.dataset.index = index.toString();
-          return plus;
-        },
-        {
-          key: cn(className, index),
-        }
-      );
-    }
-
-    const createRowDecorations = (state: EditorState) => {
-      if (!this.editor.view?.editable) {
-        return DecorationSet.empty;
-      }
-
+    const createCellDecorations = (state: EditorState) => {
       const { doc } = state;
       const decorations: Decoration[] = [];
-      const rows = getCellsInColumn(0)(state);
 
-      if (rows) {
-        rows.forEach((pos, visualIndex) => {
-          const actualRowIndex = getRowIndexInMap(visualIndex, state);
-          const index = actualRowIndex !== -1 ? actualRowIndex : visualIndex;
+      // Iterate through all tables in the document
+      doc.descendants((node: ProsemirrorNode, pos: number) => {
+        if (node.type.spec.tableRole === "table") {
+          const map = TableMap.get(node);
 
-          if (index === 0) {
-            const className = cn(EditorStyleHelper.tableGrip, {
-              selected: isTableSelected(state),
-            });
+          // Mark cells in the first column and last row of this table
+          node.descendants((cellNode: ProsemirrorNode, cellPos: number) => {
+            if (
+              cellNode.type.spec.tableRole === "cell" ||
+              cellNode.type.spec.tableRole === "header_cell"
+            ) {
+              const cellOffset = cellPos;
+              const cellIndex = map.map.indexOf(cellOffset);
 
-            decorations.push(
-              Decoration.widget(
-                pos + 1,
-                () => {
-                  const grip = document.createElement("a");
-                  grip.role = "button";
-                  grip.className = className;
-                  return grip;
-                },
-                {
-                  key: className,
+              if (cellIndex !== -1) {
+                const col = cellIndex % map.width;
+                const row = Math.floor(cellIndex / map.width);
+                const rowspan = cellNode.attrs.rowspan || 1;
+                const colspan = cellNode.attrs.colspan || 1;
+                const attrs: Record<string, string> = {};
+
+                if (col === 0) {
+                  attrs["data-first-column"] = "true";
                 }
-              )
-            );
-          }
 
-          const className = cn(EditorStyleHelper.tableGripRow, {
-            selected: isRowSelected(index)(state) || isTableSelected(state),
-            first: index === 0,
-            last: visualIndex === rows.length - 1,
-          });
+                // Mark cells that extend into the last column (accounting for colspan)
+                if (col + colspan >= map.width) {
+                  attrs["data-last-column"] = "true";
+                }
 
-          decorations.push(
-            Decoration.widget(
-              pos + 1,
-              () => {
-                const grip = document.createElement("a");
-                grip.role = "button";
-                grip.className = className;
-                grip.dataset.index = index.toString();
-                return grip;
-              },
-              {
-                key: cn(className, index),
+                // Mark cells that extend into the last row (accounting for rowspan)
+                if (row + rowspan >= map.height) {
+                  attrs["data-last-row"] = "true";
+                }
+
+                if (Object.keys(attrs).length > 0) {
+                  decorations.push(
+                    Decoration.node(
+                      pos + cellPos + 1,
+                      pos + cellPos + 1 + cellNode.nodeSize,
+                      attrs
+                    )
+                  );
+                }
               }
-            )
-          );
-
-          if (index === 0) {
-            decorations.push(buildAddRowDecoration(pos, index));
-          }
-
-          decorations.push(buildAddRowDecoration(pos, index + 1));
-        });
-      }
+            }
+          });
+        }
+      });
 
       return DecorationSet.create(doc, decorations);
     };
 
     return [
       new Plugin({
-        key: new PluginKey("table-cell-decorations"),
+        key: new PluginKey("table-cell-attributes"),
         state: {
-          init: (_, state) => createRowDecorations(state),
+          init: (_, state) => createCellDecorations(state),
           apply: (tr, pluginState, oldState, newState) => {
-            // Only recompute if selection or document changed
-            if (!tr.selectionSet && !tr.docChanged) {
+            // Only recompute if document changed
+            if (!tr.docChanged) {
               return pluginState;
             }
 
-            return createRowDecorations(newState);
+            return createCellDecorations(newState);
           },
         },
+        props: {
+          decorations(state) {
+            return this.getState(state);
+          },
+        },
+      }),
+      new Plugin({
+        key: new PluginKey("table-cell-copy-transform"),
         props: {
           transformCopied: (slice) => {
             // check if the copied selection is a single table, with a single row, with a single cell. If so,
@@ -183,54 +155,6 @@ export default class TableCell extends Node {
 
             return slice;
           },
-          handleDOMEvents: {
-            mousedown: (view, event) => {
-              if (!(event.target instanceof HTMLElement)) {
-                return false;
-              }
-
-              const targetAddRow = event.target.closest(
-                `.${EditorStyleHelper.tableAddRow}`
-              );
-              if (targetAddRow) {
-                event.preventDefault();
-                event.stopImmediatePropagation();
-                const index = Number(targetAddRow.getAttribute("data-index"));
-
-                addRowBefore({ index })(view.state, view.dispatch);
-                return true;
-              }
-
-              const targetGrip = event.target.closest(
-                `.${EditorStyleHelper.tableGrip}`
-              );
-              if (targetGrip) {
-                event.preventDefault();
-                event.stopImmediatePropagation();
-                selectTable()(view.state, view.dispatch);
-                return true;
-              }
-
-              const targetGripRow = event.target.closest(
-                `.${EditorStyleHelper.tableGripRow}`
-              );
-              if (targetGripRow) {
-                event.preventDefault();
-                event.stopImmediatePropagation();
-
-                selectRow(
-                  Number(targetGripRow.getAttribute("data-index")),
-                  event.metaKey || event.shiftKey
-                )(view.state, view.dispatch);
-                return true;
-              }
-
-              return false;
-            },
-          },
-          decorations(state) {
-            return this.getState(state);
-          },
         },
       }),
     ];

shared/editor/nodes/TableHeader.ts

@@ -4,6 +4,7 @@ import type { EditorState } from "prosemirror-state";
 import { Plugin, PluginKey } from "prosemirror-state";
 import type { EditorView } from "prosemirror-view";
 import { DecorationSet, Decoration } from "prosemirror-view";
+import { TableMap } from "prosemirror-tables";
 import { addColumnBefore, selectColumn } from "../commands/table";
 import { getCellAttrs, setCellAttrs } from "../lib/table";
 import {
@@ -114,7 +115,80 @@ export default class TableHeader extends Node {
       return DecorationSet.create(doc, decorations);
     };
 
+    const createHeaderDecorations = (state: EditorState) => {
+      const { doc } = state;
+      const decorations: Decoration[] = [];
+
+      // Iterate through all tables in the document
+      doc.descendants((node, pos) => {
+        if (node.type.spec.tableRole === "table") {
+          const map = TableMap.get(node);
+
+          // Mark cells in the first column and last row of this table
+          node.descendants((cellNode, cellPos) => {
+            if (cellNode.type.spec.tableRole === "header_cell") {
+              const cellOffset = cellPos;
+              const cellIndex = map.map.indexOf(cellOffset);
+
+              if (cellIndex !== -1) {
+                const col = cellIndex % map.width;
+                const row = Math.floor(cellIndex / map.width);
+                const rowspan = cellNode.attrs.rowspan || 1;
+                const colspan = cellNode.attrs.colspan || 1;
+                const attrs: Record<string, string> = {};
+
+                if (col === 0) {
+                  attrs["data-first-column"] = "true";
+                }
+
+                // Mark cells that extend into the last column (accounting for colspan)
+                if (col + colspan >= map.width) {
+                  attrs["data-last-column"] = "true";
+                }
+
+                // Mark cells that extend into the last row (accounting for rowspan)
+                if (row + rowspan >= map.height) {
+                  attrs["data-last-row"] = "true";
+                }
+
+                if (Object.keys(attrs).length > 0) {
+                  decorations.push(
+                    Decoration.node(
+                      pos + cellPos + 1,
+                      pos + cellPos + 1 + cellNode.nodeSize,
+                      attrs
+                    )
+                  );
+                }
+              }
+            }
+          });
+        }
+      });
+
+      return DecorationSet.create(doc, decorations);
+    };
+
     return [
+      new Plugin({
+        key: new PluginKey("table-header-first-column"),
+        state: {
+          init: (_, state) => createHeaderDecorations(state),
+          apply: (tr, pluginState, oldState, newState) => {
+            // Only recompute if document changed
+            if (!tr.docChanged) {
+              return pluginState;
+            }
+
+            return createHeaderDecorations(newState);
+          },
+        },
+        props: {
+          decorations(state) {
+            return this.getState(state);
+          },
+        },
+      }),
       new Plugin({
         key: new PluginKey("table-header-decorations"),
         state: {

shared/editor/nodes/TableRow.ts

@@ -1,5 +1,16 @@
 import type { NodeSpec } from "prosemirror-model";
+import { isInTable, selectedRect } from "prosemirror-tables";
 import Node from "./Node";
+import { cn } from "../styles/utils";
+import { EditorStyleHelper } from "../styles/EditorStyleHelper";
+import { Decoration, DecorationSet } from "prosemirror-view";
+import { Plugin } from "prosemirror-state";
+import { addRowBefore, selectRow, selectTable } from "../commands/table";
+import {
+  getRowsInTable,
+  isRowSelected,
+  isTableSelected,
+} from "../queries/table";
 
 export default class TableRow extends Node {
   get name() {
@@ -17,6 +28,173 @@ export default class TableRow extends Node {
     };
   }
 
+  get plugins() {
+    function buildAddRowDecoration(pos: number, index: number) {
+      const className = cn(EditorStyleHelper.tableAddRow, {
+        first: index === 0,
+      });
+
+      return Decoration.widget(
+        pos + 1,
+        () => {
+          const plus = document.createElement("a");
+          plus.role = "button";
+          plus.className = className;
+          plus.dataset.index = index.toString();
+          return plus;
+        },
+        {
+          key: cn(className, index),
+        }
+      );
+    }
+
+    return [
+      new Plugin({
+        props: {
+          handleDOMEvents: {
+            mousedown: (view, event) => {
+              if (!(event.target instanceof HTMLElement)) {
+                return false;
+              }
+
+              const targetAddRow = event.target.closest(
+                `.${EditorStyleHelper.tableAddRow}`
+              );
+              if (targetAddRow) {
+                event.preventDefault();
+                event.stopImmediatePropagation();
+                const index = Number(targetAddRow.getAttribute("data-index"));
+
+                addRowBefore({ index })(view.state, view.dispatch);
+                return true;
+              }
+
+              const tableGrip = event.target.closest(
+                `.${EditorStyleHelper.tableGrip}`
+              );
+              if (tableGrip) {
+                event.preventDefault();
+                event.stopImmediatePropagation();
+                selectTable()(view.state, view.dispatch);
+                return true;
+              }
+
+              const targetGripRow = event.target.closest(
+                `.${EditorStyleHelper.tableGripRow}`
+              );
+              if (targetGripRow) {
+                event.preventDefault();
+                event.stopImmediatePropagation();
+
+                selectRow(
+                  Number(targetGripRow.getAttribute("data-index")),
+                  event.metaKey || event.shiftKey
+                )(view.state, view.dispatch);
+                return true;
+              }
+
+              return false;
+            },
+          },
+          decorations: (state) => {
+            if (!this.editor.view?.editable) {
+              return;
+            }
+
+            const { doc } = state;
+            const decorations: Decoration[] = [];
+            const rows = getRowsInTable(state);
+
+            if (rows && rows.length > 0 && isInTable(state)) {
+              const rect = selectedRect(state);
+              const firstColumnCells = new Map<number, number>();
+
+              // Map each visual row index to its first column cell position
+              for (let row = 0; row < rect.map.height; row++) {
+                const cellPos =
+                  rect.tableStart + rect.map.map[row * rect.map.width];
+                firstColumnCells.set(row, cellPos);
+              }
+
+              rows.forEach((pos, visualIndex) => {
+                const index = visualIndex;
+
+                // Check if this row's first column is part of a merged cell from above
+                const currentFirstCellPos = firstColumnCells.get(visualIndex);
+                let isFirstColumnMerged = false;
+
+                for (let prevRow = 0; prevRow < visualIndex; prevRow++) {
+                  if (firstColumnCells.get(prevRow) === currentFirstCellPos) {
+                    isFirstColumnMerged = true;
+                    break;
+                  }
+                }
+
+                // Skip decorations for rows where first column is merged from above
+                if (isFirstColumnMerged) {
+                  return;
+                }
+
+                if (index === 0) {
+                  const className = cn(EditorStyleHelper.tableGrip, {
+                    selected: isTableSelected(state),
+                  });
+
+                  decorations.push(
+                    Decoration.widget(
+                      pos + 1,
+                      () => {
+                        const grip = document.createElement("a");
+                        grip.role = "button";
+                        grip.className = className;
+                        return grip;
+                      },
+                      {
+                        key: className,
+                      }
+                    )
+                  );
+                }
+
+                const className = cn(EditorStyleHelper.tableGripRow, {
+                  selected:
+                    isRowSelected(index)(state) || isTableSelected(state),
+                  first: index === 0,
+                  last: visualIndex === rows.length - 1,
+                });
+
+                decorations.push(
+                  Decoration.widget(
+                    pos + 1,
+                    () => {
+                      const grip = document.createElement("a");
+                      grip.role = "button";
+                      grip.className = className;
+                      grip.dataset.index = index.toString();
+                      return grip;
+                    },
+                    {
+                      key: cn(className, index),
+                    }
+                  )
+                );
+
+                if (index === 0) {
+                  decorations.push(buildAddRowDecoration(pos, index));
+                }
+
+                decorations.push(buildAddRowDecoration(pos, index + 1));
+              });
+            }
+
+            return DecorationSet.create(doc, decorations);
+          },
+        },
+      }),
+    ];
+  }
+
   toMarkdown() {
     // see: renderTable
   }

shared/editor/queries/table.ts

@@ -66,7 +66,6 @@ export function getRowIndexInMap(
   if (!isInTable(state)) {
     return -1;
   }
-
   const rect = selectedRect(state);
   const cells = getCellsInColumn(0)(state);
 
@@ -87,6 +86,71 @@ export function getRowIndexInMap(
   return -1;
 }
 
+/**
+ * Get the actual row positions in the table map, accounting for merged cells.
+ *
+ * Iterates through each visual row and returns the position of the first unique cell in that row.
+ * This ensures correct row identification even when cells span multiple rows (rowspan > 1).
+ *
+ * @param state The editor state
+ * @returns Array of cell positions representing the first unique cell in each row
+ */
+export function getRowsInTable(state: EditorState): number[] {
+  if (!isInTable(state)) {
+    return [];
+  }
+
+  const rect = selectedRect(state);
+  const rows: number[] = [];
+  const seenCells = new Set<number>();
+
+  for (let row = 0; row < rect.map.height; row++) {
+    // Find the leftmost cell in this row
+    for (let col = 0; col < rect.map.width; col++) {
+      const cellPos =
+        rect.tableStart + rect.map.map[row * rect.map.width + col];
+      if (!seenCells.has(cellPos)) {
+        rows.push(cellPos);
+        seenCells.add(cellPos);
+        break; // Only add the first unique cell per row
+      }
+    }
+  }
+
+  return rows;
+}
+
+/**
+ * Get the actual column positions in the table map, accounting for merged cells.
+ *
+ * @param state The editor state
+ * @returns Array of cell positions representing the first unique cell in each column
+ */
+export function getColumnsInTable(state: EditorState): number[] {
+  if (!isInTable(state)) {
+    return [];
+  }
+
+  const rect = selectedRect(state);
+  const columns: number[] = [];
+  const seenCells = new Set<number>();
+
+  for (let col = 0; col < rect.map.width; col++) {
+    // Find the topmost cell in this column
+    for (let row = 0; row < rect.map.height; row++) {
+      const cellPos =
+        rect.tableStart + rect.map.map[row * rect.map.width + col];
+      if (!seenCells.has(cellPos)) {
+        columns.push(cellPos);
+        seenCells.add(cellPos);
+        break; // Only add the first unique cell per column
+      }
+    }
+  }
+
+  return columns;
+}
+
 export function getCellsInColumn(index: number) {
   return (state: EditorState): number[] => {
     if (!isInTable(state)) {
@@ -196,14 +260,10 @@ export function isHeaderEnabled(
  * @returns Boolean indicating if the row is selected
  */
 export function isRowSelected(index: number) {
-  return (state: EditorState): boolean => {
-    if (isRowSelection(state)) {
-      const rect = selectedRect(state);
-      return rect.top <= index && rect.bottom > index;
-    }
-
-    return false;
-  };
+  return (state: EditorState): boolean =>
+    state.selection instanceof RowSelection && state.selection.isRowSelection()
+      ? state.selection.$index === index
+      : false;
 }
 
 /**

shared/editor/selection/RowSelection.ts

@@ -5,13 +5,26 @@ import { CellSelection, inSameTable, TableMap } from "prosemirror-tables";
 import type { Mappable } from "prosemirror-transform";
 
 export class RowSelection extends CellSelection {
+  constructor(
+    public $anchorCell: ResolvedPos,
+    public $headCell: ResolvedPos,
+    public $index: number = 0
+  ) {
+    super($anchorCell, $headCell);
+  }
+
   getBookmark(): RowBookmark {
-    return new RowBookmark(this.$anchorCell.pos, this.$headCell.pos);
+    return new RowBookmark(
+      this.$anchorCell.pos,
+      this.$headCell.pos,
+      this.$index
+    );
   }
 
   public static rowSelection(
     $anchorCell: ResolvedPos,
-    $headCell: ResolvedPos = $anchorCell
+    $headCell: ResolvedPos = $anchorCell,
+    $index: number = 0
   ): CellSelection {
     const table = $anchorCell.node(-1);
     const map = TableMap.get(table);
@@ -42,18 +55,23 @@ export class RowSelection extends CellSelection {
         );
       }
     }
-    return new RowSelection($anchorCell, $headCell);
+    return new RowSelection($anchorCell, $headCell, $index);
   }
 }
 
 export class RowBookmark {
   constructor(
     public anchor: number,
-    public head: number
+    public head: number,
+    public index: number = 0
   ) {}
 
   map(mapping: Mappable): RowBookmark {
-    return new RowBookmark(mapping.map(this.anchor), mapping.map(this.head));
+    return new RowBookmark(
+      mapping.map(this.anchor),
+      mapping.map(this.head),
+      this.index
+    );
   }
 
   resolve(doc: Node): CellSelection | Selection {