Jotai (not exactly Redux) Dev Tools

#reduxdevtools #debugging #jotai #timetravel

Recently, we have begun investigating Jotai, a state management alternative to Recoil by Poimandres. In addition, we documented some light investigation converting an existing Recoil app to Jotai in a previous article. While performing this exercise it got us thinking about what debug tools that might be available which are specifically targeted to Jotai. Unfortunately, as of this writing all that is available is the useDebugValue hook.

Now, like most folks, we enjoy using powerful and versatile dev tools such as React DevTools, React Query DevTools, etc. While we have not used Redux in a production capacity in the past, (for a number of reasons) we have always preferred to use ReduxDevTools for debugging React state management systems. In fact, we have created these custom internal plugins for each state management system we have used in our contracting work. As of this writing, we couldn't find a Jotai plugin yet, so naturally we thought it might be interersting to attempt to create one.

First, we setup some interfaces/types for Config, Message, ConnectionResult, Extension, etc. to match the payloads transmitted from/to ReduxDevTools. These interfaces enable strongly typed mapping for communication with the Browser Extension proper.

interface Config {
    instanceID?: number,
    name?: string,
    serialize?: boolean,
    actionCreators?: any,
    latency?: number,
    predicate?: any,
    autoPause?: boolean
}

interface Message {
    type: string,
    payload?: any,
    state?: any
}

interface IConnectionResult {
    subscribe: (dispatch: any) => {};
    unsubscribe: () => {};
    send: (action: string, state: any) => {};
    error: (payload: any) => {};
}

type ConnectionResult = IConnectionResult & Observable<Message>

interface Extension {
    connect: (options?: Config) => ConnectionResult;
}

Unfortunately, the Redux DevTools connect function used for the browser extension frontend does not return a standard RxJS Observable. Consequently, we need the wrapConnectionResult function to make a RxJS compatible observable to receive ReduxDevTools events.

const wrapConnectionResult = (result: ConnectionResult) => {
    const subject = new Subject<Message>()

    result.subscribe(
        (x: any) => subject.next(x),
        (x: any) => subject.error(x),
        () => {subject.complete()}
    );

    return subject;
}

Now we'll implement a JotaiDevtoolsProps interface so that the user can name and configure DevTools instances.

interface JotaiDevtoolsProps {
    atom: WritableAtom<any, any>,
    name?: string,
    config?: Config
}

To enable seamless integration with React Functional Components, we create a hook that will take our JotaiDevtoolsProps and instantiate an instance for a particular Jotai atom.

...
export const useJotaiDevtools = ({ atom, name, config, ...props }: JotaiDevtoolsProps) => {
...

Next, we'll use the observable-hooks library's useObservable hook to provide the value of the incoming atom as an RxJS observable.

    const atom$ = useObservable((input$) => input$.pipe(
        map(([x]) => x)
    ), [atomCurrentValue]);

In order to prevent a race condition when an atom is updated, we'll add a flag to determine whether the last state update was a ReduxDevTools Time Travel Event.

    const [wasTriggeredByDevtools, setWasTriggeredByDevtools] = useState(() => false);

Additionally, we need a flag to determine if the initial state has already been sent to the DevTools Extension.

    const [sentInitialState, setSentInitialState] = useState(() => false);

    const [devTools, setDevTools] = useState<ConnectionResult>();

Since the DevTools connection may not be availiable when the hook is initialized, we will add another useObservable to provide a sanitized stream of events when the connection is ready.

    const devTools$ = useObservable((input$) => input$.pipe(
        filter(([x]) => !!x),
        switchMap(([x]) => wrapConnectionResult(x as ConnectionResult)),
        catchError((error, observable) => {
            console.error(error);
            return observable;
        })
    ), [devTools]);

This function is called to handle State Jumps and Time Travel events.

    const jumpToState = (newState: any) => {
        setWasTriggeredByDevtools(true)
        // var oldState = atomCurrentValue();
        setAtom(newState);
        // setWasTriggeredByDevtools(false);
    };

Continuing our previous pattern, we will use observable-hook's useSubscription hook to subscribe to the DevTools Extension events and respond appropriately to either either a START or DISPATCH action.

    useSubscription(devTools$, (message) => {
        switch (message.type) {
            case 'START':
                console.log("Atom Devtools Start", options.name, atomCurrentValue)
                if(!sentInitialState) {
                    // devTools.send("\"" + options.name + "\" - Initial State", atom.getState());
                    devTools?.send(name + " - Initial State", atomCurrentValue);
                    setSentInitialState(true);
                }
                return;
            case 'DISPATCH':
                if (!message.state) {
                    return;
                }
                switch (message.payload.type) {
                    case 'JUMP_TO_ACTION':
                    case 'JUMP_TO_STATE':
                        jumpToState(JSON.parse(message.state));
                        return;
                }
                return;
        }
    });

Next, we will need another useSubscribe hook to subscribe updates coming from application state changes to the current Jotai atom.

    useSubscription(atom$, (state) => {
        if (wasTriggeredByDevtools) {
            setWasTriggeredByDevtools(false);
            return;
        }
        devTools?.send(name + " - " + moment().toISOString(), state);
    })

When the component is initialized, we will need a function to get a reference to the DevTools Extension. If the extension is not installed, an error will be logged to the console.

    const initDevtools = () => {
        const devtools = (window as any).__REDUX_DEVTOOLS_EXTENSION__ as Extension;
        // const options = config;
        const name = options.name || window.document.title;

        if (!devtools) {
            console.error('Jotai Devtools plugin: Cannot find Redux Devtools browser extension. Is it installed?');
            return atom;
        }

        const devTools = devtools.connect(options);

        console.log("Get Dev Tools", devTools, of(devTools));

        setDevTools(devTools);

        // setTimeout(() => devTools.send(name + " - Initial State", atomCurrentValue), 50)

        return atom;
    }

In order to activate our initialization function, we will utilize the useLifeCycles hook (from the excellent react-use library) to handle the component mount lifecycle event.

    useLifecycles(() => {
        initDevtools();
    });

That's it. Now just install the new Jotai Devtools plugin in any projects that utilize Jotai atoms.

Specifically, we can just call the useJotaiDevtools hook for each atom you want to view in the Devtool Browser Extension

...
    useJotaiDevtools({
        name: "Dark Mode",
        atom: darkModeState
    });
...

    useJotaiDevtools({
        name: "Tasks",
        atom: tasksAtom
    });

...

For illustration, we can re-use the Recoil example we converted to Jotai in a previous post. Once the app is started, we can open the Redux DevTools Browser Extension and we will be able to watch our state changes, time travel debug, etc.

Final Thoughts:

By leveraging the existing work done on ReduxDevTools, we have access to a useful debugging aid without reinventing the wheel.
Leveraging the observable-hooks and react-use libraries enabled clean and efficient ReduxDevTools integration.
Adapting the time travel functionality of ReduxDevTools enables full replay of the Jotai chain of state.
The result is a compelling insight into the Jotai atom lifecycle.
Check out the Jotai Devtools Repo on Github!

Below is a functioning example. If you haven't already, make sure you install the Redux DevTools extension to be able to see the state update.