createState function using React Context and ImmerJS

import React, { useContext, useEffect, useReducer } from "react";
import { useRouter } from "next/router";
import produce from "immer";
import mapValues from "lodash/mapValues";

type Children = { children?: any };
type ProviderFC = React.FC<Children>;

type MutationFn<T> = (state: Partial<T>) => void;
type Mutations<T> = {
  [name: string]: (state: T, ...args: any[]) => void;
};

type Update<T, M extends Mutations<T>> = { setState: MutationFn<T> } & M;
type UseStateResult<T, M extends Mutations<T>> = [T, MutationFn<T>, M];
type UseStateFn<T, M extends Mutations<T>> = () => UseStateResult<T, M>;

export default function createState<T, M extends Mutations<T>>(
  {
    initialState,
    loadState = () => Promise.resolve(initialState),
  }: {
    initialState: T;
    loadState?: () => Promise<T>;
  },
  mutations?: M
): [ProviderFC, UseStateFn<T, M>] {
  const StateContext = React.createContext<T>(initialState);
  const UpdateContext = React.createContext<Update<T, M>>(null as any);

  function Provider({ children }) {
    const router = useRouter();
    const [state, update] = useReducer(produce, initialState);
    const mutate = update as any;

    const setState = (newState: Partial<T>) => {
      mutate((target: any) => {
        Object.assign(target, newState);
      });
    };

    const actions = mapValues(mutations, (fn) => (...params: any[]) => {
      mutate((target: any) => {
        fn(target, ...params);
      });
    });

    const loadAsync = async () => {
      const newState = await loadState();
      mutate((target: any) => {
        Object.assign(target, newState);
      });
    };

    // TODO avoid this effect if loadAsync is not defined
    useEffect(() => {
      loadAsync();
    }, [router.pathname]);

    return (
      <UpdateContext.Provider value={{ setState, ...actions } as any}>
        <StateContext.Provider value={state as any}>
          {children}
        </StateContext.Provider>
      </UpdateContext.Provider>
    );
  }

  function useState(): UseStateResult<T, M> {
    const { setState, ...mutations } = useContext(UpdateContext);
    return [useContext(StateContext), setState, mutations as any];
  }

  return [Provider, useState];
}

// auth state
const initialState = {
  id: "",
  name: "",
  email: "",
  is_owner: false,
  is_admin: false,
  is_authenticated: !!getApiToken(),
};

type State = typeof initialState;

const [UserProvider, useUser] = createState({
  initialState,
  loadState: fetchAuthState,
});

export { UserProvider, useUser };

async function fetchAuthState(): Promise<State> {
  if (!getApiToken()) {
    return initialState;
  }
  try {
    const resp = await api.me();
    return {
      ...initialState,
      ...resp.data,
      is_authenticated: true,
    };
  } catch (err) {
    return initialState;
  }
}

// example login page
function LoginPage() {
   const [_, setUser] = useUser();
   const handleSubmit = async () => {
       const resp = await api.login(form);
       setUser({
            ...omit(resp.data, ['token', "token_expired_at"]),
          is_authenticated: true,
       });
   };
   return <div> Login Form </div>
}