I'd like to use an union type like this to represent my slice's state when using Redux Toolkit:

type AppState = { loading: true } | { loading: false; data: number };

It seems like the usage of Immerjs in the reducer files prevents me from switching between the two Union options, e.g.:

import { createSlice, PayloadAction } from "@reduxjs/toolkit";

const initialState: AppState = { loading: true };

export const sampleSlice = createSlice({
  name: "sample",
  initialState,
  reducers: {
    setup: (state) => {
      return { loading: false, data: 5 };
    },
  },
});

Which gives me the following Typescript error:

Type '(state: WritableDraft<{ loading: true; }>) => { loading: false; data: number; }' is not assignable to type 'CaseReducer<{ loading: true; }, { payload: any; type: string; }> | CaseReducerWithPrepare<{ loading: true; }, PayloadAction<any, string, any, any>>'.

Is there any way to make this work? My current workaround is not using an union type to represent the state but I'm wondering if there is a workaround. I have also tried not using booleans but string literals for the loading key but it produces the same error.