Keep row actions menu open when table actions callback changes identity

Author: david-crespo

app/table/columns/action-col.browser.spec.tsx

@@ -0,0 +1,55 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Copyright Oxide Computer Company
+ */
+import { getCoreRowModel, useReactTable, type ColumnDef } from '@tanstack/react-table'
+import { useCallback } from 'react'
+import { expect, test } from 'vitest'
+import { render } from 'vitest-browser-react'
+
+import { Table } from '../Table'
+import { useColsWithActions, type MenuAction } from './action-col'
+
+type Item = { id: string; name: string }
+
+const data: Item[] = [{ id: '1', name: 'alpha' }]
+const staticCols: ColumnDef<Item>[] = [{ accessorKey: 'name', header: 'Name' }]
+
+function ActionsTable({ version }: { version: number }) {
+  // depends on `version` so its identity changes when the prop changes, like a
+  // real page whose actions depend on a background query result
+  const makeActions = useCallback(
+    (item: Item): MenuAction[] => [
+      { label: `Rename ${item.name} v${version}`, onActivate: () => {} },
+    ],
+    [version]
+  )
+  const columns = useColsWithActions(staticCols, makeActions)
+  const table = useReactTable({
+    columns,
+    data,
+    getCoreRowModel: getCoreRowModel(),
+    getRowId: (row) => row.id,
+  })
+  return <Table aria-label="Items" table={table} />
+}
+
+test('open row actions menu survives makeActions identity change', async () => {
+  const screen = await render(<ActionsTable version={1} />)
+
+  await screen.getByRole('button', { name: 'Row actions' }).click()
+  await expect
+    .element(screen.getByRole('menuitem', { name: 'Rename alpha v1' }))
+    .toBeVisible()
+
+  // simulates a query resolving and changing a dep of makeActions while the
+  // menu is open. Before the fix, the cell remounted and the menu closed.
+  await screen.rerender(<ActionsTable version={2} />)
+
+  await expect
+    .element(screen.getByRole('menuitem', { name: 'Rename alpha v2' }))
+    .toBeVisible()
+})

app/table/columns/action-col.tsx

@@ -5,7 +5,7 @@
  *
  * Copyright Oxide Computer Company
  */
-import type { ColumnDef } from '@tanstack/react-table'
+import type { CellContext, ColumnDef } from '@tanstack/react-table'
 import cn from 'classnames'
 import { useMemo } from 'react'
 
@@ -35,13 +35,15 @@ type MenuActionLink = MenuActionBase & {
  */
 export type MenuAction = MenuActionItem | MenuActionLink
 
-type MakeActions<Item> = (item: Item) => Array<MenuAction>
+export type MakeActions<Item> = (item: Item) => Array<MenuAction>
 
-/** Convenience helper to combine regular cols with actions col and memoize */
+/**
+ * Convenience helper to combine regular cols with actions col and memoize.
+ * Memoizing `columns` and `makeActions` at the call site avoids re-renders but
+ * is no longer required for correctness (see `ActionsCell`).
+ */
 export function useColsWithActions<TData extends Record<string, unknown>>(
-  /** Should be static or memoized */
   columns: ColumnDef<TData, any>[], // eslint-disable-line @typescript-eslint/no-explicit-any
-  /** Must be memoized to avoid re-renders */
   makeActions: MakeActions<TData>,
   copyIdLabel?: string
 ) {
@@ -61,17 +63,34 @@ export const getActionsCol = <TData extends Record<string, unknown>>(
     meta: {
       thClassName: 'action-col',
       tdClassName: 'action-col',
+      makeActions,
+      copyIdLabel,
     },
-
-    cell: ({ row }) => {
-      // TODO: control flow here has always confused me, would like to straighten it out
-      const actions = makeActions(row.original)
-      const id = typeof row.original.id === 'string' ? row.original.id : null
-      return <RowActions id={id} actions={actions} copyIdLabel={copyIdLabel} />
-    },
+    cell: ActionsCell,
   }
 }
 
+/**
+ * `flexRender` renders a function `cell` as a component, so its identity is
+ * the element type. Passing `makeActions` through `meta` instead of a closure
+ * keeps the type stable when `makeActions` changes (e.g., a query it depends
+ * on refetches), which would otherwise remount the cell and close an open menu.
+ *
+ * This is why polling used to close row action menus and why call sites had to
+ * memoize `makeActions` and `columns` so carefully (#2112, #2460, #2816). Those
+ * memoizations are now only a render-count optimization.
+ */
+function ActionsCell<TData extends Record<string, unknown>>({
+  row,
+  column,
+}: CellContext<TData, unknown>) {
+  // TODO: control flow here has always confused me, would like to straighten it out
+  const { makeActions, copyIdLabel } = column.columnDef.meta ?? {}
+  const actions = makeActions?.(row.original)
+  const id = typeof row.original.id === 'string' ? row.original.id : null
+  return <RowActions id={id} actions={actions} copyIdLabel={copyIdLabel} />
+}
+
 type RowActionsProps = {
   /** If `id` is provided, a `Copy ID` menu item will be automatically included. */
   id?: string | null

types/react-table.d.ts

@@ -5,11 +5,16 @@
  *
  * Copyright Oxide Computer Company
  */
-import '@tanstack/react-table'
+import type { RowData } from '@tanstack/react-table'
+
+import type { MakeActions } from '../app/table/columns/action-col'
 
 declare module '@tanstack/react-table' {
-  interface ColumnMeta {
+  interface ColumnMeta<TData extends RowData, TValue> {
     thClassName?: string
     tdClassName?: string
+    /** Set by `getActionsCol`, read by `ActionsCell` */
+    makeActions?: MakeActions<TData>
+    copyIdLabel?: string
   }
 }

vitest.browser.config.ts

@@ -32,6 +32,7 @@ export default defineConfig({
   optimizeDeps: {
     entries: ['app/**/*.browser.spec.{ts,tsx}', 'app/util/ip.spec.ts'],
     include: [
+      '@base-ui/react/menu',
       'date-fns',
       'ip-num/IPNumber.js',
       'react-router',