feat: improve support for promises during time travel

src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts

@@ -120,7 +120,7 @@ export default function useSyncSnapshotHistory() {
     const cb = (
       action: Parameters<Parameters<typeof userStore.subscribeStore>[0]>[0],
     ) => {
-      const isWrite = action.type === 'set';
+      const isWrite = action.type === 'set' || action.type === 'async-get';
       if (
         isWrite &&
         shouldRecordSnapshotHistory &&

src/stories/Default/Playground/Async.tsx

@@ -1,5 +1,5 @@
-import React from 'react';
-import { Box, Button, Text, Title } from '@mantine/core';
+import React, { useState } from 'react';
+import { Box, Button, Flex, Text, Title } from '@mantine/core';
 import { useAtom, useAtomValue } from 'jotai/react';
 import { atom } from 'jotai/vanilla';
 
@@ -24,34 +24,56 @@ const makeRandomFetchReq = async () => {
   );
 };
 // const asyncAtom = atom<Promise<any>>(Promise.resolve(null));
-const asyncAtom = atom(Promise.resolve<any>(null));
+const asyncAtom = atom<Promise<any> | null>(null);
 asyncAtom.debugLabel = 'asyncAtom';
 
 const derivedAsyncAtom = atom(async (get) => {
   const result = await get(asyncAtom);
   return result?.userId || 'No user';
 });
 
+const ConditionalAsync = () => {
+  const userId = useAtomValue(derivedAsyncAtom);
+  return <Text>User: {userId}</Text>;
+};
+
 export const Async = () => {
   const [request, setRequest] = useAtom(asyncAtom);
-  const userId = useAtomValue(derivedAsyncAtom);
-  // const setRequest = useSetAtom(asyncAtom, demoStoreOptions);
+  const [showResult, setShowResult] = useState(false);
 
   const handleFetchClick = async () => {
     setRequest(makeRandomFetchReq); // Will suspend until request resolves
   };
 
+  const handleShowResultClick = () => {
+    setShowResult((v) => !v);
+  };
+
   return (
     <Box>
       <Title size="h5">Async</Title>
       <Text mb={10} c="dark.2">
         Out-of-the-box Suspense support. <i>Timeout: 8000 ms</i>
       </Text>
-      User: {userId}
+      {/* User: {userId} */}
+      {showResult && <ConditionalAsync />}
       <Text>Request status: {!request ? 'Ready' : '✅ Success'} </Text>
-      <Button onClick={handleFetchClick} size="xs" tt="uppercase" mt={5}>
-        Fetch
-      </Button>
+      <Flex>
+        <Button onClick={handleFetchClick} size="xs" tt="uppercase" mt={5}>
+          Fetch
+        </Button>
+
+        <Button
+          onClick={handleShowResultClick}
+          size="xs"
+          tt="uppercase"
+          mt={5}
+          ml={5}
+          color="green"
+        >
+          Toggle result
+        </Button>
+      </Flex>
     </Box>
   );
 };

src/utils/internals/compose-with-devtools.ts

@@ -10,7 +10,7 @@ import {
 } from '../../types';
 
 type DevSubscribeStoreListener = (action: {
-  type: 'get' | 'set' | 'sub' | 'unsub' | 'restore';
+  type: 'get' | 'async-get' | 'set' | 'sub' | 'unsub' | 'restore';
 }) => void;
 
 type DevToolsStoreMethods = {
@@ -56,28 +56,76 @@ const __composeV2StoreWithDevTools = (
   const storeListeners: Set<DevSubscribeStoreListener> = new Set();
   const mountedAtoms = new Set<Atom<unknown>>();
 
+  // Map to keep track of how many times an atom was set recently
+  // We mostly use this to re-collect the values for history tracking for async atoms
+  // Async atoms call `get` once they finish fetching the value, so we verify if the atom was set recently and then trigger the history tracking
+  // I hope there is a better way to do this
+  const recentlySetAtomsMap = new WeakMap<Atom<unknown>, number>();
+
+  const reduceCountOrRemoveRecentlySetAtom = (
+    atom: Atom<unknown>,
+    onFound?: () => void,
+  ) => {
+    const foundCount = recentlySetAtomsMap.get(atom);
+    if (typeof foundCount === 'number') {
+      if (foundCount > 1) {
+        recentlySetAtomsMap.set(atom, foundCount - 1);
+      } else {
+        recentlySetAtomsMap.delete(atom);
+      }
+      onFound?.();
+    }
+  };
+
+  const increaseCountRecentlySetAtom = (atom: Atom<unknown>) => {
+    const foundCount = recentlySetAtomsMap.get(atom);
+    recentlySetAtomsMap.set(atom, (foundCount || 0) + 1);
+  };
+
   store.dev4_override_method('sub', (...args) => {
     mountedAtoms.add(args[0]);
     const unsub = sub(...args);
     storeListeners.forEach((l) => l({ type: 'sub' }));
     return () => {
-      // Check if the atom has no listeners, if so, remove it from the mounted list
-      if (store.dev4_get_internal_weak_map().get(args[0])?.m?.l.size === 0) {
-        mountedAtoms.delete(args[0]);
-      }
       unsub();
+
+      // FIXME is there a better way to check if its mounted?
+      // Check if the atom has no listeners, if so, remove it from the mounted list in the next tick
+      Promise.resolve().then(() => {
+        const atomState = store.dev4_get_internal_weak_map().get(args[0]);
+        if (typeof atomState?.m === 'undefined') {
+          mountedAtoms.delete(args[0]);
+        }
+      });
+
+      // We remove the atom from the recently set map if it was set recently when it is unsubscribed
+      reduceCountOrRemoveRecentlySetAtom(args[0]);
+
       storeListeners.forEach((l) => l({ type: 'unsub' }));
     };
   });
 
   store.dev4_override_method('get', (...args) => {
     const value = get(...args);
+
+    reduceCountOrRemoveRecentlySetAtom(args[0], () => {
+      if (value instanceof Promise) {
+        value.then(() => {
+          // We wait for a tick to ensure that if there are any derived atoms, we wait for them to be flushed out as well
+          Promise.resolve().then(() => {
+            storeListeners.forEach((l) => l({ type: 'async-get' }));
+          });
+        });
+      }
+    });
+
     storeListeners.forEach((l) => l({ type: 'get' }));
     return value;
   });
 
   store.dev4_override_method('set', (...args) => {
     const value = set(...args);
+    increaseCountRecentlySetAtom(args[0]);
     storeListeners.forEach((l) => l({ type: 'set' }));
     return value;
   });