Jotai and a Pseudo-Flux Pattern Behind Custom Hooks

You don’t get far without state in a frontend app.

First, we used Redux. Then we moved from our internal framework to Next.js, and from Flow to TypeScript. Along the way, we tried Context, useReducer, and Recoil. Meta archived the Recoil repository while React 19 support remained unresolved, so we moved to Jotai.

I had put most of our Recoil calls behind custom hooks. That helped us move to Jotai without changing state access throughout the components. We still had to change the implementation, but we had fewer places to touch.

I also liked that Daishi Kato was active in open source, maintaining Jotai and other React libraries. That gave me more confidence in Jotai’s future maintenance. It was a reason to choose it, though I still wanted the hooks in between so we could handle another change later.

The Kind of App I Had in Mind

I’m thinking of medium-to-large B2B applications and internal company systems. In that setting, I want components to call operations that belong to the application, without needing to know which state library we use underneath.

I still liked the dispatch and action pattern from Redux. Giving an update a name lets me collect related changes in one place and follow what a component asks the app to do. With Jotai, I could keep that shape while letting each component subscribe to individual fields.

This is a convention I found useful for that kind of codebase. A small form may be easier to write with Jotai’s hooks alone. I’m not attached to one stack; this was a way to keep a familiar API while changing the implementation behind it.

A Small Example

The code below is a simplified example of that boundary, not the whole architecture of a large application. It leaves out validation and API calls. The atoms stay inside this module, and components use the custom hook.

I call it “pseudo-Flux” because we dispatch typed actions, but the handler writes to atoms. A Redux reducer is a pure function that returns the next state from the previous state and an action. This handler does not have that contract.

import type { Getter, Setter } from "jotai";
import { atom, useAtomValue } from "jotai";
import { useAtomCallback } from "jotai/utils";
import { useCallback, useMemo } from "react";

const name = atom("");
const age = atom<number | undefined>(undefined);
const mail = atom("");
const occupation = atom("");

// Action type definition: This follows the Redux pattern of using discriminated unions
// to ensure type safety when dispatching actions. Each action type has its own payload.
type Action =
  | {
      type: "update_name";
      name: string;
    }
  | {
      type: "update_age";
      age: number;
    }
  | {
      type: "update_mail";
      mail: string;
    }
  | {
      type: "update_occupation";
      occupation: string;
    }
  | { type: "reset_form" };

// Collection of all user form atoms for easy access
// This allows us to dynamically access atoms by key name
const userFormState = {
  name,
  age,
  mail,
  occupation,
};

// Subscribe to the selected field. Parent or local state updates can still re-render the component.
const useUserFormStateProps = <K extends keyof typeof userFormState>(key: K) => {
  const baseAtom = userFormState[key];
  const returnValue = useAtomValue(baseAtom);
  return returnValue;
};

// Using dispatch alone does not subscribe the component to the form values.
export const useUserFormState = () => {
  // _get can read current atom values if an action needs them.
  // useCallback keeps the callback passed to useAtomCallback stable.
  const dispatch = useAtomCallback(
    useCallback((_get: Getter, set: Setter, action: Action) => {
      switch (action.type) {
        case "update_name":
          // Update the name atom with the new value
          set(name, action.name);
          return;
        case "update_age":
          set(age, action.age);
          return;
        case "update_mail":
          set(mail, action.mail);
          return;
        case "update_occupation":
          set(occupation, action.occupation);
          return;
        case "reset_form":
          set(name, "");
          set(age, undefined);
          set(mail, "");
          set(occupation, "");
          return;
      }
    }, [])
  );
  // memoize the return value to prevent creating new objects on every render
  return useMemo(() => {
    return {
      dispatch,
      useUserFormStateProps,
    };
  }, [dispatch]);
};

Here, an empty string means an empty text field, and undefined means no age has been entered. A component can read one field and dispatch an operation:

// In a component module, import useUserFormState from your state module.
export const NameField = () => {
  const { dispatch, useUserFormStateProps } = useUserFormState();
  const name = useUserFormStateProps("name");

  return (
    <input
      aria-label="Name"
      value={name}
      onChange={(event) => dispatch({ type: "update_name", name: event.target.value })}
    />
  );
};

export const ResetButton = () => {
  const { dispatch } = useUserFormState();
  return (
    <button type="button" onClick={() => dispatch({ type: "reset_form" })}>
      Reset
    </button>
  );
};

NameField subscribes to the name atom. ResetButton only gets dispatch, so it does not subscribe to the form values. Call the reader hook at the top level of a component, following the usual Rules of Hooks.

The reset_form action is where the extra layer starts to help. The button asks to reset the form, and the handler knows which fields to clear. If we add another field, we update the reset operation in this module. The button does not need another setter.

These module-level atoms share values within the same Jotai store. For independent copies of the form, give each one its own store, for example with a separate Jotai Provider. This example assumes one shared form in that scope.

Why Keep This Layer?

For one setter per action, this can look like extra typing. In a larger app, I want a place to group operations that affect several fields and to keep those details out of components. The hooks also give us a smaller surface to preserve if we change the state library again. They don’t make a migration free; we still need to check subscription behavior and how updates work.

Jotai already supports this kind of operation through write-only atoms. We could put the same action handler in a write-only atom and expose it through useSetAtom. useAtomCallback is one way to implement this convention, not a requirement for it.

The callback can use get to read current atom values and set to update other atoms in the same store. That gives us room to handle related updates, but we still need to decide which state belongs to each operation. Otherwise, we can end up with a handler that knows too much about the rest of the app.

This also does not give us Redux’s action history or replay tools by itself. Using hooks with both libraries does not make their state models the same. For me, the useful part was keeping typed actions and field-level subscriptions behind an API our components already understood.

I used AI to help check the accuracy of the explanations in this article.