typescript - Tipizzare il metodo `set` dello store Zustand con il middleware Immer mantenendo l'autocompletamento

Ho problemi a digitare il metodo set dello store di Zustand mentre utilizzo il middleware Immer. Devo digitare questo metodo poiché alcune proprietà contengono funzioni che aggiorneranno lo stato, quindi hanno bisogno della funzione set:

export const useStore = create<Store>()(
     immer((set, get) => ({
         func: beginDraggingVertex(set, get, vertexID) 
     }))
)     

....

function beginDraggingVertex(set: Setter, get: Getter, vertexID: string) {
  ...ommited
  set((store) => 
   // here is where I do want to have the autocomplete working
   store.whateverProperty = 'hi'
  )

}

Se rimuovo il middleware Immer dallo store, quando passo il cursore sul metodo set ottengo il suo tipo così posso successivamente digitare gli argomenti nella dichiarazione di beginDraggingVertex. Ma con Immer applicato, il metodo set cambia e se cambio il tipo di set al nuovo tipo, perdo l'autocompletamento.

 //`set` type with Immer applied, with this one I loose autocomplete wihtin:
 type Setter = (
  nextStateOrUpdater:
    | (Store & ActionsObject)
    | Partial<Store & ActionsObject>
    | ((state: WritableDraft<Store & ActionsObject>) => void),
  shouldReplace?: boolean | undefined
) => void

//`set` type without Immer:
 export type Setter = (
  partial:
    | (Store & ActionsObject)
    | Partial<Store & ActionsObject>
    | ((state: Store & ActionsObject) => (Store & ActionsObject) | Partial<Store & ActionsObject>),
  replace?: boolean | undefined
) => void

Se applico Immer e poi do a set il vecchio tipo di quando non era applicato, ovviamente ottengo l'autocompletamento ma ricevo comunque questi errori sia nel callback di set:

I was expecting (Store & ActionsObject) | Partial<Store & ActionsObject> | ((state: Store & ActionsObject) => (Store & ActionsObject) | Partial<...>), but you passed (store: Store & ActionsObject) => void

Apprezzo davvero qualsiasi aiuto tu possa fornire.