diff --git a/packages/ui-patterns/src/CommandMenu/api/hooks/pagesHooks.ts b/packages/ui-patterns/src/CommandMenu/api/hooks/pagesHooks.ts index b1b5728e99545..5ee14039bd4b5 100644 --- a/packages/ui-patterns/src/CommandMenu/api/hooks/pagesHooks.ts +++ b/packages/ui-patterns/src/CommandMenu/api/hooks/pagesHooks.ts @@ -1,7 +1,7 @@ 'use client' -import { isEqual } from 'lodash' -import { useEffect, useMemo, useRef } from 'react' +import { useLayoutEffect, useMemo } from 'react' +import { useLatest } from 'react-use' import { useSnapshot } from 'valtio' import { useCommandContext } from '../../internal/Context' @@ -60,37 +60,18 @@ const useRegisterPage = ( const { pagesState } = useCommandContext() const { registerNewPage } = useSnapshot(pagesState) - const prevDeps = useRef(deps) - const prevEnabled = useRef(enabled) - - const unsubscribe = useRef<() => void>() - - /** - * useEffect handles the registration on first render, since React runs the - * first render twice in development. (Otherwise the first render would leave - * a dangling subscription.) - * - * It also handles final cleanup, since useMemo can't do this. - * - * useMemo handles the registration on subsequent renders, to ensure it - * happens synchronously. - */ - useMemo(() => { - if (!isEqual(prevDeps.current, deps) || prevEnabled.current !== enabled) { - unsubscribe.current?.() - - unsubscribe.current = enabled ? registerNewPage(name, definition) : undefined - - prevDeps.current = deps - prevEnabled.current = enabled - } - }, [registerNewPage, name, definition, deps, enabled]) - - useEffect(() => { - unsubscribe.current = enabled ? registerNewPage(name, definition) : undefined - - return () => unsubscribe.current?.() - }, []) + const definitionRef = useLatest(definition) + + // useLayoutEffect runs synchronously after DOM mutations but before paint. + // This ensures the page is registered before any subsequent code (e.g. in + // child components) tries to access it, avoiding a timing gap that would + // exist with useEffect. + useLayoutEffect(() => { + if (!enabled) return + + const unsubscribe = registerNewPage(name, definitionRef.current) + return unsubscribe + }, [registerNewPage, name, enabled, ...deps]) } -export { useCurrentPage, useRegisterPage, useSetPage, usePopPage, usePageComponent } +export { useCurrentPage, usePageComponent, usePopPage, useRegisterPage, useSetPage }