fix(mosaic): keep Flow transitioning for the root height transition only

.changeset/flow-height-transitioning.md

@@ -0,0 +1,2 @@
+---
+---

packages/mosaic/src/primitives/flow/README.md

@@ -80,7 +80,9 @@ The initially active step does not receive `data-starting-style`. An exiting ste
 
 `Flow.Root` carries `data-initial` through the first measured frame. Styled adapters can use it to disable viewport transitions so the initial step never animates.
 
-`Flow.Root` carries `data-transitioning` while any step is exiting. Styled adapters should only transition the viewport height under it, so a resize inside the active step (for example a field message animating open) is tracked as-is instead of being animated a second time.
+`Flow.Root` carries `data-transitioning` from a step change until its own `height` transition finishes, independent of how long the steps take to enter or exit. If the height does not transition, it clears at once, and other animations on the root are ignored. Styled adapters should only transition the viewport height under it, so a resize inside the active step (for example a field message animating open) is tracked as-is instead of being animated a second time.
+
+`Flow.Root` carries `data-height-change` set to `grow` or `shrink` after the measured step height changes, so the height transition can use a different curve in each direction.
 
 ## CSS variable
 

packages/mosaic/src/primitives/flow/flow-context.ts

@@ -8,8 +8,6 @@ export interface FlowContextValue {
   rootRef: RefObject<HTMLDivElement | null>;
   registerActiveStep: (element: HTMLElement) => void;
   unregisterActiveStep: (element: HTMLElement) => void;
-  registerExitingStep: (element: HTMLElement) => void;
-  unregisterExitingStep: (element: HTMLElement) => void;
 }
 
 export const FlowContext = createContext<FlowContextValue | null>(null);

packages/mosaic/src/primitives/flow/flow-root.tsx

@@ -7,6 +7,10 @@ import { type ComponentProps, mergeProps, useRender } from '../utils';
 import { autoUpdate, getDimensions } from '../utils/dom';
 import { FlowContext, type FlowContextValue, type FlowDirection } from './flow-context';
 
+function isHeightTransition(animation: Animation): boolean {
+  return 'transitionProperty' in animation && animation.transitionProperty === 'height';
+}
+
 export interface FlowRootProps extends ComponentProps<'div'> {
   value: string;
   direction?: FlowDirection;
@@ -18,35 +22,23 @@ export const FlowRoot = React.forwardRef<HTMLDivElement, FlowRootProps>(function
   const [activeStep, setActiveStep] = useState<HTMLElement | null>(null);
   const [measured, setMeasured] = useState(false);
   const [initial, setInitial] = useState(true);
-  const [exitingSteps, setExitingSteps] = useState<ReadonlySet<HTMLElement>>(() => new Set());
-  const [settling, setSettling] = useState(false);
+  const [transitioning, setTransitioning] = useState(false);
+  const lastActiveStepRef = useRef<HTMLElement | null>(null);
   const stepHeightRef = useRef<number | null>(null);
-  const runOnRootAnimationsFinished = useAnimationsFinished(rootRef, false);
-  const exiting = exitingSteps.size > 0;
+  const runOnHeightTransitionFinished = useAnimationsFinished(rootRef, false, isHeightTransition);
 
   const registerActiveStep = useCallback((element: HTMLElement) => {
+    if (lastActiveStepRef.current && lastActiveStepRef.current !== element) {
+      setTransitioning(true);
+    }
+    lastActiveStepRef.current = element;
     setActiveStep(element);
   }, []);
 
   const unregisterActiveStep = useCallback((element: HTMLElement) => {
     setActiveStep(current => (current === element ? null : current));
   }, []);
 
-  const registerExitingStep = useCallback((element: HTMLElement) => {
-    setExitingSteps(current => new Set(current).add(element));
-  }, []);
-
-  const unregisterExitingStep = useCallback((element: HTMLElement) => {
-    setExitingSteps(current => {
-      if (!current.has(element)) {
-        return current;
-      }
-      const next = new Set(current);
-      next.delete(element);
-      return next;
-    });
-  }, []);
-
   useLayoutEffect(() => {
     if (!activeStep) {
       return;
@@ -65,16 +57,12 @@ export const FlowRoot = React.forwardRef<HTMLDivElement, FlowRootProps>(function
   }, [activeStep]);
 
   useLayoutEffect(() => {
-    if (exiting) {
-      setSettling(true);
-      return;
-    }
-    if (!settling) {
+    if (!transitioning) {
       return;
     }
 
-    return runOnRootAnimationsFinished(() => setSettling(false));
-  }, [exiting, settling, runOnRootAnimationsFinished]);
+    return runOnHeightTransitionFinished(() => setTransitioning(false));
+  }, [transitioning, activeStep, runOnHeightTransitionFinished]);
 
   useLayoutEffect(() => {
     if (!measured || !initial) {
@@ -92,10 +80,8 @@ export const FlowRoot = React.forwardRef<HTMLDivElement, FlowRootProps>(function
       rootRef,
       registerActiveStep,
       unregisterActiveStep,
-      registerExitingStep,
-      unregisterExitingStep,
     }),
-    [value, direction, registerActiveStep, unregisterActiveStep, registerExitingStep, unregisterExitingStep],
+    [value, direction, registerActiveStep, unregisterActiveStep],
   );
 
   const element = useRender({
@@ -105,7 +91,7 @@ export const FlowRoot = React.forwardRef<HTMLDivElement, FlowRootProps>(function
     props: mergeProps<'div'>(
       {
         'data-initial': initial ? '' : undefined,
-        'data-transitioning': exiting || settling ? '' : undefined,
+        'data-transitioning': transitioning ? '' : undefined,
       },
       otherProps,
     ),

packages/mosaic/src/primitives/flow/flow-step.tsx

@@ -29,15 +29,7 @@ function firstInDocumentOrder(elements: Iterable<HTMLElement>): HTMLElement | nu
 
 export const FlowStep = React.forwardRef<HTMLDivElement, FlowStepProps>(function FlowStep(props, forwardedRef) {
   const { render, ids, children, ...otherProps } = props;
-  const {
-    value,
-    direction,
-    rootRef,
-    registerActiveStep,
-    unregisterActiveStep,
-    registerExitingStep,
-    unregisterExitingStep,
-  } = useFlowContext();
+  const { value, direction, rootRef, registerActiveStep, unregisterActiveStep } = useFlowContext();
   const open = ids.includes(value);
   const stepRef = useRef<HTMLDivElement | null>(null);
   const activeChildrenRef = useRef(children);
@@ -63,18 +55,6 @@ export const FlowStep = React.forwardRef<HTMLDivElement, FlowStepProps>(function
     return () => unregisterActiveStep(element);
   }, [open, registerActiveStep, unregisterActiveStep]);
 
-  const exiting = mounted && !open;
-
-  useLayoutEffect(() => {
-    const element = stepRef.current;
-    if (!exiting || !element) {
-      return;
-    }
-
-    registerExitingStep(element);
-    return () => unregisterExitingStep(element);
-  }, [exiting, registerExitingStep, unregisterExitingStep]);
-
   // A layout effect, so focus moves in the commit that makes the leaving step inert: field to field,
   // which iOS does without dropping the keyboard. A passive effect would leave a frame with none.
   useLayoutEffect(() => {

packages/mosaic/src/primitives/flow/flow.test.tsx

@@ -239,40 +239,101 @@ describe('Flow', () => {
 
     expect(root).not.toHaveAttribute('data-initial');
 
+    expect(root).not.toHaveAttribute('data-height-change');
+
     rerender(<TestFlow value='otp' />);
 
     expect(root.style.getPropertyValue('--cl-flow-step-height')).toBe('240px');
+    expect(root).toHaveAttribute('data-height-change', 'grow');
     expect(root).not.toHaveAttribute('data-initial');
+
+    rerender(<TestFlow value='password' />);
+
+    expect(root).toHaveAttribute('data-height-change', 'shrink');
     offsetHeight.mockRestore();
   });
-  it('marks the root as transitioning only while a step is exiting', async () => {
-    let finishAnimation!: () => void;
-    const animationFinished = new Promise<void>(resolve => {
-      finishAnimation = resolve;
+  describe('data-transitioning', () => {
+    function deferred() {
+      let resolve!: () => void;
+      const promise = new Promise<void>(r => {
+        resolve = r;
+      });
+      return { promise, resolve };
+    }
+
+    function mockRootAnimations(root: HTMLElement, animations: Array<Record<string, unknown>>) {
+      root.getAnimations = vi.fn(() => animations as unknown as Animation[]);
+    }
+
+    it('lasts from a step change until the root height transition finishes', async () => {
+      const height = deferred();
+      const { rerender } = render(<TestFlow value='password' />);
+      const root = screen.getByTestId('flow-root');
+      mockRootAnimations(root, [{ transitionProperty: 'height', finished: height.promise }]);
+
+      rerender(<TestFlow value='password-error' />);
+
+      expect(root).not.toHaveAttribute('data-transitioning');
+
+      rerender(<TestFlow value='otp' />);
+
+      expect(screen.queryByTestId('password-step')).not.toBeInTheDocument();
+      expect(root).toHaveAttribute('data-transitioning');
+
+      mockRootAnimations(root, []);
+      await act(async () => {
+        height.resolve();
+        await height.promise;
+      });
+
+      expect(root).not.toHaveAttribute('data-transitioning');
     });
-    const { rerender } = render(<TestFlow value='password' />);
-    const root = screen.getByTestId('flow-root');
-    const outgoingStep = screen.getByTestId('password-step');
-    outgoingStep.getAnimations = vi.fn(() => [{ finished: animationFinished }] as unknown as Animation[]);
 
-    expect(root).not.toHaveAttribute('data-transitioning');
+    it('clears at once when the root height does not transition', () => {
+      const { rerender } = render(<TestFlow value='password' />);
+      const root = screen.getByTestId('flow-root');
+      mockRootAnimations(root, []);
 
-    rerender(<TestFlow value='password-error' />);
+      rerender(<TestFlow value='otp' />);
 
-    expect(root).not.toHaveAttribute('data-transitioning');
+      expect(root).not.toHaveAttribute('data-transitioning');
+    });
 
-    rerender(<TestFlow value='otp' />);
+    it('is not held by other animations on the root', () => {
+      const { rerender } = render(<TestFlow value='password' />);
+      const root = screen.getByTestId('flow-root');
+      mockRootAnimations(root, [{ animationName: 'pulse', finished: new Promise<void>(() => {}) }]);
 
-    expect(root).toHaveAttribute('data-transitioning');
+      rerender(<TestFlow value='otp' />);
 
-    outgoingStep.getAnimations = vi.fn(() => []);
-    await act(async () => {
-      finishAnimation();
-      await animationFinished;
+      expect(root).not.toHaveAttribute('data-transitioning');
     });
 
-    expect(screen.queryByTestId('password-step')).not.toBeInTheDocument();
-    expect(root).not.toHaveAttribute('data-transitioning');
+    it('waits for the latest height transition when the step changes again', async () => {
+      const first = deferred();
+      const second = deferred();
+      const { rerender } = render(<TestFlow value='password' />);
+      const root = screen.getByTestId('flow-root');
+      mockRootAnimations(root, [{ transitionProperty: 'height', finished: first.promise }]);
+
+      rerender(<TestFlow value='otp' />);
+      mockRootAnimations(root, [{ transitionProperty: 'height', finished: second.promise }]);
+      rerender(<TestFlow value='password' />);
+      await act(async () => {
+        first.resolve();
+        await first.promise;
+      });
+
+      expect(root).toHaveAttribute('data-transitioning');
+
+      mockRootAnimations(root, []);
+      await act(async () => {
+        second.resolve();
+        await second.promise;
+      });
+
+      expect(root).not.toHaveAttribute('data-transitioning');
+    });
   });
 
   describe('useFlowAutoFocus', () => {

packages/mosaic/src/primitives/hooks/__tests__/use-animations-finished.test.ts

@@ -5,7 +5,7 @@ import { describe, expect, it, vi } from 'vitest';
 import { useAnimationsFinished } from '../use-animations-finished';
 
 function createMockElement(
-  animations: Array<{ finished: Promise<void> }> = [],
+  animations: Array<{ finished: Promise<void>; transitionProperty?: string; animationName?: string }> = [],
   attributes: Record<string, string> = {},
 ): HTMLElement {
   const el = document.createElement('div');
@@ -120,6 +120,66 @@ describe('useAnimationsFinished', () => {
     expect(callback).toHaveBeenCalledTimes(1);
   });
 
+  describe('with a filter', () => {
+    const neverFinishes = new Promise<void>(() => {});
+    const isHeightTransition = (animation: Animation) =>
+      'transitionProperty' in animation && animation.transitionProperty === 'height';
+
+    it('fires callback immediately when only unmatched animations are running', () => {
+      const el = createMockElement([{ animationName: 'pulse', finished: neverFinishes }]);
+      const ref = { current: el } as RefObject<HTMLElement | null>;
+
+      const { result } = renderHook(() => useAnimationsFinished(ref, false, isHeightTransition));
+
+      const callback = vi.fn();
+      act(() => result.current(callback));
+      expect(callback).toHaveBeenCalledTimes(1);
+    });
+
+    it('waits only for matched animations', async () => {
+      let resolveHeight!: () => void;
+      const heightFinished = new Promise<void>(r => {
+        resolveHeight = r;
+      });
+      const el = createMockElement([
+        { transitionProperty: 'height', finished: heightFinished },
+        { animationName: 'pulse', finished: neverFinishes },
+      ]);
+      const ref = { current: el } as RefObject<HTMLElement | null>;
+
+      const { result } = renderHook(() => useAnimationsFinished(ref, false, isHeightTransition));
+
+      const callback = vi.fn();
+      act(() => result.current(callback));
+      expect(callback).not.toHaveBeenCalled();
+
+      await act(() => resolveHeight());
+      expect(callback).toHaveBeenCalledTimes(1);
+    });
+
+    it('re-checks only matched animations when one is cancelled', async () => {
+      let rejectHeight!: () => void;
+      const cancelledHeight = new Promise<void>((_, reject) => {
+        rejectHeight = reject;
+      });
+      const el = createMockElement([{ transitionProperty: 'height', finished: cancelledHeight }]);
+      const ref = { current: el } as RefObject<HTMLElement | null>;
+
+      const { result } = renderHook(() => useAnimationsFinished(ref, false, isHeightTransition));
+
+      const callback = vi.fn();
+      act(() => result.current(callback));
+
+      el.getAnimations = vi.fn(() => [{ animationName: 'pulse', finished: neverFinishes }] as unknown as Animation[]);
+      await act(async () => {
+        rejectHeight();
+        await new Promise(r => setTimeout(r, 0));
+      });
+
+      expect(callback).toHaveBeenCalledTimes(1);
+    });
+  });
+
   it('waits for starting-style attribute removal when open=true', async () => {
     const el = createMockElement([], { 'data-starting-style': '' });
     const ref = { current: el } as RefObject<HTMLElement | null>;