Transition "leave" stutters in Firefox

[KingSora]

What package within Headless UI are you using?

@headlessui/react

What version of that package are you using?

v2.1.9

What browser are you using?

Firefox

Reproduction URL

https://stackblitz.com/edit/stackblitz-starters-wk4zwz?file=src%2Fstyle.css,src%2FApp.tsx,src%2Findex.tsx

Describe your issue

I've recently upgraded headlessui from v1 to v2 and noticed that certain transitions are stuttering in firefox. Its a bit "random" and certainly doesn't happen all the time, but its noticeable. I've added a video which shows what I mean. Please pay attention to the "leave" animation of the dialog, I'm opening the dialogin the minimal example 3 times. The first 2 times you can briefly observe the "stuttering" and the 3rd time the animation is smooth. The dialog is opened by the "trigger" button and closed with the esc key. (Please ignore all the audio in the video)

https://github.com/user-attachments/assets/d76b429f-deab-4083-b696-3efde10ddfbe

After some digging I believe firefox is briefly rendering the styles from the "leaveTo" classNames only to then switch back to the "leaveFrom" classNames.. The classNames are applied to the DOM in the correct order though.

To me it seems like the described behavior from the use-Transition hook workaround is still happening. (I could be wrong with that though)

After some tests I've found that the following changes are fixing the behavior:

function transition(
  node: HTMLElement,
  {
    prepare,
    run,
    done,
    inFlight,
  }: {
    prepare: () => void
    run: () => void
    done: () => void
    inFlight: MutableRefObject<boolean>
  }
) {
  let d = disposables()
  let originalTransition = node.style.transition;

  // same behavior as before
  if (inFlight?.current) {
    prepare()
  }
  else {
    // prepare at the end of the current frame
    d.requestAnimationFrame(() => {
      node.style.transition = "none";
      prepare();
    });
  }

  d.nextFrame(() => {
    // reset the transition in the next frame
    node.offsetHeight; // maybe not needed?
    node.style.transition = originalTransition;

    // Initiate the transition by applying the new classes.
    run()

    // Wait for the transition, once the transition is complete we can cleanup.
    // We wait for a frame such that the browser has time to flush the changes
    // to the DOM.
    d.requestAnimationFrame(() => {
      d.add(waitForTransition(node, done))
    })
  })

  return d.dispose
}

The main difference here is the node.style.transition value is modified accross mutlitple frames: Before:

1. node.style.transition = "none"
2. prepare();
3. node.offsetHeight to trigger reflow
4. node.style.transition = original

After:

1. wait for the end of the current frame
2. node.style.transition = "none"
3. prepare();
4. wait for next frame
5. node.offsetHeight to trigger reflow (maybe not needed anymore)
6. node.style.transition = original

I've done some tests and found everything working properly, although I'm certainly not an expert in the codebase and don't know all the implications the changes could have.

[KingSora]

Short addition:

I'm also observing the "enter" transitions to be affected, so that they are instantly in the "enterTo" state without any transition happening. - Its also "random" in a way where it seems to depend on Firefox's mood or some internal state.. sometimes the transition works and sometimes not.