fix(core): fix issues with PT-Input PopoverModal (#8021)

packages/sanity/src/core/form/inputs/PortableText/object/modals/PopoverModal.tsx

@@ -2,7 +2,8 @@
 
 import {CloseIcon} from '@sanity/icons'
 import {Box, Flex, Text, useClickOutsideEvent, useGlobalKeyDown} from '@sanity/ui'
-import {type ReactNode, useCallback, useEffect, useRef, useState} from 'react'
+import {type ReactNode, useCallback, useRef, useState} from 'react'
+import FocusLock from 'react-focus-lock'
 import {type PortableTextEditorElement} from 'sanity/_singletons'
 
 import {Button, type PopoverProps} from '../../../../../../ui-components'
@@ -24,18 +25,8 @@ interface PopoverEditDialogProps {
 
 const POPOVER_FALLBACK_PLACEMENTS: PopoverProps['fallbackPlacements'] = ['top', 'bottom']
 
-export function PopoverEditDialog(props: PopoverEditDialogProps) {
+export function PopoverEditDialog(props: PopoverEditDialogProps): ReactNode {
   const {floatingBoundary, referenceBoundary, referenceElement, width = 1} = props
-  const [open, setOpen] = useState(false)
-
-  // This hook is here to set open after the initial render.
-  // If rendered immediately, the popover will for a split second be
-  // visible in the top left of the boundaryElement before correctly
-  // placed pointing at the reference element.
-  useEffect(() => {
-    setOpen(true)
-  }, [])
-
   return (
     <RootPopover
       content={<Content {...props} />}
@@ -44,34 +35,41 @@ export function PopoverEditDialog(props: PopoverEditDialogProps) {
       data-ui="PopoverEditDialog"
       fallbackPlacements={POPOVER_FALLBACK_PLACEMENTS}
       floatingBoundary={floatingBoundary}
-      open={open}
+      open
       overflow="auto"
       placement="bottom"
       portal="default"
       preventOverflow
       referenceBoundary={referenceBoundary}
       referenceElement={referenceElement}
       width={width}
+      autoFocus
     />
   )
 }
 
 function Content(props: PopoverEditDialogProps) {
   const {onClose, referenceBoundary, referenceElement, title, autoFocus} = props
+  const isClosedRef = useRef(false)
+
+  const handleClose = useCallback(() => {
+    isClosedRef.current = true
+    onClose()
+  }, [onClose])
 
   useGlobalKeyDown(
     useCallback(
       (event: KeyboardEvent) => {
         if (event.key === 'Escape') {
-          onClose()
+          handleClose()
         }
       },
-      [onClose],
+      [handleClose],
     ),
   )
 
   useClickOutsideEvent(
-    onClose,
+    handleClose,
     () => [referenceElement],
     () => referenceBoundary,
   )
@@ -80,42 +78,44 @@ function Content(props: PopoverEditDialogProps) {
   const [contentElement, setContentElement] = useState<HTMLDivElement | null>(null)
   const containerElement = useRef<HTMLDivElement | null>(null)
 
-  useEffect(() => {
-    // When rendered, focus on the first input element in the content
-    if (contentElement) {
-      contentElement.querySelector('input')?.focus()
-    }
-  }, [contentElement])
+  const handleFocusLockWhiteList = useCallback((element: HTMLElement) => {
+    // This is needed in order for focusLock not to trap focus in the
+    // popover when closing the popover and focus is to be returned to the editor
+    if (isClosedRef.current) return false
+    return Boolean(element.contentEditable) || Boolean(containerElement.current?.contains(element))
+  }, [])
 
   return (
     <VirtualizerScrollInstanceProvider
       scrollElement={contentElement}
       containerElement={containerElement}
     >
-      <Flex ref={containerElement} direction="column" height="fill">
-        <ContentHeaderBox flex="none" padding={1}>
-          <Flex align="center">
-            <Box flex={1} padding={2}>
-              <Text weight="medium">{title}</Text>
-            </Box>
+      <FocusLock autoFocus whiteList={handleFocusLockWhiteList}>
+        <Flex ref={containerElement} direction="column" height="fill">
+          <ContentHeaderBox flex="none" padding={1}>
+            <Flex align="center">
+              <Box flex={1} padding={2}>
+                <Text weight="medium">{title}</Text>
+              </Box>
 
-            <Button
-              autoFocus={Boolean(autoFocus)}
-              icon={CloseIcon}
-              mode="bleed"
-              onClick={onClose}
-              tooltipProps={{content: 'Close'}}
-            />
-          </Flex>
-        </ContentHeaderBox>
-        <ContentScrollerBox flex={1}>
-          <PresenceOverlay margins={[0, 0, 1, 0]}>
-            <Box padding={3} ref={setContentElement}>
-              {props.children}
-            </Box>
-          </PresenceOverlay>
-        </ContentScrollerBox>
-      </Flex>
+              <Button
+                autoFocus={Boolean(autoFocus)}
+                icon={CloseIcon}
+                mode="bleed"
+                onClick={handleClose}
+                tooltipProps={{content: 'Close'}}
+              />
+            </Flex>
+          </ContentHeaderBox>
+          <ContentScrollerBox flex={1}>
+            <PresenceOverlay margins={[0, 0, 1, 0]}>
+              <Box padding={3} ref={setContentElement}>
+                {props.children}
+              </Box>
+            </PresenceOverlay>
+          </ContentScrollerBox>
+        </Flex>
+      </FocusLock>
     </VirtualizerScrollInstanceProvider>
   )
 }