fix: regression in #1176 (#1206)

Author: dai-shi

src/react.ts

@@ -127,25 +127,21 @@ export function useSnapshot<T extends object>(
     [proxyObject],
   )
   const lastSnapshot = useRef<Snapshot<T>>(undefined)
-  const subscribed = useRef(false)
+  let inRender = true
   const currSnapshot = useSyncExternalStore(
     useCallback(
       (callback) => {
-        subscribed.current = true
         const unsub = subscribe(proxyObject, callback, notifyInSync)
         callback() // Note: do we really need this?
-        return () => {
-          unsub()
-          subscribed.current = false
-        }
+        return unsub
       },
       [proxyObject, notifyInSync],
     ),
     () => {
       const nextSnapshot = snapshot(proxyObject)
       try {
         if (
-          subscribed.current &&
+          !inRender &&
           lastSnapshot.current &&
           !isChanged(
             lastSnapshot.current,
@@ -164,6 +160,9 @@ export function useSnapshot<T extends object>(
     },
     () => snapshot(proxyObject),
   )
+  // TODO how could we bypass this?
+  // eslint-disable-next-line react-hooks/immutability
+  inRender = false
   useLayoutEffect(() => {
     lastSnapshot.current = currSnapshot
   })

tests/basic.test.tsx

@@ -1,4 +1,4 @@
-import { StrictMode, useEffect, useState } from 'react'
+import { StrictMode, useEffect, useLayoutEffect, useState } from 'react'
 import { act, fireEvent, render, screen } from '@testing-library/react'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { proxy, snapshot, useSnapshot } from 'valtio'
@@ -399,6 +399,47 @@ describe('basic', () => {
     expect(screen.getByText('count: 1')).toBeInTheDocument()
   })
 
+  it('should not commit stale value for newly accessed keys on local rerender (regression in #1176)', async () => {
+    const obj = proxy({ count: 0, anotherValue: 0 })
+
+    const commitFn = vi.fn()
+    const Component = () => {
+      const [showAnotherValue, setShowAnotherValue] = useState(false)
+      const snap = useSnapshot(obj)
+      const value = showAnotherValue ? snap.anotherValue : 'hidden'
+      useLayoutEffect(() => {
+        commitFn(value)
+      }, [value])
+      return (
+        <>
+          <div>count: {snap.count}</div>
+          {showAnotherValue && <div>anotherValue: {snap.anotherValue}</div>}
+          <button onClick={() => setShowAnotherValue(true)}>
+            showAnotherValue
+          </button>
+        </>
+      )
+    }
+
+    render(<Component />)
+
+    expect(screen.getByText('count: 0')).toBeInTheDocument()
+    expect(screen.queryByText('anotherValue: 0')).not.toBeInTheDocument()
+    expect(screen.queryByText('anotherValue: 1')).not.toBeInTheDocument()
+    expect(commitFn).toBeCalledTimes(1)
+    expect(commitFn).toHaveBeenNthCalledWith(1, 'hidden')
+
+    obj.anotherValue += 1
+
+    await act(() => vi.advanceTimersByTimeAsync(0))
+    expect(commitFn).toBeCalledTimes(1)
+
+    fireEvent.click(screen.getByText('showAnotherValue'))
+
+    expect(screen.getByText('anotherValue: 1')).toBeInTheDocument()
+    expect(commitFn.mock.calls.map(([value]) => value)).toEqual(['hidden', 1])
+  })
+
   it('counter with sync option', async () => {
     const obj = proxy({ count: 0 })