Lazy-Load Headless UI Without Breaking Transitions
How to Lazy-Load Headless UI Components in React Without Breaking Transitions and Animations
I have a component called CreatePokemon, and I want to load it lazily. Below is how we lazy-load components in React using the lazy API.
import { lazy, Suspense, useRef, useState } from "react";
import { createPortal } from "react-dom";
const load = () => import("./create-pokemon");
const CreatePokemon = lazy(load);
export function App() {
const [open, setOpen] = useState(false);
const buttonRef = useRef<HTMLButtonElement | null>(null);
const spinner = createPortal(<Spinner />, buttonRef.current || document.body);
return (
<div>
<Button ref={buttonRef} onClick={() => setOpen(true)}>
<Plus className="w-6 h-6 mr-2" />
Create New Pokémon
</Button>
<Suspense fallback={spinner}>{open && <CreatePokemon open={open} onClose={() => setOpen(false)} />}</Suspense>
</div>
);
}The end result will look like this.
But here's the problem: the dialog component has enter and exit animations, but they aren't playing at all. The reason is that we're rendering it conditionally - which is how we typically load a component lazily when we only want to render it in response to specific actions. For example, I want to show this dialog component only when the user clicks the "Create New Pokémon" button.
We can't remove conditional rendering; otherwise, that chunk will be loaded as soon as the JavaScript file loads. To fix this issue, we need to keep the component mounted and let Headless UI handle the enter and exit animations as well as open and close states.
To do that, we can use a ref to track whether the component has been loaded or not. Here's how we can do it.
import { lazy, Suspense, useCallback, useRef, useState } from "react";
import { createPortal } from "react-dom";
const load = () => import("./create-pokemon");
const CreatePokemon = lazy(load);
export function App() {
const [open, setOpen] = useState(false);
const buttonRef = useRef<HTMLButtonElement | null>(null);
// To keep track of wether the component has been requested to load
const isLoadedRef = useRef(false);
const isLoaded = useCallback(() => isLoadedRef.current, []);
// To trigger the loading of the component
const trigger = useCallback(() => {
isLoadedRef.current = true;
}, []);
const spinner = createPortal(<Spinner />, buttonRef.current || document.body);
return (
<div>
<Button
ref={buttonRef}
onClick={() => {
trigger();
setOpen(true);
}}
>
<Plus />
Create New Pokémon
</Button>
<Suspense fallback={spinner}>
{isLoaded() && <CreatePokemon open={open} onClose={() => setOpen(false)} />}
</Suspense>
</div>
);
}The end result will look like this.
Bam! The transition is now working — but how?
Initially, the isLoadedRef is set to false, so the first render returns JSX that looks like this.
<Suspense fallback={spinner}>{false && <CreatePokemon open={open} onClose={() => setOpen(false)} />}</Suspense>Now, when we click the "Create New Pokémon" button, we set isLoadedRef to true, and then update the open state using setOpen. This triggers a re-render, and on the next render, React receives new JSX that looks like this.
<Suspense fallback={spinner}>{true && <CreatePokemon open={open} onClose={() => setOpen(false)} />}</Suspense>Here, React sees that a new component needs to be mounted, so it loads the CreatePokemon. Since this is the first time the component is requested, the bundler downloads the CreatePokemon bundle along with all its dependencies and loads them for us. During this process, the component is suspended, so the Suspense boundary shows the fallback UI. Once all the bundles are downloaded and parsed, we can finally see our sweet Pokémon dialog appear.
Now, if we close and open the dialog again, we can see that the transition works correctly. The reason is that every time the open state changes, the resulting JSX remains the same as before, since isLoadedRef is now always true. This means the component stays mounted (i.e., the condition is always true, and it always returns CreatePokemon rather than false), allowing Headless UI to handle the enter and exit animations properly.
Custom hook
We can wrap this logic in a custom hook to make it reusable.
function useLazyLoad(props?: { forceUpdate?: boolean }) {
const { forceUpdate = false } = props || {};
const isLoadedRef = useRef(false);
const [_, force] = useState(0);
const isLoaded = useCallback(() => isLoadedRef.current, []);
const trigger = useCallback(() => {
isLoadedRef.current = true;
if (forceUpdate) {
force((prev) => prev + 1);
}
}, [forceUpdate]);
return { isLoaded, trigger };
}Now, we can use this hook like this.
export function App() {
const [open, setOpen] = useState(false);
const buttonRef = useRef<HTMLButtonElement | null>(null);
const { isLoaded, trigger } = useLazyLoad();
const spinner = createPortal(<Spinner />, buttonRef.current || document.body);
return (
<div>
<Button
ref={buttonRef}
onClick={() => {
trigger();
setOpen(true);
}}
>
<Plus className="w-6 h-6 mr-2" />
Create New Pokémon
</Button>
<Suspense fallback={spinner}>
{isLoaded() && <CreatePokemon open={open} onClose={() => setOpen(false)} />}
</Suspense>
</div>
);
}The result is the same as before.
We can take this a step further. We can eagerly load the component when the user hovers over or focuses on the button. This improves the experience, because by the time the user clicks the button, the component is already loaded and ready to be displayed.
export function App() {
const [open, setOpen] = useState(false);
const buttonRef = useRef<HTMLButtonElement | null>(null);
const { isLoaded, trigger } = useLazyLoad({
forceUpdate: true
});
const spinner = createPortal(<Spinner />, buttonRef.current || document.body);
return (
<div>
<Button
ref={buttonRef}
onMouseEnter={() => load().then(trigger)}
onFocus={() => load().then(trigger)}
onClick={() => {
trigger();
setOpen(true);
}}
>
<Plus className="w-6 h-6 mr-2" />
Create New Pokémon
</Button>
<Suspense fallback={spinner}>
{isLoaded() && <CreatePokemon open={open} onClose={() => setOpen(false)} />}
</Suspense>
</div>
);
}The result is the same as before, but now the component is eagerly loaded on hover and focus.
You can see this by inspecting the Network tab in DevTools, or you might notice a spinner appear for a brief moment when you hover over the button.
You can use this technique with other headless libraries like React Aria or Radix UI as well.
For example, let's use our fancy useLazyLoad hook with Radix UI dialog.
import { lazy, Suspense, useCallback, useRef, useState } from "react";
import { createPortal } from "react-dom";
const load = () => import("./radix-dialog");
const RadixDialog = lazy(load);
export function App() {
const [open, setOpen] = useState(false);
const buttonRef = useRef<HTMLButtonElement | null>(null);
const spinner = createPortal(<Spinner />, buttonRef.current || document.body);
const { isLoaded, trigger } = useLazyLoad({});
return (
<>
<Button
ref={buttonRef}
onClick={() => {
trigger();
setOpen(true);
}}
>
<Plus className="w-6 h-6 mr-2 sr-only" />
Open Dialog
</Button>
<Suspense fallback={spinner}>{isLoaded() && <RadixDialog open={open} onOpenChange={setOpen} />}</Suspense>
</>
);
}And the dialog will look like this.
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
export default function RadixDialog(props: React.ComponentProps<typeof Dialog>) {
return (
<Dialog {...props}>
<DialogContent>
<DialogHeader>
<DialogTitle>Are you absolutely sure?</DialogTitle>
<DialogDescription>
This action cannot be undone. This will permanently delete your account and remove your data from our
servers.
</DialogDescription>
</DialogHeader>
</DialogContent>
</Dialog>
);
}The result will look like this.
Side note
It is okay to call dynamic import multiple times. Bundlers likes Webpack, Vite, or the Browser if we use native ESM import, they maintain a cache of all the promises that are created for these dynamic imports and resolved values of these promises. So, it doesn't matter if we call the dynamic import multiple times; it will only download and parse the module once. After that, it will return the cached resolved value of the promise.