fix react composed ref stability

Co-authored-by: Segun Adebayo <joseshegs@gmail.com>

Author: cursoragent

packages/react/src/components/color-picker/color-picker-content.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresenceContext } from '../presence/index.ts'
 import { useColorPickerContext } from './use-color-picker-context.ts'
@@ -14,12 +14,13 @@ export const ColorPickerContent = forwardRef<HTMLDivElement, ColorPickerContentP
   const colorPicker = useColorPickerContext()
   const presence = usePresenceContext()
   const mergedProps = mergeProps(colorPicker.getContentProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 ColorPickerContent.displayName = 'ColorPickerContent'

packages/react/src/components/combobox/combobox-content.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresenceContext } from '../presence/index.ts'
 import { useComboboxContext } from './use-combobox-context.ts'
@@ -14,12 +14,13 @@ export const ComboboxContent = forwardRef<HTMLDivElement, ComboboxContentProps>(
   const combobox = useComboboxContext()
   const presence = usePresenceContext()
   const mergedProps = mergeProps(combobox.getContentProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 ComboboxContent.displayName = 'ComboboxContent'

packages/react/src/components/date-picker/date-picker-content.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresenceContext } from '../presence/index.ts'
 import { useDatePickerContext } from './use-date-picker-context.ts'
@@ -14,12 +14,13 @@ export const DatePickerContent = forwardRef<HTMLDivElement, DatePickerContentPro
   const datePicker = useDatePickerContext()
   const presence = usePresenceContext()
   const mergedProps = mergeProps(datePicker.getContentProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 DatePickerContent.displayName = 'DatePickerContent'

packages/react/src/components/dialog/dialog-backdrop.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { useRenderStrategyPropsContext } from '../../utils/render-strategy.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresence } from '../presence/index.ts'
@@ -16,12 +16,13 @@ export const DialogBackdrop = forwardRef<HTMLDivElement, DialogBackdropProps>((p
   const renderStrategyProps = useRenderStrategyPropsContext()
   const presence = usePresence({ ...renderStrategyProps, present: dialog.open })
   const mergedProps = mergeProps(dialog.getBackdropProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 DialogBackdrop.displayName = 'DialogBackdrop'

packages/react/src/components/dialog/dialog-content.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresenceContext } from '../presence/index.ts'
 import { useDialogContext } from './use-dialog-context.ts'
@@ -14,12 +14,13 @@ export const DialogContent = forwardRef<HTMLDivElement, DialogContentProps>((pro
   const dialog = useDialogContext()
   const presence = usePresenceContext()
   const mergedProps = mergeProps(dialog.getContentProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 DialogContent.displayName = 'DialogContent'

packages/react/src/components/drawer/drawer-backdrop.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { useRenderStrategyPropsContext } from '../../utils/render-strategy.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresence } from '../presence/index.ts'
@@ -16,12 +16,13 @@ export const DrawerBackdrop = forwardRef<HTMLDivElement, DrawerBackdropProps>((p
   const renderStrategyProps = useRenderStrategyPropsContext()
   const presence = usePresence({ ...renderStrategyProps, present: drawer.open })
   const mergedProps = mergeProps(drawer.getBackdropProps(), presence.getPresenceProps(), props)
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 DrawerBackdrop.displayName = 'DrawerBackdrop'

packages/react/src/components/drawer/drawer-content.tsx

@@ -3,7 +3,7 @@
 import { mergeProps } from '@zag-js/react'
 import type { ContentProps } from '@zag-js/drawer'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { usePresenceContext } from '../presence/index.ts'
 import { useDrawerContext } from './use-drawer-context.ts'
@@ -23,12 +23,13 @@ export const DrawerContent = forwardRef<HTMLDivElement, DrawerContentProps>((pro
     presence.getPresenceProps(),
     localProps,
   )
+  const composedRefs = useComposedRefs(presence.ref, ref)
 
   if (presence.unmounted) {
     return null
   }
 
-  return <ark.div {...mergedProps} ref={composeRefs(presence.ref, ref)} />
+  return <ark.div {...mergedProps} ref={composedRefs} />
 })
 
 DrawerContent.displayName = 'DrawerContent'

packages/react/src/components/factory.test.tsx

@@ -1,5 +1,6 @@
 import { render, screen } from '@testing-library/react'
 import user from '@testing-library/user-event'
+import { useCallback, useReducer } from 'react'
 import { ark } from './factory.ts'
 
 const ComponentUnderTest = () => (
@@ -69,4 +70,34 @@ describe('Ark Factory', () => {
     )
     expect(screen.getByText('Ark UI')).not.toHaveAttribute('data-testid', 'parent')
   })
+
+  it('should not detach stable asChild callback ref on parent re-render', async () => {
+    const callbackRef = vi.fn()
+
+    const AsChildTest = () => {
+      const [, rerender] = useReducer((count) => count + 1, 0)
+      const setRef = useCallback((node: HTMLSpanElement | null) => {
+        callbackRef(node)
+      }, [])
+
+      return (
+        <>
+          <ark.span ref={setRef} asChild>
+            <span data-testid="child">Ark UI</span>
+          </ark.span>
+          <button type="button" onClick={() => rerender()}>
+            re-render
+          </button>
+        </>
+      )
+    }
+
+    render(<AsChildTest />)
+    expect(callbackRef).toHaveBeenCalledWith(screen.getByTestId('child'))
+
+    const callsAfterMount = callbackRef.mock.calls.length
+    await user.click(screen.getByRole('button', { name: /re-render/i }))
+
+    expect(callbackRef).toHaveBeenCalledTimes(callsAfterMount)
+  })
 })

packages/react/src/components/factory.ts

@@ -11,7 +11,7 @@ import {
   isValidElement,
   memo,
 } from 'react'
-import { composeRefs } from '../utils/compose-refs.ts'
+import { useComposedRefs } from '../utils/compose-refs.ts'
 
 export interface PolymorphicProps {
   /**
@@ -47,22 +47,22 @@ const withAsChild = (Component: React.ElementType) => {
   const Comp = memo(
     forwardRef<unknown, ArkPropsWithRef<typeof Component>>((props, ref) => {
       const { asChild, children, ...restProps } = props
+      const onlyChild =
+        asChild && isValidElement<Record<string, unknown>>(children) ? Children.only(children) : undefined
+      const childRef = onlyChild ? getRef(onlyChild) : undefined
+      const composedRef = useComposedRefs(ref, childRef)
 
       if (!asChild) {
         return createElement(Component, { ...restProps, ref }, children)
       }
 
-      if (!isValidElement<Record<string, unknown>>(children)) {
+      if (!onlyChild) {
         return null
       }
 
-      const onlyChild: React.ReactElement<Record<string, unknown>> = Children.only(children)
-
-      const childRef = getRef(onlyChild)
-
       return cloneElement(onlyChild, {
         ...mergeProps(restProps, onlyChild.props),
-        ref: ref ? composeRefs(ref, childRef) : childRef,
+        ref: ref ? composedRef : childRef,
       })
     }),
   )

packages/react/src/components/field/field-root.tsx

@@ -2,7 +2,7 @@
 
 import { mergeProps } from '@zag-js/react'
 import { forwardRef } from 'react'
-import { composeRefs } from '../../utils/compose-refs.ts'
+import { useComposedRefs } from '../../utils/compose-refs.ts'
 import { createSplitProps } from '../../utils/create-split-props.ts'
 import { type HTMLProps, type PolymorphicProps, ark } from '../factory.ts'
 import { type UseFieldProps, useField } from './use-field.ts'
@@ -26,10 +26,11 @@ export const FieldRoot = forwardRef<HTMLDivElement, FieldRootProps>((props, ref)
 
   const field = useField(useFieldProps)
   const mergedProps = mergeProps<HTMLProps<'div'>>(field.getRootProps(), localProps)
+  const composedRefs = useComposedRefs(ref, field.refs.rootRef)
 
   return (
     <FieldProvider value={field}>
-      <ark.div {...mergedProps} ref={composeRefs(ref, field.refs.rootRef)} />
+      <ark.div {...mergedProps} ref={composedRefs} />
     </FieldProvider>
   )
 })