feat(dialogs): implement focus restoration for dialogs and update LoginModal to support it

src/components/LoginModal.tsx

@@ -9,18 +9,21 @@ import {
   DialogContent,
   DialogHeader,
 } from '~/components/ds/ui'
+import type { RestoreFocusRef } from '~/components/ds/ui/Dialog'
 
 interface LoginModalProps {
   open: boolean
   description?: string
   onOpenChange: (open: boolean) => void
+  restoreFocusRef?: RestoreFocusRef
   onSocialSignIn: (provider: OAuthProvider) => void
 }
 
 export function LoginModal({
   open,
   description,
   onOpenChange,
+  restoreFocusRef,
   onSocialSignIn,
 }: LoginModalProps) {
   const providers = useQuery({
@@ -30,7 +33,7 @@ export function LoginModal({
   })
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
-      <DialogContent size="xs">
+      <DialogContent size="xs" restoreFocusRef={restoreFocusRef}>
         <DialogHeader title="Sign in to continue" description={description} />
         <DialogBody className="pb-6">
           <div className="space-y-3">

src/components/ds/ui/Dialog.tsx

@@ -23,6 +23,31 @@ import { twMerge } from 'tailwind-merge'
  *       <DialogFooter>…</DialogFooter>
  *     </DialogContent>
  *   </Dialog>
+ *
+ * FOCUS RESTORATION CONTRACT
+ *
+ * On close, focus must return to the control that opened the dialog.
+ *
+ *   1. One colocated opener  -> wrap it in <DialogTrigger asChild>.
+ *      Radix registers it and restores focus itself.
+ *
+ *   2. External or dynamic opener (a parent component, a table row, …)
+ *      -> pass `restoreFocusRef` to <DialogContent>. Without a
+ *      DialogTrigger, Radix's internal trigger ref is null and focus would
+ *      fall to <body>.
+ *
+ *   // Parent-owned opener
+ *   const openerRef = React.useRef<HTMLButtonElement>(null)
+ *   <Button ref={openerRef} onClick={() => setOpen(true)}>Open</Button>
+ *   <Dialog open={open} onOpenChange={setOpen}>
+ *     <DialogContent restoreFocusRef={openerRef}>…</DialogContent>
+ *   </Dialog>
+ *
+ *   // One opener per table row: record whichever row was clicked
+ *   const openerRef = React.useRef<HTMLElement | null>(null)
+ *   <Button onClick={(e) => { openerRef.current = e.currentTarget; setRow(r) }}>
+ *     Remove
+ *   </Button>
  */
 
 export const Dialog = DialogPrimitive.Root
@@ -39,19 +64,74 @@ const sizeStyles: Record<DialogSize, string> = {
   xl: 'max-w-xl',
 }
 
+/* ------------------------------------------------------- focus restoration -- */
+
+// Type-only export: erased at build time, so it is safe for Fast Refresh.
+export type RestoreFocusRef = React.RefObject<HTMLElement | null>
+
+/**
+ * Builds an `onCloseAutoFocus` handler for Radix modal content.
+ *
+ * Radix's own close handler calls `preventDefault()` and then focuses
+ * `context.triggerRef.current`, which is null when no <DialogTrigger> was
+ * rendered. Radix composes the caller's handler first and skips its own when
+ * the event is default-prevented, so we take over only when we hold a live
+ * element to restore to. Otherwise Radix's trigger logic still runs.
+ *
+ * NOT exported on purpose: exporting a plain function from a component file
+ * breaks React Fast Refresh.
+ */
+function createCloseAutoFocus(
+  restoreFocusRef: RestoreFocusRef | undefined,
+  userHandler?: (event: Event) => void,
+) {
+  return (event: Event) => {
+    userHandler?.(event)
+    if (event.defaultPrevented) return
+
+    const el = restoreFocusRef?.current
+    // `isConnected` guards against a stale node (e.g. a table row removed by
+    // the action the dialog just confirmed).
+    if (el && el.isConnected) {
+      event.preventDefault() // skip Radix's null-trigger focus
+      el.focus()
+    }
+  }
+}
+
+/* ------------------------------------------------------------ DialogContent -- */
+
 type DialogContentProps = {
   children: React.ReactNode
   size?: DialogSize
   className?: string
   /** Escape hatch for content that manages its own dismissal (e.g. a wizard mid-submit). */
   onInteractOutside?: DialogPrimitive.DialogContentProps['onInteractOutside']
+  /**
+   * Element to focus when the dialog closes. Required for controlled dialogs
+   * that have no <DialogTrigger> (external or dynamic openers). Not needed
+   * when a colocated <DialogTrigger asChild> opens the dialog.
+   */
+  restoreFocusRef?: RestoreFocusRef
+  /**
+   * Runs before focus restoration. Call `event.preventDefault()` to take full
+   * control of where focus goes; restoreFocusRef is then ignored.
+   */
+  onCloseAutoFocus?: DialogPrimitive.DialogContentProps['onCloseAutoFocus']
 }
 
 export const DialogContent = React.forwardRef<
   HTMLDivElement,
   DialogContentProps
 >(function DialogContent(
-  { children, size = 'sm', className, onInteractOutside },
+  {
+    children,
+    size = 'sm',
+    className,
+    onInteractOutside,
+    restoreFocusRef,
+    onCloseAutoFocus,
+  },
   ref,
 ) {
   return (
@@ -64,6 +144,10 @@ export const DialogContent = React.forwardRef<
         ref={ref}
         data-ds-dialog-panel=""
         onInteractOutside={onInteractOutside}
+        onCloseAutoFocus={createCloseAutoFocus(
+          restoreFocusRef,
+          onCloseAutoFocus,
+        )}
         className={twMerge(
           // Centring uses the independent `translate` property (that is what
           // Tailwind v4 compiles these to), which leaves `transform` free for
@@ -89,6 +173,8 @@ export const DialogContent = React.forwardRef<
   )
 })
 
+/* ------------------------------------------------------------ DialogHeader -- */
+
 type DialogHeaderProps = {
   title: React.ReactNode
   /**

src/components/npm-stats/BaselineSection.tsx

@@ -21,6 +21,7 @@ import {
   DialogBody,
   DialogContent,
   DialogHeader,
+  DialogTrigger,
 } from '~/components/ds/ui'
 import { PackageSearch } from './PackageSearch'
 import { getBaselineDisplayName, type PackageGroup } from './shared'
@@ -142,14 +143,16 @@ export function BaselineSection({
 
   const addButton = (
     <Tooltip content="Search for a package to add as baseline">
-      <button
-        onClick={() => setShowSearch(true)}
-        className="flex items-center gap-1 px-1.5 py-0.5 text-xs rounded
+      <DialogTrigger asChild>
+        <button
+          type="button"
+          className="flex items-center gap-1 px-1.5 py-0.5 text-xs rounded
           text-blue-700 dark:text-blue-300 hover:bg-blue-500/10 font-medium"
-      >
-        <PlusIcon className="w-3 h-3" />
-        Add
-      </button>
+        >
+          <PlusIcon className="w-3 h-3" />
+          Add
+        </button>
+      </DialogTrigger>
     </Tooltip>
   )
 
@@ -240,7 +243,7 @@ export function BaselineSection({
   )
 
   return (
-    <>
+    <Dialog open={showSearch} onOpenChange={setShowSearch}>
       <div className="flex flex-wrap items-center gap-x-2 gap-y-1 text-xs">
         {labelButton}
 
@@ -256,29 +259,27 @@ export function BaselineSection({
         </div>
       </div>
 
-      <Dialog open={showSearch} onOpenChange={setShowSearch}>
-        <DialogContent size="md">
-          <DialogHeader
-            title={
-              <span className="flex items-center gap-2">
-                <PushPinIcon className="w-4 h-4 text-icon-accent" />
-                Add baseline package
-              </span>
-            }
+      <DialogContent size="md">
+        <DialogHeader
+          title={
+            <span className="flex items-center gap-2">
+              <PushPinIcon className="w-4 h-4 text-icon-accent" />
+              Add baseline package
+            </span>
+          }
+        />
+        <DialogBody className="pb-6">
+          <PackageSearch
+            onSelect={(pkg) => {
+              onAddBaseline(pkg)
+              setShowSearch(false)
+            }}
+            placeholder="Search for baseline package..."
+            // eslint-disable-next-line jsx-a11y/no-autofocus
+            autoFocus={true}
           />
-          <DialogBody className="pb-6">
-            <PackageSearch
-              onSelect={(pkg) => {
-                onAddBaseline(pkg)
-                setShowSearch(false)
-              }}
-              placeholder="Search for baseline package..."
-              // eslint-disable-next-line jsx-a11y/no-autofocus
-              autoFocus={true}
-            />
-          </DialogBody>
-        </DialogContent>
-      </Dialog>
-    </>
+        </DialogBody>
+      </DialogContent>
+    </Dialog>
   )
 }

src/contexts/LoginModalContext.tsx

@@ -52,8 +52,16 @@ export function LoginModalProvider({ children }: LoginModalProviderProps) {
   > | null>(null)
   const pendingOnSuccessRef = React.useRef<(() => void) | undefined>(undefined)
 
+  const openerRef = React.useRef<HTMLElement | null>(null)
+
   const openLoginModal = React.useCallback(
     (options?: { description?: string; onSuccess?: () => void }) => {
+      const active = document.activeElement
+      openerRef.current =
+        active instanceof HTMLElement && active !== document.body
+          ? active
+          : null
+
       pendingOnSuccessRef.current = options?.onSuccess
       setDescription(options?.description)
       setHasLoadedModal(true)
@@ -129,6 +137,7 @@ export function LoginModalProvider({ children }: LoginModalProviderProps) {
             open={isOpen}
             description={description}
             onOpenChange={handleOpenChange}
+            restoreFocusRef={openerRef}
             onSocialSignIn={openSocialPopup}
           />
         </React.Suspense>

src/routes/admin/roles.$roleId.tsx

@@ -1,5 +1,5 @@
 import { Link, redirect, createFileRoute } from '@tanstack/react-router'
-import { useState, useMemo, useCallback } from 'react'
+import { useState, useMemo, useCallback, useRef } from 'react'
 import { useRemoveUsersFromRole } from '~/utils/mutations'
 import { useQuery } from '@tanstack/react-query'
 import { getRole, getUsersWithRole } from '~/utils/roles.functions'
@@ -47,6 +47,9 @@ function RoleDetailPage() {
     name: string
   } | null>(null)
 
+  const openerRef = useRef<HTMLElement | null>(null)
+  const tableFocusRef = useRef<HTMLDivElement>(null)
+
   const userQuery = useCurrentUserQuery()
   const user = userQuery.data
   const roleQuery = useQuery({
@@ -200,12 +203,16 @@ function RoleDetailPage() {
         header: 'Actions',
         cell: ({ row }) => {
           const user = row.original
+          const name = user.name || user.email
           return (
             <button
-              onClick={() => {
+              type="button"
+              aria-label={`Remove ${name} from role`}
+              onClick={(e) => {
+                openerRef.current = e.currentTarget
                 setConfirmRemove({
                   userId: user._id,
-                  name: user.name || user.email,
+                  name,
                 })
               }}
               className="text-red-600 hover:text-red-900 dark:text-red-400 dark:hover:text-red-300"
@@ -336,7 +343,7 @@ function RoleDetailPage() {
             if (!open) setConfirmRemove(null)
           }}
         >
-          <DialogContent size="sm">
+          <DialogContent size="sm" restoreFocusRef={openerRef}>
             <DialogHeader
               title="Confirm Removal"
               description={
@@ -360,6 +367,7 @@ function RoleDetailPage() {
                       roleId: roleId,
                       userIds: [confirmRemove.userId],
                     })
+                    openerRef.current = tableFocusRef.current
                     setConfirmRemove(null)
                   } catch (error) {
                     console.error(
@@ -381,7 +389,11 @@ function RoleDetailPage() {
           </DialogContent>
         </Dialog>
 
-        <div className="bg-white dark:bg-gray-800 rounded-lg shadow-lg overflow-hidden">
+        <div
+          ref={tableFocusRef}
+          tabIndex={-1}
+          className="bg-white dark:bg-gray-800 rounded-lg shadow-lg overflow-hidden"
+        >
           <div className="overflow-x-auto">
             <table className="w-full">
               <thead className="bg-gray-50 dark:bg-gray-700">