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

@@ -7,15 +7,9 @@ import type { Atom } from 'jotai/vanilla'
  * @returns an atom with an array of history states
  */
 export function atomWithHistory<T>(targetAtom: Atom<T>, limit: number) {
-  const refAtom = atom(
-    () => ({
-      history: [] as T[],
-    }),
-    (get) => () => {
-      get(refAtom).history.length = 0
-    }
-  )
-  refAtom.onMount = (mount) => mount()
+  const refAtom = atom(() => ({
+    history: [] as T[],
+  }))
   refAtom.debugPrivate = true
   return atom((get) => {
     const ref = get(refAtom)

src/atomWithUndo.ts

@@ -21,7 +21,7 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
   type DoAction = typeof UNDO | typeof REDO
   const refAtom = atom(() => ({
     index: 0,
-    stack: [] as T[][],
+    stack: [] as T[],
     action: null as DoAction | null,
   }))
   refAtom.debugPrivate = true
@@ -36,7 +36,7 @@ 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
@@ -45,7 +45,7 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
       return null
     },
     (get) => {
-      get(refAtom).stack = [[get(targetAtom)]]
+      get(refAtom).stack = [get(targetAtom)]
       return () => {
         const ref = get(refAtom)
         ref.index = 0
@@ -73,9 +73,9 @@ export function atomWithUndo<T>(targetAtom: PrimitiveAtom<T>, limit: number) {
     (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) {