fix(file-explorer): batch dirCache writes on tree refresh to remove O(N²) freeze (#6321)

Co-authored-by: Orca <help@stably.ai>
Co-authored-by: brennanb2025 <brennankbenson@gmail.com>

src/renderer/src/components/right-sidebar/useFileExplorerTree.test.ts

@@ -0,0 +1,210 @@
+import { describe, expect, it, vi } from 'vitest'
+import type { SetStateAction } from 'react'
+import type { DirEntry } from '../../../../shared/types'
+import type { DirCache } from './file-explorer-types'
+import { createFileExplorerDirLoadTracker } from './file-explorer-dir-load-tracker'
+import { refreshFileExplorerExpandedDirs } from './useFileExplorerTree'
+
+type CacheUpdate = SetStateAction<Record<string, DirCache>>
+
+function entry(name: string, isDirectory = false): DirEntry {
+  return { name, isDirectory, isSymlink: false }
+}
+
+describe('refreshFileExplorerExpandedDirs', () => {
+  it('reloads expanded directories with one loading cache commit and one result cache commit', async () => {
+    let cache: Record<string, DirCache> = {
+      '/repo': {
+        children: [
+          { name: 'old', path: '/repo/old', relativePath: 'old', isDirectory: false, depth: 0 }
+        ],
+        loading: false
+      },
+      '/repo/src': { children: [], loading: false },
+      '/repo/docs': { children: [], loading: false }
+    }
+    const committedCaches: Record<string, DirCache>[] = []
+    const setDirCache = vi.fn((update: CacheUpdate) => {
+      cache = typeof update === 'function' ? update(cache) : update
+      committedCaches.push(cache)
+    })
+    const readDirectory = vi.fn(async (dirPath: string) => {
+      const entriesByPath: Record<string, DirEntry[]> = {
+        '/repo/src': [entry('index.ts')],
+        '/repo/docs': [entry('guide.md')]
+      }
+      return entriesByPath[dirPath] ?? []
+    })
+
+    const refreshed = await refreshFileExplorerExpandedDirs({
+      dirs: [
+        { dirPath: '/repo/src', depth: 0 },
+        { dirPath: '/repo/docs', depth: 0 }
+      ],
+      worktreePath: '/repo',
+      dirLoadTracker: createFileExplorerDirLoadTracker(),
+      setDirCache,
+      readDirectory
+    })
+
+    expect(refreshed).toBe(true)
+    expect(setDirCache).toHaveBeenCalledTimes(2)
+    expect(committedCaches[0]).toMatchObject({
+      '/repo': { loading: false, children: [{ name: 'old' }] },
+      '/repo/src': { loading: true },
+      '/repo/docs': { loading: true }
+    })
+    expect(committedCaches[1]).toMatchObject({
+      '/repo': { loading: false, children: [{ name: 'old' }] },
+      '/repo/src': {
+        loading: false,
+        children: [
+          {
+            name: 'index.ts',
+            path: '/repo/src/index.ts',
+            relativePath: 'src/index.ts',
+            isDirectory: false,
+            depth: 1
+          }
+        ]
+      },
+      '/repo/docs': {
+        loading: false,
+        children: [
+          {
+            name: 'guide.md',
+            path: '/repo/docs/guide.md',
+            relativePath: 'docs/guide.md',
+            isDirectory: false,
+            depth: 1
+          }
+        ]
+      }
+    })
+    expect(readDirectory).toHaveBeenCalledTimes(2)
+  })
+
+  it('drops a superseded directory result so a newer concurrent load is not clobbered', async () => {
+    const tracker = createFileExplorerDirLoadTracker()
+    let cache: Record<string, DirCache> = {
+      '/repo/src': { children: [], loading: false },
+      '/repo/docs': { children: [], loading: false }
+    }
+    const setDirCache = vi.fn((update: CacheUpdate) => {
+      cache = typeof update === 'function' ? update(cache) : update
+    })
+    const newerSrcCache: DirCache = {
+      loading: true,
+      children: [
+        {
+          name: 'fresh.ts',
+          path: '/repo/src/fresh.ts',
+          relativePath: 'src/fresh.ts',
+          isDirectory: false,
+          depth: 1
+        }
+      ]
+    }
+    const readDirectory = vi.fn(async (dirPath: string) => {
+      if (dirPath === '/repo/src') {
+        // Simulate a concurrent newer load (e.g. a watcher-driven refreshDir)
+        // superseding this directory while its refresh read is still in flight:
+        // bump the load token and commit fresher children.
+        tracker.begin('/repo/src')
+        setDirCache((prev) => ({
+          ...prev,
+          '/repo/src': newerSrcCache
+        }))
+        return [entry('stale.ts')]
+      }
+      return [entry('guide.md')]
+    })
+
+    const refreshed = await refreshFileExplorerExpandedDirs({
+      dirs: [
+        { dirPath: '/repo/src', depth: 0 },
+        { dirPath: '/repo/docs', depth: 0 }
+      ],
+      worktreePath: '/repo',
+      dirLoadTracker: tracker,
+      setDirCache,
+      readDirectory
+    })
+
+    // Not every dir was still current, so the refresh reports partial completion.
+    expect(refreshed).toBe(false)
+    // The superseded dir keeps the newer load state; the stale read is
+    // dropped from the batched commit instead of clobbering fresher data.
+    expect(cache['/repo/src']).toEqual(newerSrcCache)
+    // The still-current dir is committed normally.
+    expect(cache['/repo/docs']).toMatchObject({
+      loading: false,
+      children: [{ name: 'guide.md' }]
+    })
+  })
+
+  it('drops a result superseded after its read resolved but before the batch commit', async () => {
+    const tracker = createFileExplorerDirLoadTracker()
+    let cache: Record<string, DirCache> = {
+      '/repo/src': { children: [], loading: false },
+      '/repo/docs': { children: [], loading: false }
+    }
+    const setDirCache = vi.fn((update: CacheUpdate) => {
+      cache = typeof update === 'function' ? update(cache) : update
+    })
+    let releaseDocs!: () => void
+    const docsGate = new Promise<void>((resolve) => {
+      releaseDocs = resolve
+    })
+    const readDirectory = vi.fn(async (dirPath: string) => {
+      if (dirPath === '/repo/src') {
+        return [entry('stale.ts')]
+      }
+      await docsGate
+      return [entry('guide.md')]
+    })
+
+    const refreshPromise = refreshFileExplorerExpandedDirs({
+      dirs: [
+        { dirPath: '/repo/src', depth: 0 },
+        { dirPath: '/repo/docs', depth: 0 }
+      ],
+      worktreePath: '/repo',
+      dirLoadTracker: tracker,
+      setDirCache,
+      readDirectory
+    })
+
+    // Let /repo/src resolve (and pass its resolve-time token check) while
+    // /repo/docs is still in flight — the batch commit is gated on docs.
+    await new Promise((resolve) => setTimeout(resolve, 0))
+
+    // A newer load (e.g. a watcher-driven refreshDir) supersedes /repo/src in
+    // the window between its resolved read and the final batched commit.
+    tracker.begin('/repo/src')
+    const newerSrcCache: DirCache = {
+      loading: false,
+      children: [
+        {
+          name: 'fresh.ts',
+          path: '/repo/src/fresh.ts',
+          relativePath: 'src/fresh.ts',
+          isDirectory: false,
+          depth: 1
+        }
+      ]
+    }
+    setDirCache((prev) => ({ ...prev, '/repo/src': newerSrcCache }))
+
+    releaseDocs()
+    const refreshed = await refreshPromise
+
+    expect(refreshed).toBe(false)
+    // The stale /repo/src read must not clobber the newer committed cache.
+    expect(cache['/repo/src']).toEqual(newerSrcCache)
+    expect(cache['/repo/docs']).toMatchObject({
+      loading: false,
+      children: [{ name: 'guide.md' }]
+    })
+  })
+})

src/renderer/src/components/right-sidebar/useFileExplorerTree.ts

@@ -1,17 +1,21 @@
-import type React from 'react'
+import type { Dispatch, SetStateAction } from 'react'
 import { useCallback, useRef, useState } from 'react'
 import { joinPath, normalizeRelativePath } from '@/lib/path'
 import { getConnectionId } from '@/lib/connection-context'
+import type { DirEntry } from '../../../../shared/types'
 import type { DirCache, TreeNode } from './file-explorer-types'
 import { splitPathSegments } from './path-tree'
 import { shouldIncludeFileExplorerEntry } from './file-explorer-entries'
 import { readRuntimeDirectory, statRuntimePath } from '@/runtime/runtime-file-client'
-import { createFileExplorerDirLoadTracker } from './file-explorer-dir-load-tracker'
+import {
+  createFileExplorerDirLoadTracker,
+  type FileExplorerDirLoadTracker
+} from './file-explorer-dir-load-tracker'
 import { getRightSidebarWorktreeRuntimeSettings } from './file-explorer-runtime-owner'
 
 type UseFileExplorerTreeResult = {
   dirCache: Record<string, DirCache>
-  setDirCache: React.Dispatch<React.SetStateAction<Record<string, DirCache>>>
+  setDirCache: Dispatch<SetStateAction<Record<string, DirCache>>>
   rootCache: DirCache | undefined
   rootError: string | null
   loadDir: (
@@ -26,6 +30,122 @@ type UseFileExplorerTreeResult = {
   resetAndLoad: () => void
 }
 
+type RefreshFileExplorerTreeDir = {
+  dirPath: string
+  depth: number
+}
+
+type RefreshFileExplorerExpandedDirsParams = {
+  dirs: RefreshFileExplorerTreeDir[]
+  worktreePath: string
+  dirLoadTracker: FileExplorerDirLoadTracker
+  setDirCache: Dispatch<SetStateAction<Record<string, DirCache>>>
+  readDirectory: (dirPath: string) => Promise<DirEntry[]>
+}
+
+function entriesToTreeNodes(
+  entries: DirEntry[],
+  dirPath: string,
+  depth: number,
+  worktreePath: string | null
+): TreeNode[] {
+  return entries.filter(shouldIncludeFileExplorerEntry).map((entry) => {
+    const path = joinPath(dirPath, entry.name)
+    return {
+      name: entry.name,
+      path,
+      relativePath: worktreePath
+        ? normalizeRelativePath(path.slice(worktreePath.length + 1))
+        : entry.name,
+      isDirectory: entry.isDirectory,
+      isSymlink: entry.isSymlink,
+      depth: depth + 1
+    }
+  })
+}
+
+export async function refreshFileExplorerExpandedDirs({
+  dirs,
+  worktreePath,
+  dirLoadTracker,
+  setDirCache,
+  readDirectory
+}: RefreshFileExplorerExpandedDirsParams): Promise<boolean> {
+  if (dirs.length === 0) {
+    return true
+  }
+
+  const uniqueDirs = Array.from(new Map(dirs.map((dir) => [dir.dirPath, dir])).values())
+  const loadTokens = new Map(
+    uniqueDirs.map((dir) => [dir.dirPath, dirLoadTracker.begin(dir.dirPath)])
+  )
+
+  // Why: expanded refresh can touch many directories; commit the loading and
+  // result states in two batched setDirCache writes (rather than per-directory)
+  // so refreshing large worktrees stays O(N) instead of O(N²) cache spreads.
+  setDirCache((prev) => {
+    const next = { ...prev }
+    for (const { dirPath } of uniqueDirs) {
+      next[dirPath] = {
+        children: prev[dirPath]?.children ?? [],
+        loading: true
+      }
+    }
+    return next
+  })
+
+  const results = await Promise.all(
+    uniqueDirs.map(async ({ dirPath, depth }) => {
+      const loadToken = loadTokens.get(dirPath)!
+      try {
+        const entries = await readDirectory(dirPath)
+        if (!dirLoadTracker.isCurrent(loadToken)) {
+          return { current: false as const }
+        }
+        return {
+          current: true as const,
+          dirPath,
+          cache: {
+            children: entriesToTreeNodes(entries, dirPath, depth, worktreePath),
+            loading: false
+          }
+        }
+      } catch {
+        if (!dirLoadTracker.isCurrent(loadToken)) {
+          return { current: false as const }
+        }
+        return {
+          current: true as const,
+          dirPath,
+          cache: { children: [], loading: false }
+        }
+      }
+    })
+  )
+
+  // Why: the batch commits only after the slowest read, so a dir can be
+  // superseded (watcher refreshDir, worktree reset) after its own read
+  // resolved. Re-check tokens at commit time so the batched write never
+  // clobbers a newer load — preserving the old per-dir commit ordering.
+  const currentResults = results.filter(
+    (result): result is Extract<typeof result, { current: true }> =>
+      result.current && dirLoadTracker.isCurrent(loadTokens.get(result.dirPath)!)
+  )
+  if (currentResults.length === 0) {
+    return false
+  }
+
+  setDirCache((prev) => {
+    const next = { ...prev }
+    for (const result of currentResults) {
+      next[result.dirPath] = result.cache
+    }
+    return next
+  })
+
+  return currentResults.length === uniqueDirs.length
+}
+
 export function useFileExplorerTree(
   worktreePath: string | null,
   expanded: Set<string>,
@@ -76,18 +196,7 @@ export function useFileExplorerTree(
         if (depth === -1) {
           setRootError(null)
         }
-        const children: TreeNode[] = entries
-          .filter(shouldIncludeFileExplorerEntry)
-          .map((entry) => ({
-            name: entry.name,
-            path: joinPath(dirPath, entry.name),
-            relativePath: worktreePath
-              ? normalizeRelativePath(joinPath(dirPath, entry.name).slice(worktreePath.length + 1))
-              : entry.name,
-            isDirectory: entry.isDirectory,
-            isSymlink: entry.isSymlink,
-            depth: depth + 1
-          }))
+        const children = entriesToTreeNodes(entries, dirPath, depth, worktreePath)
         setDirCache((prev) => ({ ...prev, [dirPath]: { children, loading: false } }))
         return true
       } catch (error) {
@@ -156,13 +265,33 @@ export function useFileExplorerTree(
     if (!rootLoadCompleted || !dirLoadTrackerRef.current.isSessionCurrent(refreshSession)) {
       return
     }
-    await Promise.all(
-      Array.from(expanded).map(async (dirPath) => {
-        const depth = splitPathSegments(dirPath.slice(worktreePath.length + 1)).length - 1
-        await loadDir(dirPath, depth, { force: true })
-      })
-    )
-  }, [expanded, loadDir, worktreePath])
+    // Why: root (worktreePath) was just force-loaded above; exclude it here so
+    // refreshFileExplorerExpandedDirs doesn't queue a duplicate read of root.
+    const expandedDirs = Array.from(expanded)
+      .filter((dirPath) => dirPath !== worktreePath)
+      .map((dirPath) => ({
+        dirPath,
+        depth: splitPathSegments(dirPath.slice(worktreePath.length + 1)).length - 1
+      }))
+    await refreshFileExplorerExpandedDirs({
+      dirs: expandedDirs,
+      worktreePath,
+      dirLoadTracker: dirLoadTrackerRef.current,
+      setDirCache,
+      readDirectory: async (dirPath) => {
+        const connectionId = getConnectionId(activeWorktreeId ?? null) ?? undefined
+        return readRuntimeDirectory(
+          {
+            settings: getRightSidebarWorktreeRuntimeSettings(activeWorktreeId),
+            worktreeId: activeWorktreeId,
+            worktreePath,
+            connectionId
+          },
+          dirPath
+        )
+      }
+    })
+  }, [activeWorktreeId, expanded, loadDir, worktreePath])
 
   const refreshDir = useCallback(
     async (dirPath: string) => {