setting up React store using Immer - reactjs

Question

I am trying to set up store in React using immer and borrowed the code from another project. However it doesn't work for me. Could be that immer version on borrowed project is lower than it is now.
Store looks like this:
const initialState: MyState = {
search: "No results yet"
}
const useValue = () => useState(initialState);
const { Provider, useTrackedState, useUpdate: useSetState } = createContainer(
useValue
);
const useSetDraft = () => {
const setState = useSetState();
return useCallback(
(draftUpdater: MyState) => {
setState(produce(draftUpdater));
},
[setState]
);
};
export { Provider, useTrackedState, useSetDraft, initialState }
and hook that i use to mutate state is:
const UseSetSearch = () => {
const setDraft = useSetDraft();
return useCallback((searchString: string) => {
setDraft((draft: MyState) => {
draft.search = searchString
return draft
})
}, [setDraft]);
}
With the store, I get error
Type 'MyState' provides no match for the signature '(state: WritableDraft): ValidRecipeReturnType'
and hook error is
Value of type '(draft: MyState) => MyState' has no properties in common with type 'MyState'. Did you mean to call it?
Store error is fixed when I do:
const useSetDraft = () => {
const setState = useSetState();
return useCallback(
(draftUpdater: MyState) => {
setState(produce(draftUpdater, draft => {
}));
},
[setState]
);
};
but hook error remains. What am I missing?