Y
A
S
H
Mar 04, 2025

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.

1import { lazy, Suspense, useRef, useState } from "react";
2import { createPortal } from "react-dom";
3
4const CreatePokemon = lazy(
5  () => import("./create-pokemon")
6);
7
8export function App() {
9  const [open, setOpen] = useState(false);
10  const buttonRef = useRef<HTMLButtonElement | null>(null);
11  const spinner = createPortal(
12    <Spinner />,
13    buttonRef.current || document.body
14  );
15  return (
16    <div>
17      <Button ref={buttonRef} onClick={() => setOpen(true)}>
18        <Plus className="w-6 h-6 mr-2" />
19        Create New Pokémon
20      </Button>
21      <Suspense fallback={spinner}>
22        {open && (
23          <CreatePokemon
24            open={open}
25            onClose={() => setOpen(false)}
26          />
27        )}
28      </Suspense>
29    </div>
30  );
31}

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.

1import {
2  lazy,
3  Suspense,
4  useCallback,
5  useRef,
6  useState
7} from "react";
8import { createPortal } from "react-dom";
9
10const CreatePokemon = lazy(
11  () => import("./create-pokemon")
12);
13
14export function App() {
15  const [open, setOpen] = useState(false);
16  const buttonRef = useRef<HTMLButtonElement | null>(null);
17
18  // To keep track if the component is requested to be loaded
19  const isLoadedRef = useRef(false);
20  const isLoaded = useCallback(
21    () => isLoadedRef.current,
22    []
23  );
24
25  // To trigger the loading of the component
26  const trigger = useCallback(() => {
27    isLoadedRef.current = true;
28  }, []);
29
30  const spinner = createPortal(
31    <Spinner />,
32    buttonRef.current || document.body
33  );
34
35  return (
36    <div>
37      <Button
38        ref={buttonRef}
39        onClick={() => {
40          trigger();
41          setOpen(true);
42        }}
43      >
44        <Plus />
45        Create New Pokémon
46      </Button>
47      <Suspense fallback={spinner}>
48        {isLoaded() && (
49          <CreatePokemon
50            open={open}
51            onClose={() => setOpen(false)}
52          />
53        )}
54      </Suspense>
55    </div>
56  );
57}

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.

1<Suspense fallback={spinner}>
2	{false && (
3	  <CreatePokemon
4		open={open}
5		onClose={() => setOpen(false)}
6	  />
7	)}
8</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.

1<Suspense fallback={spinner}>
2	{true && (
3	  <CreatePokemon
4		open={open}
5		onClose={() => setOpen(false)}
6	  />
7	)}
8</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 it's 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 download 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 work 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 then 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.

1function useLazyLoad(props?: { forceUpdate?: boolean }) {
2  const { forceUpdate = false } = props || {};
3
4  const isLoadedRef = useRef(false);
5  const [_, force] = useState(0);
6
7  const isLoaded = useCallback(
8    () => isLoadedRef.current,
9    []
10  );
11  const trigger = useCallback(() => {
12    isLoadedRef.current = true;
13    if (forceUpdate) {
14      force((prev) => prev + 1);
15    }
16  }, [forceUpdate]);
17
18  return { isLoaded, trigger };
19}

Now, we can use this hook like this.

1export function App() {
2  const [open, setOpen] = useState(false);
3  const buttonRef = useRef<HTMLButtonElement | null>(null);
4  const { isLoaded, trigger } = useLazyLoad();
5
6  const spinner = createPortal(
7    <Spinner />,
8    buttonRef.current || document.body
9  );
10
11  return (
12    <div>
13      <Button
14        ref={buttonRef}
15        onClick={() => {
16          trigger();
17          setOpen(true);
18        }}
19      >
20        <Plus className="w-6 h-6 mr-2" />
21        Create New Pokémon
22      </Button>
23      <Suspense fallback={spinner}>
24        {isLoaded() && (
25          <CreatePokemon
26            open={open}
27            onClose={() => setOpen(false)}
28          />
29        )}
30      </Suspense>
31    </div>
32  );
33}

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.

1export function App() {
2  const [open, setOpen] = useState(false);
3  const buttonRef = useRef<HTMLButtonElement | null>(null);
4  const { isLoaded, trigger } = useLazyLoad({
5    forceUpdate: true
6  });
7
8  const spinner = createPortal(
9    <Spinner />,
10    buttonRef.current || document.body
11  );
12
13  return (
14    <div>
15      <Button
16        ref={buttonRef}
17        onMouseEnter={() => load().then(trigger)}
18        onFocus={() => load().then(trigger)}
19        onClick={() => {
20          trigger();
21          setOpen(true);
22        }}
23      >
24        <Plus className="w-6 h-6 mr-2" />
25        Create New Pokémon
26      </Button>
27      <Suspense fallback={spinner}>
28        {isLoaded() && (
29          <CreatePokemon
30            open={open}
31            onClose={() => setOpen(false)}
32          />
33        )}
34      </Suspense>
35    </div>
36  );
37}

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.

1import {
2  lazy,
3  Suspense,
4  useCallback,
5  useRef,
6  useState
7} from "react";
8import { createPortal } from "react-dom";
9
10const load = () => import("./radix-dialog");
11const RadixDialog = lazy(load);
12
13export function App() {
14  const [open, setOpen] = useState(false);
15  const buttonRef = useRef<HTMLButtonElement | null>(null);
16  const spinner = createPortal(
17    <Spinner />,
18    buttonRef.current || document.body
19  );
20  const { isLoaded, trigger } = useLazyLoad({});
21  return (
22    <>
23      <Button
24        ref={buttonRef}
25        onClick={() => {
26          trigger();
27          setOpen(true);
28        }}
29      >
30        <Plus className="w-6 h-6 mr-2 sr-only" />
31        Open Dialog
32      </Button>
33      <Suspense fallback={spinner}>
34        {isLoaded() && (
35          <RadixDialog
36            open={open}
37            onOpenChange={setOpen}
38          />
39        )}
40      </Suspense>
41    </>
42  );
43}

And the dialog will look like this.

1import {
2  Dialog,
3  DialogContent,
4  DialogDescription,
5  DialogHeader,
6  DialogTitle
7} from "@/components/ui/dialog";
8
9export default function RadixDialog(
10  props: React.ComponentProps<typeof Dialog>
11) {
12  return (
13    <Dialog {...props}>
14      <DialogContent>
15        <DialogHeader>
16          <DialogTitle>
17            Are you absolutely sure?
18          </DialogTitle>
19          <DialogDescription>
20            This action cannot be undone. This will
21            permanently delete your account and remove your
22            data from our servers.
23          </DialogDescription>
24        </DialogHeader>
25      </DialogContent>
26    </Dialog>
27  );
28}

The result will look like this.

Side note

It is okay to call dynamic import multiple times. The bundler likes webpack or vite, or 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.