Fix navigation

app/gui/src/dashboard/components/dashboard/column/PathColumn.tsx

@@ -11,7 +11,6 @@ import { useSetExpandedDirectoryIds, useSetSelectedKeys } from '#/providers/Driv
 import type { DirectoryId } from '#/services/Backend'
 import { isDirectoryId } from '#/services/Backend'
 import { Fragment, useTransition } from 'react'
-import { flushSync } from 'react-dom'
 import invariant from 'tiny-invariant'
 import type { AssetColumnProps } from '../column'
 
@@ -68,16 +67,8 @@ export default function PathColumn(props: AssetColumnProps) {
     if (targetDirectoryNode == null && rootDirectoryInThePath.categoryId != null) {
       // we reassign the variable only to make ts happy here.
       const categoryId = rootDirectoryInThePath.categoryId
-      // We need to set the category first, because setting a category
-      // resets the list of expanded folders and selected keys
-      // Drive store resets it's state when a category change
-      // calling `setCategory` inside the `flushSync` guaranties that the all side effects will
-      // be executed and state is fresh and clean.
-      // This comes with a cost though - it might be slow as it executes everything synchronously.
-      flushSync(() => {
-        setCategory(categoryId)
-      })
 
+      setCategory(categoryId)
       setExpandedDirectoryIds(pathToDirectory.map(({ id }) => id).concat(targetDirectory))
     }
 

app/gui/src/dashboard/hooks/storeHooks.ts

@@ -7,6 +7,7 @@ import type { DispatchWithoutAction, Reducer, RefObject } from 'react'
 import { useEffect, useReducer, useRef } from 'react'
 import { type StoreApi } from 'zustand'
 import { useStoreWithEqualityFn } from 'zustand/traditional'
+
 import { objectEquality, refEquality, shallowEquality } from '../utilities/equalities'
 
 /**

app/gui/src/dashboard/layouts/Drive/Categories/categoriesHooks.tsx

@@ -335,19 +335,20 @@ const CategoriesContext = createContext<CategoriesContextValue | null>(null)
  */
 export interface CategoriesProviderProps {
   readonly children: ReactNode | ((contextValue: CategoriesContextValue) => ReactNode)
+  readonly onCategoryChange?: (previousCategory: Category | null, newCategory: Category) => void
 }
 
 /**
  * Provider for the categories.
  */
 export function CategoriesProvider(props: CategoriesProviderProps): React.JSX.Element {
-  const { children } = props
+  const { children, onCategoryChange = () => {} } = props
 
   const { cloudCategories, localCategories, findCategoryById } = useCategories()
   const localBackend = useLocalBackend()
   const { isOffline } = useOffline()
 
-  const [categoryId, setCategoryId, resetCategoryId] = useSearchParamsState<CategoryId>(
+  const [categoryId, privateSetCategoryId, resetCategoryId] = useSearchParamsState<CategoryId>(
     'driveCategory',
     () => {
       if (isOffline && localBackend != null) {
@@ -361,6 +362,14 @@ export function CategoriesProvider(props: CategoriesProviderProps): React.JSX.El
     (value): value is CategoryId => findCategoryById(value as CategoryId) != null,
   )
 
+  const setCategoryId = useEventCallback((nextCategoryId: CategoryId) => {
+    const previousCategory = findCategoryById(categoryId)
+    privateSetCategoryId(nextCategoryId)
+    // This is safe, because we know that the result will have the correct type.
+    // eslint-disable-next-line no-restricted-syntax
+    onCategoryChange(previousCategory, findCategoryById(nextCategoryId) as Category)
+  })
+
   const category = findCategoryById(categoryId)
 
   // This is safe, because category is always set

app/gui/src/dashboard/pages/dashboard/Dashboard.tsx

@@ -69,19 +69,19 @@ export interface DashboardProps {
 /** The component that contains the entire UI. */
 export default function Dashboard(props: DashboardProps) {
   return (
-    <CategoriesProvider>
-      {/* Ideally this would be in `Drive.tsx`, but it currently must be all the way out here
-       * due to modals being in `TheModal`. */}
-      {({ category }) => (
-        <DriveProvider currentCategoryId={category.id}>
+    /* Ideally this would be in `Drive.tsx`, but it currently must be all the way out here
+     * due to modals being in `TheModal`. */
+    <DriveProvider>
+      {({ resetAssetTableState }) => (
+        <CategoriesProvider onCategoryChange={resetAssetTableState}>
           <EventListProvider>
             <ProjectsProvider>
               <DashboardInner {...props} />
             </ProjectsProvider>
           </EventListProvider>
-        </DriveProvider>
+        </CategoriesProvider>
       )}
-    </CategoriesProvider>
+    </DriveProvider>
   )
 }
 

app/gui/src/dashboard/providers/DriveProvider.tsx

@@ -5,7 +5,7 @@ import * as zustand from '#/utilities/zustand'
 import invariant from 'tiny-invariant'
 
 import { useEventCallback } from '#/hooks/eventCallbackHooks'
-import type { Category, CategoryId } from '#/layouts/CategorySwitcher/Category'
+import type { Category } from '#/layouts/CategorySwitcher/Category'
 import type AssetTreeNode from '#/utilities/AssetTreeNode'
 import type { PasteData } from '#/utilities/pasteData'
 import { EMPTY_SET } from '#/utilities/set'
@@ -59,8 +59,13 @@ export type ProjectsContextType = zustand.StoreApi<DriveStore>
 const DriveContext = React.createContext<ProjectsContextType | null>(null)
 
 /** Props for a {@link DriveProvider}. */
-export type ProjectsProviderProps = Readonly<React.PropsWithChildren> & {
-  readonly currentCategoryId: CategoryId
+export interface ProjectsProviderProps {
+  readonly children:
+    | React.ReactNode
+    | ((context: {
+        readonly store: ProjectsContextType
+        readonly resetAssetTableState: () => void
+      }) => React.ReactNode)
 }
 
 // ========================
@@ -72,7 +77,7 @@ export type ProjectsProviderProps = Readonly<React.PropsWithChildren> & {
  * containing the current element is focused.
  */
 export default function DriveProvider(props: ProjectsProviderProps) {
-  const { children, currentCategoryId } = props
+  const { children } = props
 
   const [store] = React.useState(() =>
     zustand.createStore<DriveStore>((set, get) => ({
@@ -135,14 +140,13 @@ export default function DriveProvider(props: ProjectsProviderProps) {
     })),
   )
 
-  // Reset the asset table state when the category changes
-  // TODO: This is a bit of a hack, but should stay for now. Eventually we should just recreate the store
-  // when the category changes.
-  React.useEffect(() => {
-    store.getState().resetAssetTableState()
-  }, [currentCategoryId, store])
+  const resetAssetTableState = zustand.useStore(store, (state) => state.resetAssetTableState)
 
-  return <DriveContext.Provider value={store}>{children}</DriveContext.Provider>
+  return (
+    <DriveContext.Provider value={store}>
+      {typeof children === 'function' ? children({ store, resetAssetTableState }) : children}
+    </DriveContext.Provider>
+  )
 }
 
 /** The drive store. */
@@ -223,15 +227,8 @@ export function useExpandedDirectoryIds() {
 /** A function to set the expanded directoyIds in the Asset Table. */
 export function useSetExpandedDirectoryIds() {
   const store = useDriveStore()
-  const privateSetExpandedDirectoryIds = zustand.useStore(
-    store,
-    (state) => state.setExpandedDirectoryIds,
-    { unsafeEnableTransition: true },
-  )
-  return useEventCallback((expandedDirectoryIds: readonly DirectoryId[]) => {
-    React.startTransition(() => {
-      privateSetExpandedDirectoryIds(expandedDirectoryIds)
-    })
+  return zustand.useStore(store, (state) => state.setExpandedDirectoryIds, {
+    unsafeEnableTransition: true,
   })
 }
 

app/gui/src/dashboard/providers/__test__/DriveProvider.test.tsx

@@ -12,6 +12,7 @@ function renderDriveProviderHook<Result, Props>(
 ): RenderHookResult<Result, Props> {
   let currentCategoryId: CategoryId = 'cloud'
   let setCategoryId: (categoryId: CategoryId) => void
+  let doResetAssetTableState: () => void
 
   return renderHook(
     (props) => {
@@ -23,9 +24,19 @@ function renderDriveProviderHook<Result, Props>(
         // eslint-disable-next-line react-hooks/rules-of-hooks
         const [category, setCategory] = useState(() => currentCategoryId)
         currentCategoryId = category
-        setCategoryId = setCategory
+        setCategoryId = (nextCategoryId) => {
+          setCategory(nextCategoryId)
+          doResetAssetTableState()
+        }
 
-        return <DriveProvider currentCategoryId={currentCategoryId}>{children}</DriveProvider>
+        return (
+          <DriveProvider>
+            {({ resetAssetTableState }) => {
+              doResetAssetTableState = resetAssetTableState
+              return children
+            }}
+          </DriveProvider>
+        )
       },
       ...options,
     },
@@ -36,11 +47,21 @@ describe('<DriveProvider />', () => {
   it('Should reset expanded directory ids when category changes', () => {
     const driveAPI = renderDriveProviderHook((setCategoryId: (categoryId: string) => void) => {
       const store = useDriveStore()
-      return useStore(store, ({ setExpandedDirectoryIds, expandedDirectoryIds }) => ({
-        expandedDirectoryIds,
-        setExpandedDirectoryIds,
-        setCategoryId,
-      }))
+      return useStore(
+        store,
+        ({
+          setExpandedDirectoryIds,
+          expandedDirectoryIds,
+          selectedKeys,
+          visuallySelectedKeys,
+        }) => ({
+          expandedDirectoryIds,
+          setExpandedDirectoryIds,
+          setCategoryId,
+          selectedKeys,
+          visuallySelectedKeys,
+        }),
+      )
     })
 
     act(() => {
@@ -56,5 +77,7 @@ describe('<DriveProvider />', () => {
     })
 
     expect(driveAPI.result.current.expandedDirectoryIds).toEqual([])
+    expect(driveAPI.result.current.selectedKeys).toEqual(new Set())
+    expect(driveAPI.result.current.visuallySelectedKeys).toEqual(null)
   })
 })

app/gui/src/dashboard/utilities/equalities.ts

@@ -3,6 +3,7 @@
  *
  * This file contains functions for checking equality between values.
  */
+import { shallow } from 'zustand/shallow'
 
 /**
  * Strict equality check.
@@ -22,30 +23,5 @@ export function objectEquality<T>(a: T, b: T) {
  * Shallow equality check.
  */
 export function shallowEquality<T>(a: T, b: T) {
-  if (Object.is(a, b)) {
-    return true
-  }
-
-  if (typeof a !== 'object' || a == null || typeof b !== 'object' || b == null) {
-    return false
-  }
-
-  const keysA = Object.keys(a)
-
-  if (keysA.length !== Object.keys(b).length) {
-    return false
-  }
-
-  for (let i = 0; i < keysA.length; i++) {
-    const key = keysA[i]
-
-    if (key != null) {
-      // @ts-expect-error Typescript doesn't know that key is in a and b, but it doesn't matter here
-      if (!Object.prototype.hasOwnProperty.call(b, key) || !Object.is(a[key], b[key])) {
-        return false
-      }
-    }
-  }
-
-  return true
+  return shallow(a, b)
 }