Use Controlled State hook manages a value that can be either controlled or uncontrolled. It returns the current state and a setter that updates internal state when uncontrolled and always calls an optional onChange callback.

use-controlled-state.ts

import { useCallback, useLayoutEffect, useRef, useState } from 'react';
import { flushSync } from 'react-dom';

export function useControlledState<T>(
    initialValue: T,
    controlledValue?: T,
    onChange?: (value: T) => void,
): [T, (value: T | ((prev: T) => T)) => void] {
    const [internalValue, setInternalValue] = useState(initialValue);

    const isControlled = controlledValue !== undefined;

    const stateRef = useRef<T>(initialValue);

    const controlledRef = useRef<T | undefined>(controlledValue);

    useLayoutEffect(() => {
        stateRef.current = internalValue;
    }, [internalValue]);

    useLayoutEffect(() => {
        controlledRef.current = controlledValue;
    }, [controlledValue]);

    return [
        isControlled ? controlledValue : internalValue,
        useCallback(
            (next: T | ((prev: T) => T)) => {
                const prev = isControlled ? controlledRef.current! : stateRef.current;

                const resolved = typeof next === 'function' ? (next as (p: T) => T)(prev) : next;

                flushSync(() => (isControlled ? onChange?.(resolved) : setInternalValue(resolved)));

                if (!isControlled) {
                    onChange?.(resolved);
                }
            },
            [isControlled, onChange, setInternalValue],
        ),
    ] as const;
}