Handle table selection style with an overlay element

Signed-off-by: Émile Ré <emile@getprobo.com>

packages/ui/src/RichEditor/BlockMenu/variants.ts

@@ -7,9 +7,9 @@ import { tv } from "tailwind-variants";
 export const blockMenuVariants = tv({
   slots: {
     trigger: [
-      "z-10 flex size-6 items-center justify-center",
+      "z-20 flex size-6 items-center justify-center",
       "rounded text-txt-tertiary hover:bg-subtle hover:text-txt-primary text-xl font-light cursor-pointer",
     ],
-    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-20"],
+    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-30"],
   },
 });

packages/ui/src/RichEditor/FocusedCellExtension.ts

@@ -1,40 +0,0 @@
-// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
-// Use of this source code is governed by the ISC license
-// that can be found in the LICENSE file.
-
-import { Extension } from "@tiptap/core";
-import { Plugin, PluginKey } from "@tiptap/pm/state";
-import { cellAround } from "@tiptap/pm/tables";
-import { Decoration, DecorationSet } from "@tiptap/pm/view";
-
-const focusedCellKey = new PluginKey("focusedCell");
-
-export const FocusedCellExtension = Extension.create({
-  name: "focusedCell",
-
-  addProseMirrorPlugins() {
-    return [
-      new Plugin({
-        key: focusedCellKey,
-
-        props: {
-          decorations(state) {
-            const { selection, doc } = state;
-            const $pos = doc.resolve(selection.from);
-            const cell = cellAround($pos);
-            if (!cell) return DecorationSet.empty;
-
-            const node = doc.nodeAt(cell.pos);
-            if (!node) return DecorationSet.empty;
-
-            return DecorationSet.create(doc, [
-              Decoration.node(cell.pos, cell.pos + node.nodeSize, {
-                class: "focusedCell",
-              }),
-            ]);
-          },
-        },
-      }),
-    ];
-  },
-});

packages/ui/src/RichEditor/OptionsMenu/variants.ts

@@ -7,9 +7,9 @@ import { tv } from "tailwind-variants";
 export const optionsMenuVariants = tv({
   slots: {
     trigger: [
-      "z-10 flex size-6 items-center justify-center",
+      "z-20 flex size-6 items-center justify-center",
       "rounded text-txt-tertiary hover:bg-subtle hover:text-txt-primary cursor-grab",
     ],
-    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-20"],
+    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-30"],
   },
 });

packages/ui/src/RichEditor/RichEditor.tsx

@@ -24,7 +24,6 @@ import { tv } from "tailwind-variants";
 
 import { BlockMenu } from "./BlockMenu/BlockMenu";
 import { BubbleMenu } from "./BubbleMenu";
-import { FocusedCellExtension } from "./FocusedCellExtension";
 import { LinkExtension } from "./LinkExtension";
 import { OptionsMenu } from "./OptionsMenu/OptionsMenu";
 import { PlaceholderExtension } from "./PlaceholderExtension";
@@ -32,6 +31,7 @@ import { SlashCommandExtension } from "./SlashCommandExtension";
 import { TableCellMenu } from "./TableCellMenu/TableCellMenu";
 import { TableColumnMenu } from "./TableColumnMenu/TableColumnMenu";
 import { TableRowMenu } from "./TableRowMenu/TableRowMenu";
+import { TableSelectionOverlay } from "./TableSelectionOverlay";
 
 const extensions = [
   Document,
@@ -63,7 +63,6 @@ const extensions = [
   TableKit.configure({
     table: { resizable: true },
   }),
-  FocusedCellExtension,
 ];
 
 const richEditorVariants = tv({
@@ -113,6 +112,7 @@ export function RichEditor(props: RichEditorProps) {
       <BubbleMenu editor={editor} />
       <BlockMenu editor={editor} />
       <OptionsMenu editor={editor} />
+      <TableSelectionOverlay editor={editor} />
       <TableCellMenu editor={editor} />
       <TableColumnMenu editor={editor} />
       <TableRowMenu editor={editor} />

packages/ui/src/RichEditor/TableCellMenu/TableCellMenuTrigger.tsx

@@ -57,9 +57,12 @@ export function TableCellMenuTrigger({
   useLayoutEffect(() => {
     const ed = editor;
     const fallback = activeCellEl;
+    const wrapper = fallback.closest(".tableWrapper");
 
     handleRefs.setReference({
       getBoundingClientRect() {
+        let rect: DOMRect;
+
         const { selection } = ed.state;
         if (selection instanceof CellSelection) {
           let top = Infinity;
@@ -69,17 +72,31 @@ export function TableCellMenuTrigger({
           selection.forEachCell((_node, pos) => {
             const el = cellDomElement(ed, pos);
             if (!el) return;
-            const rect = el.getBoundingClientRect();
+            const r = el.getBoundingClientRect();
-            top = Math.min(top, rect.top);
+            top = Math.min(top, r.top);
-            left = Math.min(left, rect.left);
+            left = Math.min(left, r.left);
-            bottom = Math.max(bottom, rect.bottom);
+            bottom = Math.max(bottom, r.bottom);
-            right = Math.max(right, rect.right);
+            right = Math.max(right, r.right);
           });
-          if (top !== Infinity) {
+          rect = top !== Infinity
-            return new DOMRect(left, top, right - left, bottom - top);
+            ? new DOMRect(left, top, right - left, bottom - top)
+            : fallback.getBoundingClientRect();
+        } else {
+          rect = fallback.getBoundingClientRect();
         }
+
+        if (wrapper) {
+          const clip = wrapper.getBoundingClientRect();
+          const clampedLeft = Math.max(rect.left, clip.left);
+          const clampedTop = Math.max(rect.top, clip.top);
+          const clampedRight = Math.min(rect.right, clip.right);
+          const clampedBottom = Math.min(rect.bottom, clip.bottom);
+          const w = Math.max(0, clampedRight - clampedLeft);
+          const h = Math.max(0, clampedBottom - clampedTop);
+          return new DOMRect(clampedLeft, clampedTop, w, h);
         }
-        return fallback.getBoundingClientRect();
+
+        return rect;
       },
     });
   }, [activeCellEl, editor, handleRefs]);

packages/ui/src/RichEditor/TableCellMenu/variants.ts

@@ -7,9 +7,9 @@ import { tv } from "tailwind-variants";
 export const tableCellMenuVariants = tv({
   slots: {
     trigger: [
-      "z-10 flex size-5 items-center justify-center",
+      "z-20 flex size-5 items-center justify-center",
       "rounded text-border-info cursor-pointer",
     ],
-    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-20"],
+    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-30"],
   },
 });

packages/ui/src/RichEditor/TableColumnMenu/variants.ts

@@ -7,10 +7,10 @@ import { tv } from "tailwind-variants";
 export const tableColumnMenuVariants = tv({
   slots: {
     trigger: [
-      "z-10 flex items-center justify-center",
+      "z-20 flex items-center justify-center",
       "rounded text-txt-tertiary bg-subtle hover:bg-border-solid cursor-grab",
       "py-0.5 h-3",
     ],
-    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-20"],
+    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-30"],
   },
 });

packages/ui/src/RichEditor/TableRowMenu/variants.ts

@@ -7,10 +7,10 @@ import { tv } from "tailwind-variants";
 export const tableRowMenuVariants = tv({
   slots: {
     trigger: [
-      "z-10 flex flex-col items-center justify-center",
+      "z-20 flex flex-col items-center justify-center",
       "rounded text-txt-tertiary bg-subtle hover:bg-border-solid cursor-grab",
       "px-0.5 w-3",
     ],
-    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-20"],
+    menu: ["rounded-lg border border-border-mid bg-level-0 p-1 shadow-md z-30"],
   },
 });

packages/ui/src/RichEditor/TableSelectionOverlay.tsx

@@ -0,0 +1,142 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+// Use of this source code is governed by the ISC license
+// that can be found in the LICENSE file.
+
+import {
+  autoUpdate,
+  type Middleware,
+  size,
+  useFloating,
+} from "@floating-ui/react";
+import { cellAround, CellSelection } from "@tiptap/pm/tables";
+import { type Editor, useEditorState } from "@tiptap/react";
+import { useLayoutEffect } from "react";
+import { createPortal } from "react-dom";
+
+import { cellDomElement } from "./_lib/cellDomElement";
+
+const cover: Middleware = {
+  name: "cover",
+  fn({ rects }) {
+    return { x: rects.reference.x, y: rects.reference.y };
+  },
+};
+
+type TableSelectionOverlayProps = {
+  editor: Editor;
+};
+
+function getActiveCellEl(editor: Editor): HTMLElement | null {
+  const { selection } = editor.state;
+
+  if (selection instanceof CellSelection) {
+    return cellDomElement(editor, selection.$headCell.pos);
+  }
+
+  const $pos = editor.state.doc.resolve(selection.from);
+  const cell = cellAround($pos);
+  if (!cell) return null;
+  return cellDomElement(editor, cell.pos);
+}
+
+export function TableSelectionOverlay({ editor }: TableSelectionOverlayProps) {
+  const activeCellEl = useEditorState({
+    editor,
+    selector: ({ editor: e }) => {
+      if (e.isDestroyed || !e.isEditable) return null;
+      return getActiveCellEl(e);
+    },
+  });
+
+  const isCellSelection = useEditorState({
+    editor,
+    selector: ({ editor: e }) =>
+      !e.isDestroyed && e.state.selection instanceof CellSelection,
+  });
+
+  const wrapperEl = activeCellEl?.closest(".tableWrapper") as HTMLElement | null;
+
+  const {
+    refs,
+    floatingStyles,
+    isPositioned,
+  } = useFloating({
+    strategy: "absolute",
+    placement: "bottom-start",
+    middleware: [
+      cover,
+      size({
+        apply({ rects, elements }) {
+          Object.assign(elements.floating.style, {
+            width: `${rects.reference.width}px`,
+            height: `${rects.reference.height}px`,
+          });
+        },
+      }),
+    ],
+    whileElementsMounted: (ref, floating, update) =>
+      autoUpdate(ref, floating, update, { animationFrame: true }),
+  });
+
+  useLayoutEffect(() => {
+    const ed = editor;
+    const fallback = activeCellEl;
+
+    if (!fallback) {
+      refs.setReference(null);
+      return;
+    }
+
+    refs.setReference({
+      getBoundingClientRect() {
+        const { selection } = ed.state;
+        if (selection instanceof CellSelection) {
+          let top = Infinity;
+          let left = Infinity;
+          let bottom = -Infinity;
+          let right = -Infinity;
+          selection.forEachCell((_node, pos) => {
+            const el = cellDomElement(ed, pos);
+            if (!el) return;
+            const r = el.getBoundingClientRect();
+            top = Math.min(top, r.top);
+            left = Math.min(left, r.left);
+            bottom = Math.max(bottom, r.bottom);
+            right = Math.max(right, r.right);
+          });
+          if (top !== Infinity) {
+            return new DOMRect(left, top, right - left, bottom - top);
+          }
+        }
+        return fallback.getBoundingClientRect();
+      },
+    });
+  }, [activeCellEl, editor, refs]);
+
+  if (!activeCellEl || !wrapperEl) return null;
+
+  return createPortal(
+    <div
+      ref={(node) => { refs.setFloating(node); }}
+      style={{
+        ...floatingStyles,
+        visibility: isPositioned ? "visible" : "hidden",
+        pointerEvents: "none",
+        zIndex: 10,
+      }}
+    >
+      <div
+        style={{
+          width: "100%",
+          height: "100%",
+          backgroundColor: isCellSelection
+            ? "color-mix(in srgb, var(--color-info) 50%, transparent)"
+            : undefined,
+          border: "2px solid var(--color-border-info)",
+          borderRadius: 2,
+        }}
+      />
+    </div>,
+    wrapperEl,
+  );
+}

packages/ui/src/rich-editor.css

@@ -83,25 +83,6 @@
         }
     }
 
-    th.selectedCell, td.selectedCell,
-    th.focusedCell, td.focusedCell {
-        background-color: color-mix(in srgb, var(--color-info) 50%, transparent) !important;
-        border: 2px solid var(--color-border-info) !important;
-    }
-
-    .selectedCell:has(+ .selectedCell) {
-        border-right: 1px solid var(--color-border-mid) !important;
-    }
-    .selectedCell + .selectedCell {
-        border-left: 1px solid var(--color-border-mid) !important;
-    }
-    tr:has(+ tr .selectedCell) > .selectedCell {
-        border-bottom: 1px solid var(--color-border-mid) !important;
-    }
-    tr:has(.selectedCell) + tr > .selectedCell {
-        border-top: 1px solid var(--color-border-mid) !important;
-    }
-
     .column-resize-handle {
         @apply absolute top-0 -right-px bottom-0 w-0.5 bg-border-info pointer-events-none z-20;
     }
@@ -111,6 +92,6 @@
     }
 
     .tableWrapper {
-        @apply overflow-x-auto my-4;
+        @apply relative overflow-x-auto my-4;
     }
 }