feat(ui): move focus to the entering Flow step

.changeset/mosaic-flow-autofocus.md

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

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

@@ -48,6 +48,25 @@ Multiple ids can select the same step. Moving between those ids updates the exis
 
 `Flow.Step` also accepts standard `<div>` attributes and the package's `render` prop.
 
+## Focus
+
+`useFlowAutoFocus()` returns a ref. Attach it to the element a step should focus when it enters:
+
+```tsx
+import { useFlowAutoFocus } from '@clerk/headless/flow';
+
+function PasswordView() {
+  return (
+    <input
+      ref={useFlowAutoFocus()}
+      type='password'
+    />
+  );
+}
+```
+
+Focus moves only for a step that transitions in; the initially active step is left to whatever container opened it. Focus is applied with `preventScroll` as soon as the step enters, and only when focus is currently on the body or inside `Flow.Root`, so it never steals from elsewhere on the page. Style `Flow.Root` with `overflow: clip` rather than `hidden` so focusing an element that is still sliding in cannot scroll the viewport. When several mounted elements are marked, the first in DOM order is focused. Outside a `Flow.Step` the hook returns a no-op ref.
+
 ## Transition attributes
 
 | Attribute             | Description                                              |

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

@@ -1,10 +1,11 @@
-import { createContext, useContext } from 'react';
+import { createContext, type RefObject, useContext } from 'react';
 
 export type FlowDirection = -1 | 1;
 
 export interface FlowContextValue {
   value: string;
   direction: FlowDirection;
+  rootRef: RefObject<HTMLDivElement | null>;
   registerActiveStep: (element: HTMLElement) => void;
   unregisterActiveStep: (element: HTMLElement) => void;
 }

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

@@ -55,7 +55,7 @@ export const FlowRoot = React.forwardRef<HTMLDivElement, FlowRootProps>(function
   }, [activeStepHeight, initial]);
 
   const contextValue = useMemo<FlowContextValue>(
-    () => ({ value, direction, registerActiveStep, unregisterActiveStep }),
+    () => ({ value, direction, rootRef, registerActiveStep, unregisterActiveStep }),
     [value, direction, registerActiveStep, unregisterActiveStep],
   );
 

packages/headless/src/primitives/flow/flow-step-context.ts

@@ -0,0 +1,35 @@
+'use client';
+
+import { createContext, type RefCallback, useCallback, useContext, useRef } from 'react';
+
+export interface FlowStepContextValue {
+  registerFocusTarget: (element: HTMLElement) => void;
+  unregisterFocusTarget: (element: HTMLElement) => void;
+}
+
+export const FlowStepContext = createContext<FlowStepContextValue | null>(null);
+
+/**
+ * Marks an element as the one to focus when the enclosing `Flow.Step` enters.
+ * When several mounted elements are marked, the first in DOM order is focused. Outside a
+ * step the ref is a no-op, so a view can render standalone without a wrapper.
+ */
+export function useFlowAutoFocus<T extends HTMLElement = HTMLElement>(): RefCallback<T> {
+  const context = useContext(FlowStepContext);
+  const elementRef = useRef<T | null>(null);
+
+  return useCallback(
+    (element: T | null) => {
+      if (element) {
+        elementRef.current = element;
+        context?.registerFocusTarget(element);
+        return;
+      }
+      if (elementRef.current) {
+        context?.unregisterFocusTarget(elementRef.current);
+        elementRef.current = null;
+      }
+    },
+    [context],
+  );
+}

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

@@ -1,23 +1,41 @@
 'use client';
 
 import { inertProps } from '@clerk/shared/inert';
-import React, { useLayoutEffect, useRef } from 'react';
+import React, { useCallback, useEffect, useLayoutEffect, useMemo, useRef } from 'react';
 
 import { useTransition } from '../../hooks/use-transition';
 import { type ComponentProps, mergeProps, useRender } from '../../utils';
 import { useFlowContext } from './flow-context';
+import { FlowStepContext, type FlowStepContextValue } from './flow-step-context';
 
 export interface FlowStepProps extends ComponentProps<'div'> {
   ids: readonly string[];
 }
 
+function focusIsWithin(root: HTMLElement): boolean {
+  const active = root.ownerDocument.activeElement;
+  return active === null || active === root.ownerDocument.body || root.contains(active);
+}
+
+function firstInDocumentOrder(elements: Iterable<HTMLElement>): HTMLElement | null {
+  let first: HTMLElement | null = null;
+  for (const element of elements) {
+    if (!first || first.compareDocumentPosition(element) & Node.DOCUMENT_POSITION_PRECEDING) {
+      first = element;
+    }
+  }
+  return first;
+}
+
 export const FlowStep = React.forwardRef<HTMLDivElement, FlowStepProps>(function FlowStep(props, forwardedRef) {
   const { render, ids, children, ...otherProps } = props;
-  const { value, direction, registerActiveStep, unregisterActiveStep } = useFlowContext();
+  const { value, direction, rootRef, registerActiveStep, unregisterActiveStep } = useFlowContext();
   const open = ids.includes(value);
   const stepRef = useRef<HTMLDivElement | null>(null);
   const activeChildrenRef = useRef(children);
   const hasBeenClosed = useRef(false);
+  const focusTargetsRef = useRef(new Set<HTMLElement>());
+  const wasOpenRef = useRef(open);
 
   if (open) {
     activeChildrenRef.current = children;
@@ -37,6 +55,31 @@ export const FlowStep = React.forwardRef<HTMLDivElement, FlowStepProps>(function
     return () => unregisterActiveStep(element);
   }, [open, registerActiveStep, unregisterActiveStep]);
 
+  useEffect(() => {
+    const entering = open && !wasOpenRef.current;
+    wasOpenRef.current = open;
+    if (!entering) {
+      return;
+    }
+
+    const target = firstInDocumentOrder(focusTargetsRef.current);
+    const root = rootRef.current;
+    if (target && root && root.contains(target) && focusIsWithin(root)) {
+      target.focus({ preventScroll: true });
+    }
+  }, [open, rootRef]);
+
+  const registerFocusTarget = useCallback((element: HTMLElement) => {
+    focusTargetsRef.current.add(element);
+  }, []);
+  const unregisterFocusTarget = useCallback((element: HTMLElement) => {
+    focusTargetsRef.current.delete(element);
+  }, []);
+  const stepContext = useMemo<FlowStepContextValue>(
+    () => ({ registerFocusTarget, unregisterFocusTarget }),
+    [registerFocusTarget, unregisterFocusTarget],
+  );
+
   const effectiveTransitionProps = !hasBeenClosed.current
     ? { ...transitionProps, 'data-starting-style': undefined, style: undefined }
     : transitionProps;
@@ -52,11 +95,13 @@ export const FlowStep = React.forwardRef<HTMLDivElement, FlowStepProps>(function
     children: open ? children : activeChildrenRef.current,
   };
 
-  return useRender({
+  const element = useRender({
     defaultTagName: 'div',
     enabled: mounted,
     render,
     ref: [stepRef, forwardedRef],
     props: mergeProps<'div'>(defaultProps, otherProps),
   });
+
+  return <FlowStepContext.Provider value={stepContext}>{element}</FlowStepContext.Provider>;
 });

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

@@ -1,8 +1,18 @@
 import { act, cleanup, render, screen } from '@testing-library/react';
-import { createRef } from 'react';
+import React, { createRef } from 'react';
+import { createPortal } from 'react-dom';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { Flow } from './index';
+import { Flow, useFlowAutoFocus } from './index';
+
+function AutoFocusInput(props: React.ComponentProps<'input'>) {
+  return (
+    <input
+      ref={useFlowAutoFocus()}
+      {...props}
+    />
+  );
+}
 
 interface TestFlowProps {
   value: string;
@@ -22,12 +32,14 @@ function TestFlow({ value, direction = 1, passwordContent = 'Password' }: TestFl
         data-testid='password-step'
       >
         {passwordContent}
+        <AutoFocusInput data-testid='password-input' />
       </Flow.Step>
       <Flow.Step
         ids={['otp', 'otp-pending', 'otp-error']}
         data-testid='otp-step'
       >
         OTP
+        <AutoFocusInput data-testid='otp-input' />
       </Flow.Step>
     </Flow.Root>
   );
@@ -233,4 +245,145 @@ describe('Flow', () => {
     expect(root).not.toHaveAttribute('data-initial');
     offsetHeight.mockRestore();
   });
+  describe('useFlowAutoFocus', () => {
+    it('does not focus the initially active step', () => {
+      render(<TestFlow value='password' />);
+
+      expect(screen.getByTestId('password-input')).not.toHaveFocus();
+    });
+
+    it('focuses the registered element without scrolling when the step enters', () => {
+      const focus = vi.spyOn(HTMLElement.prototype, 'focus');
+      const { rerender } = render(<TestFlow value='password' />);
+
+      rerender(<TestFlow value='otp' />);
+
+      expect(screen.getByTestId('otp-input')).toHaveFocus();
+      expect(focus).toHaveBeenCalledWith({ preventScroll: true });
+      focus.mockRestore();
+    });
+
+    it('leaves focus alone when it is outside the flow', () => {
+      const { rerender } = render(
+        <>
+          <button
+            type='button'
+            data-testid='outside'
+          >
+            Outside
+          </button>
+          <TestFlow value='password' />
+        </>,
+      );
+      screen.getByTestId('outside').focus();
+
+      rerender(
+        <>
+          <button
+            type='button'
+            data-testid='outside'
+          >
+            Outside
+          </button>
+          <TestFlow value='otp' />
+        </>,
+      );
+
+      expect(screen.getByTestId('outside')).toHaveFocus();
+      expect(screen.getByTestId('otp-input')).not.toHaveFocus();
+    });
+
+    it('focuses the returning step when navigation reverses mid-transition', () => {
+      const { rerender } = render(<TestFlow value='password' />);
+
+      rerender(<TestFlow value='otp' />);
+      const otpInput = screen.getByTestId('otp-input');
+      expect(otpInput).toHaveFocus();
+
+      rerender(
+        <TestFlow
+          value='password'
+          direction={-1}
+        />,
+      );
+
+      expect(otpInput).not.toHaveFocus();
+      expect(screen.getByTestId('password-input')).toHaveFocus();
+    });
+
+    it('focuses the first marked element in DOM order when several are marked', () => {
+      const otpStep = (
+        <Flow.Step ids={['otp']}>
+          <AutoFocusInput data-testid='first' />
+          <AutoFocusInput data-testid='second' />
+        </Flow.Step>
+      );
+      const { rerender } = render(
+        <Flow.Root value='password'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep}
+        </Flow.Root>,
+      );
+
+      rerender(
+        <Flow.Root value='otp'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep}
+        </Flow.Root>,
+      );
+
+      expect(screen.getByTestId('first')).toHaveFocus();
+    });
+
+    it('skips a marked element that is not rendered', () => {
+      const otpStep = (showFirst: boolean) => (
+        <Flow.Step ids={['otp']}>
+          {showFirst ? <AutoFocusInput data-testid='first' /> : null}
+          <AutoFocusInput data-testid='second' />
+        </Flow.Step>
+      );
+      const { rerender } = render(
+        <Flow.Root value='password'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep(false)}
+        </Flow.Root>,
+      );
+
+      rerender(
+        <Flow.Root value='otp'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep(false)}
+        </Flow.Root>,
+      );
+
+      expect(screen.queryByTestId('first')).not.toBeInTheDocument();
+      expect(screen.getByTestId('second')).toHaveFocus();
+    });
+
+    it('ignores a marked element portaled outside the root', () => {
+      const otpStep = (
+        <Flow.Step ids={['otp']}>{createPortal(<AutoFocusInput data-testid='portaled' />, document.body)}</Flow.Step>
+      );
+      const { rerender } = render(
+        <Flow.Root value='password'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep}
+        </Flow.Root>,
+      );
+
+      rerender(
+        <Flow.Root value='otp'>
+          <Flow.Step ids={['password']}>Password</Flow.Step>
+          {otpStep}
+        </Flow.Root>,
+      );
+
+      expect(screen.getByTestId('portaled')).not.toHaveFocus();
+    });
+
+    it('returns a no-op ref outside a step', () => {
+      expect(() => render(<AutoFocusInput data-testid='lone-input' />)).not.toThrow();
+      expect(screen.getByTestId('lone-input')).toBeInTheDocument();
+    });
+  });
 });

packages/headless/src/primitives/flow/index.ts

@@ -1,3 +1,4 @@
 export * as Flow from './parts';
+export { useFlowAutoFocus } from './flow-step-context';
 
 export type { FlowDirection, FlowRootProps, FlowStepProps } from './parts';

packages/headless/src/primitives/flow/parts.ts

@@ -1,3 +1,4 @@
 export { type FlowRootProps, FlowRoot as Root } from './flow-root';
 export { type FlowStepProps, FlowStep as Step } from './flow-step';
+export { useFlowAutoFocus } from './flow-step-context';
 export type { FlowDirection } from './flow-context';

packages/ui/src/mosaic/components/flow/flow.styles.ts

@@ -4,7 +4,7 @@ import { durationVars, easingVars } from '../../tokens.stylex';
 
 export const styles = stylex.create({
   root: {
-    overflow: 'hidden',
+    overflow: 'clip',
     position: 'relative',
     transitionDuration: {
       default: durationVars['--cl-duration-slow'],

packages/ui/src/mosaic/components/flow/index.ts

@@ -1,2 +1,3 @@
+export { useFlowAutoFocus } from '@clerk/headless/flow';
 export { Flow } from './flow';
 export type { FlowDirection, FlowRootProps, FlowStepProps } from './flow';