React Suspense と Zustand の createStore を使って型安全に状態管理を行う

こんにちは、カケハシのAI在庫管理チームでフロントエンドエンジニアをしている Nokogiri です。
AI在庫ではフロントエンドをReactで実装しておりますが、サーバーとの通信には GraphQL（Apollo Client）を利用し、状態管理の一部に Zustand を利用しています。

AI在庫の開発時に GraphQL と Zustand を使ったことで発生した課題とそれに対してどのように改善したか？ということについてご紹介したいと思います。 Zustandに限らずフロントエンドで状態管理を行うときに頻繁に発生する課題だと思うので参考になれば幸いです。

課題1：サーバーから取得した値を状態管理する場合、初期値が定まらない

例えば医薬品情報を参照・編集する機能があるとします。医薬品情報はサーバーから取得して表示し編集後にまた更新後の値をサーバーに送信したいとします。

サーバーから取得した値を一度 Zustand で状態管理し保存時に取り出すことを想定します。

Zustand の create を利用することでファイル読み込み時に store を初期化します。 medicine はサーバーから取得した結果を受け取るので store を初期化する時点では値がありません。

ここでは値がないことを表現するために null を設定しています。

またサーバーから取得した値を Zustand の Store で管理するために useEffect を使って medicine を更新しています。

オーソドックスな Zustand を使った状態管理の例

// Zustand の Stateの初期化処理
const useMedicineStore = create<State>((set) => ({
  medicine: null,
  medicineLoaded: (medicine: Medicine) => set({ medicine }),
}));

// 医薬品情報を更新して保存するコンポーネント
function MedicineForm() {
  const { medicine } = useMedicineStore();
  if (!medicine) {
    return null;
  }
  ...
}

function MedicineInfo() {
  const { medicineLoaded } = useMedicineStore();
  const { data, loading } = useQuery(GET_MEDICINE);

  useEffect(() => {
    if (data && !loading) {
      medicineLoaded(data.medicine);
    }
  }, [data, loading, medicineLoaded]);

  if (loading) {
    return <Loading />
  }

  return (
    <MedicineForm />
  );
}

この実装の問題点としては2つです。

MedicineForm で medicine が nullable であるかもしれないので、事前に null を返すガードを書く必要があること
Apollo Client の useQuery をせっかく使っているのに、そのまま戻り値の data を扱うことができず Zustand で管理するために useEffect が必要になる

課題2：非同期処理の終了と Zustand の状態の初期化処理タイミングが合わない

課題1を解決するために Suspense を使って特定の階層のコンポーネント以下では data があるという状態を実現しようとしました。しかしこの場合でもMedicineFormからするとZustandで管理するmedicineはnullableであり、ガードが必要になります。

もちろん medicine の型は nullable ですが、実際も MedicineContainer が Suspend（一時停止）している状態が終了しても、Zustand の Store の初期化が終了していない（一瞬遅れてmedicineが入る）ためコンポーネント側ではmedicineをnullableとして扱う必要があります。

// 医薬品情報を取得している間Promiseをthrowするコンポーネント
function MedicineContainer() {
  const { data } = useSuspenseQuery(GET_MEDICINE);

  useEffect(() => {
    if (data) {
      medicineLoaded(data.medicine);
    }
  }, [data, medicineLoaded]);

  return <MedicineForm />;
}

function useMedicine() {
  const { medicine } = useMedicineStore();
  // ❌ Suspense を利用したとしても、medicine は null になるタイミングがあるので例外がスローされる
  // この書き方はできない
  if (medicine === null) {
    throw new Error('medicine is Null!!!')
  }
  return medicine;
}

function MedicineForm() {
  const medicine = useMedicine();
  ...
}

// Suspend中のfallbackを指定
function App() {
  return (
    <React.Suspense fallback={<Loading />}>
      <MedicineContainer />
    </React.Suspense>
  );
}

解決策：Suspense + Provider の組み合わせ

実現したいこととしては以下二つです。

Suspense を利用して特定のコンポーネント以下では必ずデータがある前提としたい
データがある前提であれば型情報を nullable で扱いたくない

Zustand の `createStore` を利用して State の初期値を Suspend 終了後に作成する

Zustand の公式サイトに Initialize state with props という記述があります。

Zustand の create 関数を利用して hooks を作成する通常の手順の場合、JavaScriptが画面にロードされたタイミングで初期化されます。

例）通常の create 関数を使った State の初期化

// create 関数は、Storeを扱う hooks を生成します。
const useMedicineStore = create<State>((set) => ({
  medicine: null,
  updateMedicine: (medicine: Medicine) => set({ medicine }),
}));

function Medicine() {
  const medicine =  useMedicineStore()
  ...
}

createStore を使うことでhooksではなく、store のインスタンス自体を生成することができます。 hooksとは異なり、store の生成を任意のタイミングで呼び出せるようなるため、データ取得が完了した後で Store を生成することができ、データを nullable にすることを防ぐことができます。

const createMedicineStore = (medicine: Medicine) => {
  return createStore<State>((set) => ({
    medicine,
    updateMedicine: (medicine: Medicine) => set({ medicine }),
  }))
}

これで createMedicineStore を任意のタイミングで呼び出すことができ、useSuspenseQuery と組み合わせることで medicine を nullable でなく扱うことができるようになりました。

function MedicineContainer() {
  const { data } = useSuspenseQuery(GET_MEDICINE);
  const store = createMedicineStore(data.medicine);
  ...
}

ただ実際にこのやり方では取得した store 自体を子供に渡さないと扱えません。

Zustand を使いたいユースケースとしては Props のバケツリレーが回避が目的の一つになっていると思うのでこのままでは本末転倒です。

Zustand の公式では createContextを使った方法で子供のコンポーネントから store にアクセスする方法を紹介しています。

AI在庫でもこの方法に従いました。

// Store の型情報は ReturnType を利用すると便利です
type MedicineStore = ReturnType<typeof createMedicineStore>;

const MedicineStoreContext = createContext<MedicineStore | null>(null);

function MedicineStoreProvider({ children, medicine }: Props) {
  const storeRef = useRef<MedicineStore>();
  if (!storeRef.current) {
    storeRef.current = createMedicineStore({ medicine });
  }

  return (
    <MedicineStoreContext.Provider value={storeRef.current}>
      {children}
    </MedicineStoreContext.Provider>
  );
}

子供のコンポーネントからはContextから取り出す処理を意識させないために以下のような汎用的な hooks を作成しています。

function useMedicineStore<T>(
  selector: (state: MedicineState, ...args: unknown[]) => T,
): T {
  const store = useContext(MedicineStoreContext);
  if (!store) throw new Error('Missing MedicineProvider in the tree');
  return useStore(store, selector);
}

Provider と組み合わせて子供のコンポーネントで取り出すには以下のように記述します。

function App() {
  return (
    <React.Suspense fallback={<Loading />}>
      <MedicineContainer />
    </React.Suspense>
  );
}

function MedicineContainer() {
  const { data } = useSuspenseQuery(GET_MEDICINE);
  const store = createMedicineStore(data.medicine)

  return (
    <MedicineStoreProvider medicine={data.medicine}>
      <MedicineForm />
    </MedicineStoreProvider>
  )
}


function MedicineForm() {
  // ここでの medicine は null ではないので安全に扱える
  const medicine = useMedicineStore(state => state.medicine);
  ...
}

Suspense + Providerを使うことで実装が少し増えますが、サーバーから取得した値を state の初期化に利用することができ Zustand の State にアクセスする際に nullable な状態を避けることができます。

まとめ

サーバー通信に GraphQL（Apollo Client）と Zustand による状態管理を組み合わせた時に以下の課題発生しました。

取得データの初期値が不定
サーバーからデータが返ってくるまで medicine などの値が null となるため、コンポーネントでガードが必要になり、型安全性が下がる。
サーバー通信完了と Zustand Stateの初期化のタイミングが異なる
Suspense を利用しても、Zustand の初期値が画面ロード時に設定されるため、nullable な状態を完全には避けられない。

これらを解決するために以下のアプローチを採用しました。

Zustand の createStore を使ってストアを生成
通常の create 関数では画面ロード時にストアが初期化されるが、createStore なら任意のタイミングでストア生成が可能。
→ サーバーからデータを取得した後にストアを初期化すれば、ストア内の値を非 nullable として扱える。
Context + Provider でストアを下層コンポーネントへ提供
createStore で生成したストアを Context にセットし、子コンポーネントからは useStore で直接アクセス。
→ Props のバケツリレーを避けつつ、Nullable なデータを使う必要がなくなる。
Suspense との組み合わせ
useSuspenseQuery でデータ取得が完了したら createStore を実行し、Provider を通じて下層コンポーネントに渡す。
→ 「データが必ずある」前提でコンポーネントを実装でき、型安全に状態を参照できる。