From dce5349e7c29690c4657060a56059699ad79e3c6 Mon Sep 17 00:00:00 2001
From: wanxiankai <xiankai.wan@gmail.com>
Date: Sat, 18 Jul 2026 22:48:36 +0800
Subject: [PATCH] fix: support callback forwarded refs

---
 src/useEnsuredForwardedRef.ts         | 28 ++++++++++++++++++++-------
 tests/useEnsuredForwardedRef.test.tsx | 21 ++++++++++++++++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/src/useEnsuredForwardedRef.ts b/src/useEnsuredForwardedRef.ts
index ad26fd1a72..0005c01780 100644
--- a/src/useEnsuredForwardedRef.ts
+++ b/src/useEnsuredForwardedRef.ts
@@ -1,4 +1,5 @@
 import {
+  ForwardedRef,
   forwardRef,
   ForwardRefExoticComponent,
   MutableRefObject,
@@ -10,16 +11,29 @@ import {
   useRef,
 } from 'react';
 
+const assignRef = <T>(ref: ForwardedRef<T>, value: T | null) => {
+  if (!ref) {
+    return;
+  }
+
+  if (typeof ref === 'function') {
+    ref(value);
+  } else {
+    ref.current = value;
+  }
+};
+
 export default function useEnsuredForwardedRef<T>(
-  forwardedRef: MutableRefObject<T>
+  forwardedRef: ForwardedRef<T>
 ): MutableRefObject<T> {
-  const ensuredRef = useRef(forwardedRef && forwardedRef.current);
+  const ensuredRef = useRef<T>(
+    (forwardedRef && typeof forwardedRef !== 'function' ? forwardedRef.current : null) as T
+  );
 
   useEffect(() => {
-    if (!forwardedRef) {
-      return;
-    }
-    forwardedRef.current = ensuredRef.current;
+    assignRef(forwardedRef, ensuredRef.current);
+
+    return () => assignRef(forwardedRef, null);
   }, [forwardedRef]);
 
   return ensuredRef;
@@ -29,7 +43,7 @@ export function ensuredForwardRef<T, P = {}>(
   Component: RefForwardingComponent<T, P>
 ): ForwardRefExoticComponent<PropsWithoutRef<P> & RefAttributes<T>> {
   return forwardRef((props: PropsWithChildren<P>, ref) => {
-    const ensuredRef = useEnsuredForwardedRef(ref as MutableRefObject<T>);
+    const ensuredRef = useEnsuredForwardedRef(ref);
     return Component(props, ensuredRef);
   });
 }
diff --git a/tests/useEnsuredForwardedRef.test.tsx b/tests/useEnsuredForwardedRef.test.tsx
index d6ba4fe63c..fffae4cbf1 100644
--- a/tests/useEnsuredForwardedRef.test.tsx
+++ b/tests/useEnsuredForwardedRef.test.tsx
@@ -72,3 +72,24 @@ test('should return a valid ref when using the wrapper function style', () => {
   expect(initialRef.current).toBeTruthy();
   expect(initialRef.current?.id).toBe('test_id');
 });
+
+test('should call callback refs when using the wrapper function style', () => {
+  const callbackRef = jest.fn();
+
+  const WrappedComponent = ensuredForwardRef<HTMLDivElement>((_props, ref) => {
+    return <div id="test_id" ref={ref} />;
+  });
+
+  TestUtils.act(() => {
+    ReactDOM.render(<WrappedComponent ref={callbackRef} />, container);
+  });
+
+  expect(callbackRef).toHaveBeenLastCalledWith(expect.any(HTMLDivElement));
+  expect(callbackRef.mock.calls[0][0].id).toBe('test_id');
+
+  TestUtils.act(() => {
+    ReactDOM.unmountComponentAtNode(container);
+  });
+
+  expect(callbackRef).toHaveBeenLastCalledWith(null);
+});