diff --git a/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts b/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts index cfac3d27..73973f53 100644 --- a/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts +++ b/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts @@ -120,7 +120,7 @@ export default function useSyncSnapshotHistory() { const cb = ( action: Parameters[0]>[0], ) => { - const isWrite = action.type === 'set'; + const isWrite = action.type === 'set' || action.type === 'async-get'; if ( isWrite && shouldRecordSnapshotHistory && diff --git a/src/stories/Default/Playground/Async.tsx b/src/stories/Default/Playground/Async.tsx index 907abc96..ebd53749 100644 --- a/src/stories/Default/Playground/Async.tsx +++ b/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,7 +24,7 @@ const makeRandomFetchReq = async () => { ); }; // const asyncAtom = atom>(Promise.resolve(null)); -const asyncAtom = atom(Promise.resolve(null)); +const asyncAtom = atom | null>(null); asyncAtom.debugLabel = 'asyncAtom'; const derivedAsyncAtom = atom(async (get) => { @@ -32,26 +32,48 @@ const derivedAsyncAtom = atom(async (get) => { return result?.userId || 'No user'; }); +const ConditionalAsync = () => { + const userId = useAtomValue(derivedAsyncAtom); + return User: {userId}; +}; + 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 ( Async Out-of-the-box Suspense support. Timeout: 8000 ms - User: {userId} + {/* User: {userId} */} + {showResult && } Request status: {!request ? 'Ready' : '✅ Success'} - + + + + + ); }; diff --git a/src/utils/internals/compose-with-devtools.ts b/src/utils/internals/compose-with-devtools.ts index eca0ffe7..54ec83e6 100644 --- a/src/utils/internals/compose-with-devtools.ts +++ b/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 = new Set(); const mountedAtoms = new Set>(); + // 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, number>(); + + const reduceCountOrRemoveRecentlySetAtom = ( + atom: Atom, + 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) => { + 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; });