atomWithUndo does not return the target value

README.md

@@ -24,14 +24,15 @@ The history atom tracks changes to the `targetAtom` and maintains a list of prev
 ### Usage
 
 ```tsx
-import { atom } from 'jotai'
+import { atom, useAtomValue, useSetAtom } from 'jotai'
 import { atomWithHistory } from 'jotai/utils'
 
 const countAtom = atom(0)
 const countWithPrevious = atomWithHistory(myAtom, 2)
 
 export function CountComponent() {
-  const [[count, previousCount], setCount] = useAtom(countWithPrevious)
+  const [count, previousCount] = useAtomValue(countWithPrevious)
+  const setCount = useSetAtom(countAtom)
 
   return (
     <>
@@ -69,7 +70,7 @@ The returned object includes:
 ### Usage
 
 ```tsx
-import { atom } from 'jotai'
+import { atom, useAtom, useAtomValue } from 'jotai'
 import { atomWithUndo } from 'jotai/utils'
 
 const counterAtom = atom(0)

src/atomWithHistory.ts

@@ -8,19 +8,13 @@ import type { Atom } from 'jotai/vanilla'
  */
 export function atomWithHistory<T>(targetAtom: Atom<T>, limit: number) {
   const refAtom = atom(
-    () => ({
-      history: [] as T[],
-    }),
-    (get) => () => {
-      get(refAtom).history.length = 0
-    }
+    () => ({ history: [] as T[] }),
+    (get, _set) => () => void (get(refAtom).history.length = 0)
   )
   refAtom.onMount = (mount) => mount()
   refAtom.debugPrivate = true
   return atom((get) => {
     const ref = get(refAtom)
-    const value = get(targetAtom)
-    ref.history = [value, ...ref.history].slice(0, limit)
-    return ref.history
+    return (ref.history = [get(targetAtom), ...ref.history].slice(0, limit))
   })
 }

src/atomWithUndo.ts

@@ -19,11 +19,16 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
   const UNDO = Symbol('undo')
   const REDO = Symbol('redo')
   type DoAction = typeof UNDO | typeof REDO
-  const refAtom = atom(() => ({
+  const createRef = () => ({
     index: 0,
-    stack: [] as T[][],
+    stack: [] as T[],
     action: null as DoAction | null,
-  }))
+  })
+  const refAtom = atom(
+    createRef,
+    (get, _set) => () => void Object.assign(get(refAtom), createRef())
+  )
+  refAtom.onMount = (mount) => mount()
   refAtom.debugPrivate = true
   const updateRefAtom = atom(
     (get) => {
@@ -36,16 +41,16 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
         // Remove future states if any
         ref.stack = ref.stack.slice(0, ref.index + 1)
         // Push the current state to the history
-        ref.stack.push(history.slice())
+        ref.stack.push(history[0] as T)
         // Limit the history
         ref.stack = ref.stack.slice(-limit)
         // Move the current index to the end
         ref.index = ref.stack.length - 1
       }
-      return null
+      return { ...ref }
     },
     (get) => {
-      get(refAtom).stack = [[get(targetAtom)]]
+      get(refAtom).stack = [get(targetAtom)]
       return () => {
         const ref = get(refAtom)
         ref.index = 0
@@ -55,27 +60,24 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
   )
   updateRefAtom.onMount = (mount) => mount()
   updateRefAtom.debugPrivate = true
+  const canUndoAtom = atom((get) => get(updateRefAtom).index > 0)
+  const canRedoAtom = atom((get) => {
+    const ref = get(updateRefAtom)
+    return ref.index < ref.stack.length - 1
+  })
   const baseAtom = atom<Undoable, [DoAction], void>(
-    (get, { setSelf }) => {
-      get(updateRefAtom)
-
-      return {
-        undo: () => setSelf(UNDO),
-        redo: () => setSelf(REDO),
-        get canUndo() {
-          return get(refAtom).index > 0
-        },
-        get canRedo() {
-          return get(refAtom).index < get(refAtom).stack.length - 1
-        },
-      }
-    },
+    (get, { setSelf }) => ({
+      undo: () => setSelf(UNDO),
+      redo: () => setSelf(REDO),
+      canUndo: get(canUndoAtom),
+      canRedo: get(canRedoAtom),
+    }),
     (get, set, update) => {
       const ref = get(refAtom)
       const setCurrentState = () => {
-        const currentSlice = ref.stack[ref.index]
-        if (currentSlice?.[0] === undefined) return
-        set(targetAtom, currentSlice[0])
+        const value = ref.stack[ref.index]
+        if (value === undefined) return
+        set(targetAtom, value as T)
       }
       if (update === UNDO) {
         if (get(baseAtom).canUndo) {