feat: impl dir content table sorting

apps/drive-web/package.json

@@ -7,7 +7,9 @@
 		"dev": "vite",
 		"build": "vite build",
 		"preview": "vite preview",
-		"format": "biome format --write"
+		"format": "bunx biome format --write",
+		"lint": "bunx biome lint",
+		"typecheck": "bunx tsc -p tsconfig.json --noEmit"
 	},
 	"dependencies": {
 		"@convex-dev/better-auth": "^0.8.9",

apps/drive-web/src/directories/directory-page/defaults.ts

@@ -0,0 +1,9 @@
+import {
+	DIRECTORY_CONTENT_ORDER_BY,
+	DIRECTORY_CONTENT_ORDER_DIRECTION,
+} from "@/vfs/api"
+
+export const DEFAULT_DIRECTORY_CONTENT_ORDER_BY =
+	DIRECTORY_CONTENT_ORDER_BY.name
+export const DEFAULT_DIRECTORY_CONTENT_ORDER_DIRECTION =
+	DIRECTORY_CONTENT_ORDER_DIRECTION.asc

apps/drive-web/src/directories/directory-page/directory-content-table.tsx

@@ -1,4 +1,6 @@
-import { useInfiniteQuery } from "@tanstack/react-query"
+// TODO: make table sorting work (right now not working probably because different params share same query key)
+
+import { useInfiniteQuery, useMutation } from "@tanstack/react-query"
 import { Link, useNavigate } from "@tanstack/react-router"
 import {
 	type ColumnDef,
@@ -10,9 +12,11 @@ import {
 	useReactTable,
 } from "@tanstack/react-table"
 import { useVirtualizer, type VirtualItem } from "@tanstack/react-virtual"
-import { type PrimitiveAtom, useSetAtom, useStore } from "jotai"
+import { type PrimitiveAtom, useAtomValue, useSetAtom, useStore } from "jotai"
+import { ArrowDownIcon, ArrowUpIcon } from "lucide-react"
 import type React from "react"
-import { useEffect, useMemo, useRef } from "react"
+import { createContext, useContext, useEffect, useMemo, useRef } from "react"
+import { toast } from "sonner"
 import { DirectoryIcon } from "@/components/icons/directory-icon"
 import { TextFileIcon } from "@/components/icons/text-file-icon"
 import { Checkbox } from "@/components/ui/checkbox"
@@ -31,25 +35,56 @@ import {
 	keyboardModifierAtom,
 } from "@/lib/keyboard"
 import { cn } from "@/lib/utils"
-import type { DirectoryContentQuery } from "@/vfs/api"
+import {
+	DIRECTORY_CONTENT_ORDER_BY,
+	DIRECTORY_CONTENT_ORDER_DIRECTION,
+	type DirectoryContentOrderBy,
+	type DirectoryContentOrderDirection,
+	type DirectoryContentQuery,
+	type MoveDirectoryItemsResult,
+	moveDirectoryItemsMutationAtom,
+} from "@/vfs/api"
 import type { DirectoryInfo, DirectoryItem, FileInfo } from "@/vfs/vfs"
+import {
+	DEFAULT_DIRECTORY_CONTENT_ORDER_BY,
+	DEFAULT_DIRECTORY_CONTENT_ORDER_DIRECTION,
+} from "./defaults"
 import { DirectoryContentTableSkeleton } from "./directory-content-table-skeleton"
 
 type DirectoryContentTableItemIdFilter = Set<string>
 
+export type DirectoryContentTableSortChangeCallback = (
+	orderBy: DirectoryContentOrderBy,
+	direction: DirectoryContentOrderDirection,
+) => void
+
 export type DirectoryContentTableProps = {
+	query: DirectoryContentQuery
 	directoryUrlFn: (directory: DirectoryInfo) => string
 	fileDragInfoAtom: PrimitiveAtom<FileDragInfo | null>
+	loadingComponent?: React.ReactNode
+	debugBanner?: React.ReactNode
+
 	onContextMenu: (
 		row: Row<DirectoryItem>,
 		table: TableType<DirectoryItem>,
 	) => void
 	onOpenFile: (file: FileInfo) => void
-	query: DirectoryContentQuery
+	onSortChange: DirectoryContentTableSortChangeCallback
-	loadingComponent?: React.ReactNode
-	debugBanner?: React.ReactNode
 }
 
+export type DirectoryContentTableContext = {
+	orderBy: DirectoryContentOrderBy
+	direction: DirectoryContentOrderDirection
+	onSortChange: DirectoryContentTableSortChangeCallback
+}
+const DirectoryContentTableContext =
+	createContext<DirectoryContentTableContext>({
+		orderBy: DEFAULT_DIRECTORY_CONTENT_ORDER_BY,
+		direction: DEFAULT_DIRECTORY_CONTENT_ORDER_DIRECTION,
+		onSortChange: () => {},
+	})
+
 function formatFileSize(bytes: number): string {
 	if (bytes === 0) return "0 B"
 
@@ -92,7 +127,7 @@ function useTableColumns(
 				size: 24,
 			},
 			{
-				header: "Name",
+				header: () => <NameHeaderCell />,
 				accessorKey: "doc.name",
 				cell: ({ row }) => {
 					switch (row.original.kind) {
@@ -144,15 +179,15 @@ function useTableColumns(
 	)
 }
 
-// Shared table component that accepts query options as props
 export function DirectoryContentTable({
 	directoryUrlFn,
-	onContextMenu,
 	fileDragInfoAtom,
-	onOpenFile,
 	query,
 	loadingComponent,
 	debugBanner,
+	onOpenFile,
+	onContextMenu,
+	onSortChange,
 }: DirectoryContentTableProps) {
 	const {
 		data: directoryContent,
@@ -164,6 +199,33 @@ export function DirectoryContentTable({
 
 	const store = useStore()
 	const navigate = useNavigate()
+	const moveDirectoryItemsMutation = useAtomValue(
+		moveDirectoryItemsMutationAtom,
+	)
+
+	const { mutate: moveDroppedItems } = useMutation({
+		...moveDirectoryItemsMutation,
+		onSuccess: (data: MoveDirectoryItemsResult) => {
+			const conflictCount = data.errors.length
+			if (conflictCount > 0) {
+				toast.warning(
+					`${data.moved.length} items moved${conflictCount > 0 ? `, ${conflictCount} conflicts` : ""}`,
+				)
+			} else {
+				toast.success(`${data.moved.length} items moved!`)
+			}
+		},
+	})
+
+	const handleFileDrop = (
+		items: import("@/vfs/vfs").DirectoryItem[],
+		targetDirectory: import("@/vfs/vfs").DirectoryInfo | string,
+	) => {
+		moveDroppedItems({
+			targetDirectory,
+			items,
+		})
+	}
 
 	const table = useReactTable({
 		data: useMemo(
@@ -292,52 +354,61 @@ export function DirectoryContentTable({
 				onDoubleClick={() => {
 					handleRowDoubleClick(row)
 				}}
+				onFileDrop={handleFileDrop}
 			/>
 		)
 	}
 
 	return (
-		<div className="h-full flex flex-col">
+		<DirectoryContentTableContext
-			{debugBanner}
+			value={{
-			<TableContainer
+				orderBy: query.initialPageParam.orderBy,
-				className={debugBanner ? "flex-1" : "h-full"}
+				direction: query.initialPageParam.direction,
-				ref={containerRef}
+				onSortChange,
-			>
+			}}
-				<Table className="h-full min-h-0">
+		>
-					<TableHeader>
+			<div className="h-full flex flex-col">
-						{table.getHeaderGroups().map((headerGroup) => (
+				{debugBanner}
-							<TableRow
+				<TableContainer
-								className="px-4 border-b-0!"
+					className={debugBanner ? "flex-1" : "h-full"}
-								key={headerGroup.id}
+					ref={containerRef}
-							>
+				>
-								{headerGroup.headers.map((header) => (
+					<Table className="h-full min-h-0">
-									<TableHead
+						<TableHeader>
-										className="first:pl-4 last:pr-4 sticky top-0 bg-background z-1 inset-shadow-[0_-1px_0_0_var(--border)]"
+							{table.getHeaderGroups().map((headerGroup) => (
-										key={header.id}
+								<TableRow
-										style={{ width: header.getSize() }}
+									className="px-4 border-b-0!"
-									>
+									key={headerGroup.id}
-										{header.isPlaceholder
+								>
-											? null
+									{headerGroup.headers.map((header) => (
-											: flexRender(
+										<TableHead
-													header.column.columnDef
+											className="first:pl-4 last:pr-4 sticky top-0 bg-background z-1 inset-shadow-[0_-1px_0_0_var(--border)]"
-														.header,
+											key={header.id}
-													header.getContext(),
+											style={{ width: header.getSize() }}
-												)}
+										>
-									</TableHead>
+											{header.isPlaceholder
-								))}
+												? null
-							</TableRow>
+												: flexRender(
-						))}
+														header.column.columnDef
-					</TableHeader>
+															.header,
-					<TableBody className="overflow-auto">
+														header.getContext(),
-						{rows.length > 0 ? (
+													)}
-							virtualItems.map(renderRow)
+										</TableHead>
-						) : (
+									))}
-							<NoResultsRow />
+								</TableRow>
-						)}
+							))}
-					</TableBody>
+						</TableHeader>
-				</Table>
+						<TableBody className="overflow-auto">
-			</TableContainer>
+							{rows.length > 0 ? (
-		</div>
+								virtualItems.map(renderRow)
+							) : (
+								<NoResultsRow />
+							)}
+						</TableBody>
+					</Table>
+				</TableContainer>
+			</div>
+		</DirectoryContentTableContext>
 	)
 }
 
@@ -358,6 +429,7 @@ function FileItemRow({
 	onContextMenu,
 	onDoubleClick,
 	fileDragInfoAtom,
+	onFileDrop,
 	...rowProps
 }: React.ComponentProps<typeof TableRow> & {
 	table: TableType<DirectoryItem>
@@ -366,6 +438,10 @@ function FileItemRow({
 	onContextMenu: (e: React.MouseEvent) => void
 	onDoubleClick: () => void
 	fileDragInfoAtom: PrimitiveAtom<FileDragInfo | null>
+	onFileDrop: (
+		items: import("@/vfs/vfs").DirectoryItem[],
+		targetDirectory: import("@/vfs/vfs").DirectoryInfo | string,
+	) => void
 }) {
 	const ref = useRef<HTMLTableRowElement>(null)
 	const setFileDragInfo = useSetAtom(fileDragInfoAtom)
@@ -374,6 +450,7 @@ function FileItemRow({
 		enabled: row.original.kind === "directory",
 		destDir: row.original.kind === "directory" ? row.original : undefined,
 		dragInfoAtom: fileDragInfoAtom,
+		onDrop: onFileDrop,
 	})
 
 	const handleDragStart = (_e: React.DragEvent) => {
@@ -433,6 +510,54 @@ function FileItemRow({
 	)
 }
 
+function NameHeaderCell() {
+	const { orderBy, direction, onSortChange } = useContext(
+		DirectoryContentTableContext,
+	)
+
+	let arrow: React.ReactNode = null
+	if (orderBy === DIRECTORY_CONTENT_ORDER_BY.name) {
+		switch (direction) {
+			case DIRECTORY_CONTENT_ORDER_DIRECTION.asc:
+				arrow = <ArrowUpIcon className="size-4" />
+				break
+			case DIRECTORY_CONTENT_ORDER_DIRECTION.desc:
+				arrow = <ArrowDownIcon className="size-4" />
+				break
+		}
+	}
+
+	let directionLabel: string
+	switch (direction) {
+		case DIRECTORY_CONTENT_ORDER_DIRECTION.asc:
+			directionLabel = "ascending"
+			break
+		case DIRECTORY_CONTENT_ORDER_DIRECTION.desc:
+			directionLabel = "descending"
+			break
+	}
+
+	return (
+		<button
+			type="button"
+			className="hover:underline cursor-pointer flex items-center gap-2 w-full"
+			onClick={() => {
+				onSortChange(
+					DIRECTORY_CONTENT_ORDER_BY.name,
+					direction === DIRECTORY_CONTENT_ORDER_DIRECTION.asc
+						? DIRECTORY_CONTENT_ORDER_DIRECTION.desc
+						: DIRECTORY_CONTENT_ORDER_DIRECTION.asc,
+				)
+			}}
+		>
+			<span className="sr-only">Sort by</span>
+			<span>Name</span>
+			<span className="sr-only">in {directionLabel} order</span>
+			{arrow}
+		</button>
+	)
+}
+
 function DirectoryNameCell({
 	directory,
 	directoryUrlFn,

apps/drive-web/src/files/upload-file-dialog.tsx

@@ -29,7 +29,7 @@ import {
 	TooltipTrigger,
 } from "@/components/ui/tooltip"
 import { formatError } from "@/lib/error"
-import { directoryContentQueryAtom } from "@/vfs/api"
+import { directoryContentQueryKey } from "@/vfs/api"
 import type { DirectoryInfoWithPath } from "@/vfs/vfs"
 import { currentAccountAtom } from "../account/account"
 import {
@@ -135,11 +135,16 @@ function useUploadFilesAtom({
 						toast.success("All files uploaded successfully")
 					}
 
-					client.invalidateQueries(
+					// Invalidate all queries for the target directory (with any params)
-						store.get(
+					const account = store.get(currentAccountAtom)
-							directoryContentQueryAtom(targetDirectory.id),
+					if (account) {
-						),
+						client.invalidateQueries({
-					)
+							queryKey: directoryContentQueryKey(
+								account.id,
+								targetDirectory.id,
+							),
+						})
+					}
 				},
 				onError: (error) => {
 					console.error(error)

apps/drive-web/src/files/use-file-drop.ts

@@ -1,12 +1,6 @@
-import { useMutation } from "@tanstack/react-query"
 import type { PrimitiveAtom } from "jotai"
-import { useAtomValue, useSetAtom, useStore } from "jotai"
+import { useSetAtom, useStore } from "jotai"
 import { useState } from "react"
-import { toast } from "sonner"
-import {
-	type MoveDirectoryItemsResult,
-	moveDirectoryItemsMutationAtom,
-} from "@/vfs/api"
 import type { DirectoryInfo, DirectoryItem } from "@/vfs/vfs"
 
 export interface FileDragInfo {
@@ -18,6 +12,10 @@ export interface UseFileDropOptions {
 	destDir?: DirectoryInfo | string
 	dragInfoAtom: PrimitiveAtom<FileDragInfo | null>
 	enabled?: boolean
+	onDrop?: (
+		items: DirectoryItem[],
+		targetDirectory: DirectoryInfo | string,
+	) => void
 }
 
 export interface UseFileDropReturn {
@@ -33,43 +31,22 @@ export function useFileDrop({
 	destDir,
 	dragInfoAtom,
 	enabled,
+	onDrop,
 }: UseFileDropOptions): UseFileDropReturn {
 	const [isDraggedOver, setIsDraggedOver] = useState(false)
 	const setDragInfo = useSetAtom(dragInfoAtom)
 	const store = useStore()
 
-	const moveDirectoryItemsMutation = useAtomValue(
-		moveDirectoryItemsMutationAtom,
-	)
-
-	const { mutate: moveDroppedItems } = useMutation({
-		...moveDirectoryItemsMutation,
-		onSuccess: (data: MoveDirectoryItemsResult, vars, result, ctx) => {
-			moveDirectoryItemsMutation.onSuccess?.(data, vars, result, ctx)
-			const conflictCount = data.errors.length
-			if (conflictCount > 0) {
-				toast.warning(
-					`${data.moved.length} items moved${conflictCount > 0 ? `, ${conflictCount} conflicts` : ""}`,
-				)
-			} else {
-				toast.success(`${data.moved.length} items moved!`)
-			}
-		},
-	})
-
 	const dirId = typeof destDir === "string" ? destDir : destDir?.id
 
 	const handleDrop = (_e: React.DragEvent) => {
-		if (!enabled || !destDir || !dirId) return
+		if (!enabled || !destDir || !dirId || !onDrop) return
 
 		const dragInfo = store.get(dragInfoAtom)
 		if (dragInfo) {
 			const items = dragInfo.items.filter((item) => item.id !== dirId)
 			if (items.length > 0) {
-				moveDroppedItems({
+				onDrop(items, destDir)
-					targetDirectory: destDir,
-					items,
-				})
 			}
 		}
 		setIsDraggedOver(false)

apps/drive-web/src/routes/_authenticated/_sidebar-layout/directories.$directoryId.tsx

@@ -13,6 +13,7 @@ import {
 } from "lucide-react"
 import { lazy, Suspense, useCallback, useContext } from "react"
 import { toast } from "sonner"
+import { currentAccountAtom } from "@/account/account"
 import { DirectoryIcon } from "@/components/icons/directory-icon"
 import { TextFileIcon } from "@/components/icons/text-file-icon"
 import { Button } from "@/components/ui/button"
@@ -31,6 +32,10 @@ import {
 import { WithAtom } from "@/components/with-atom"
 import { backgroundTaskProgressAtom } from "@/dashboard/state"
 import { DirectoryPageContext } from "@/directories/directory-page/context"
+import {
+	DEFAULT_DIRECTORY_CONTENT_ORDER_BY,
+	DEFAULT_DIRECTORY_CONTENT_ORDER_DIRECTION,
+} from "@/directories/directory-page/defaults"
 import { DirectoryContentTable } from "@/directories/directory-page/directory-content-table"
 import { DirectoryPageSkeleton } from "@/directories/directory-page/directory-page-skeleton"
 import { NewDirectoryDialog } from "@/directories/directory-page/new-directory-dialog"
@@ -41,13 +46,22 @@ import { cutItemsAtom, inProgressFileUploadCountAtom } from "@/files/store"
 import { UploadFileDialog } from "@/files/upload-file-dialog"
 import type { FileDragInfo } from "@/files/use-file-drop"
 import { formatError } from "@/lib/error"
+import type {
+	DirectoryContentOrderBy,
+	DirectoryContentOrderDirection,
+} from "@/vfs/api"
 import {
 	DIRECTORY_CONTENT_ORDER_BY,
 	DIRECTORY_CONTENT_ORDER_DIRECTION,
 	directoryContentQueryAtom,
+	directoryContentQueryKey,
 	directoryInfoQueryAtom,
 	moveToTrashMutationAtom,
 } from "@/vfs/api"
+import {
+	optimisticallyRemoveDirectoryItems,
+	rollbackDirectoryContentOptimisticUpdate,
+} from "@/vfs/optimistic"
 import type {
 	DirectoryInfo,
 	DirectoryInfoWithPath,
@@ -70,10 +84,10 @@ const MockDirectoryContentTable = import.meta.env.DEV
 const DirectoryContentPageParams = type({
 	orderBy: type
 		.valueOf(DIRECTORY_CONTENT_ORDER_BY)
-		.default(DIRECTORY_CONTENT_ORDER_BY.name),
+		.default(DEFAULT_DIRECTORY_CONTENT_ORDER_BY),
 	direction: type
 		.valueOf(DIRECTORY_CONTENT_ORDER_DIRECTION)
-		.default(DIRECTORY_CONTENT_ORDER_DIRECTION.asc),
+		.default(DEFAULT_DIRECTORY_CONTENT_ORDER_DIRECTION),
 })
 
 export const Route = createFileRoute(
@@ -218,6 +232,7 @@ function _DirectoryContentTable() {
 	// Always call the hook - in production the atom always returns false
 	const useMock = useAtomValue(mockTableAtom)
 	const { directory } = useContext(DirectoryPageContext)
+	const navigate = Route.useNavigate()
 
 	const search = Route.useSearch()
 	const query = useAtomValue(
@@ -257,6 +272,21 @@ function _DirectoryContentTable() {
 		[setContextMenuTargetItems],
 	)
 
+	const applySorting = useCallback(
+		(
+			orderBy: DirectoryContentOrderBy,
+			direction: DirectoryContentOrderDirection,
+		) => {
+			navigate({
+				search: {
+					orderBy,
+					direction,
+				},
+			})
+		},
+		[navigate],
+	)
+
 	// In production, MockDirectoryContentTable is null, so this always renders real table
 	if (import.meta.env.DEV && useMock && MockDirectoryContentTable) {
 		return (
@@ -278,6 +308,7 @@ function _DirectoryContentTable() {
 			fileDragInfoAtom={fileDragInfoAtom}
 			onContextMenu={handleContextMenuRequest}
 			onOpenFile={onTableOpenFile}
+			onSortChange={applySorting}
 		/>
 	)
 }
@@ -320,29 +351,36 @@ function DirectoryContentContextMenu({
 	const [target, setTarget] = useAtom(contextMenuTargetItemsAtom)
 	const setBackgroundTaskProgress = useSetAtom(backgroundTaskProgressAtom)
 	const setCutItems = useSetAtom(cutItemsAtom)
+	const account = useAtomValue(currentAccountAtom)
+	const { directory } = useContext(DirectoryPageContext)
+	const search = Route.useSearch()
 
 	const moveToTrashMutation = useAtomValue(moveToTrashMutationAtom)
 	const { mutate: moveToTrash } = useMutation({
 		...moveToTrashMutation,
-		onMutate: (vars, ctx) => {
+		onMutate: (items, { client }) => {
 			setBackgroundTaskProgress({
 				label: "Moving items to trash…",
 			})
-			return (
+			if (!account) {
-				moveToTrashMutation.onMutate?.(vars, ctx) ?? {
+				return null
-					prevDirContentMap: new Map(),
+			}
-				}
+			return optimisticallyRemoveDirectoryItems(client, {
-			)
+				queryKey: directoryContentQueryKey(account.id, directory.id, {
+					orderBy: search.orderBy,
+					direction: search.direction,
+				}),
+				items,
+			})
 		},
-		onSuccess: (data, vars, result, ctx) => {
+		onSuccess: (data) => {
-			moveToTrashMutation.onSuccess?.(data, vars, result, ctx)
 			setBackgroundTaskProgress(null)
 			toast.success(`Moved ${data.length} items to trash`)
 		},
-		onError: (err, vars, mutateResult, context) => {
+		onError: (err, _vars, mutateResult, { client }) => {
-			moveToTrashMutation.onError?.(err, vars, mutateResult, context)
 			toast.error(formatError(err))
 			setBackgroundTaskProgress(null)
+			rollbackDirectoryContentOptimisticUpdate(client, mutateResult)
 		},
 	})
 

apps/drive-web/src/vfs/api.ts

@@ -18,10 +18,11 @@ import {
 	FileInfo,
 } from "./vfs"
 
-const DirectoryContentResponse = type({
+export const DirectoryContentResponse = type({
 	items: DirectoryContent,
 	"nextCursor?": "string",
 })
+export type DirectoryContentResponseType = typeof DirectoryContentResponse.infer
 
 /**
  * This atom derives the file url for a given file.
@@ -81,7 +82,7 @@ export const DIRECTORY_CONTENT_ORDER_DIRECTION = {
 	asc: "asc",
 	desc: "desc",
 } as const
-type DirectoryContentOrderDirection =
+export type DirectoryContentOrderDirection =
 	(typeof DIRECTORY_CONTENT_ORDER_DIRECTION)[keyof typeof DIRECTORY_CONTENT_ORDER_DIRECTION]
 
 type DirectoryContentQueryParams = {
@@ -98,22 +99,27 @@ type DirectoryContentPageParam = {
 	cursor: string
 }
 
-const directoryContentQueryKey = (
+export const directoryContentQueryKey = (
 	accountId: string | undefined,
 	directoryId: string,
-): readonly (string | undefined)[] => [
+	params?: {
+		orderBy?: DirectoryContentOrderBy
+		direction?: DirectoryContentOrderDirection
+	},
+): readonly unknown[] => [
 	"accounts",
 	accountId,
 	"directories",
 	directoryId,
 	"content",
+	...(params ? [{ orderBy: params.orderBy, dir: params.direction }] : []),
 ]
 export type DirectoryContentQuery = ReturnType<
 	typeof infiniteQueryOptions<
 		typeof DirectoryContentResponse.infer,
 		Error,
 		InfiniteData<typeof DirectoryContentResponse.infer>,
-		readonly (string | undefined)[],
+		readonly unknown[],
 		DirectoryContentPageParam
 	>
 >
@@ -122,7 +128,10 @@ export const directoryContentQueryAtom = atomFamily(
 		atom((get) => {
 			const account = get(currentAccountAtom)
 			return infiniteQueryOptions({
-				queryKey: directoryContentQueryKey(account?.id, directoryId),
+				queryKey: directoryContentQueryKey(account?.id, directoryId, {
+					orderBy,
+					direction,
+				}),
 				initialPageParam: {
 					orderBy,
 					direction,
@@ -221,53 +230,6 @@ export const moveDirectoryItemsMutationAtom = atom((get) =>
 			)
 			return result
 		},
-		onMutate: ({ items }, { client }) => {
-			const account = get(currentAccountAtom)
-			if (!account) {
-				return null
-			}
-
-			const movedItems = new Map<string, Set<string>>()
-
-			for (const item of items) {
-				if (item.parentId) {
-					const s = movedItems.get(item.parentId)
-					if (!s) {
-						movedItems.set(item.parentId, new Set([item.id]))
-					} else {
-						s.add(item.id)
-					}
-				}
-			}
-
-			const prevDirContentMap = new Map<
-				string,
-				InfiniteData<typeof DirectoryContentResponse.infer> | undefined
-			>()
-
-			movedItems.forEach((s, parentId) => {
-				const key = directoryContentQueryKey(account.id, parentId)
-				const prevDirContent =
-					client.getQueryData<
-						InfiniteData<typeof DirectoryContentResponse.infer>
-					>(key)
-				client.setQueryData<
-					InfiniteData<typeof DirectoryContentResponse.infer>
-				>(key, (prev) => {
-					if (!prev) return prev
-					return {
-						...prev,
-						pages: prev.pages.map((page) => ({
-							...page,
-							items: page.items.filter((it) => !s.has(it.id)),
-						})),
-					}
-				})
-				prevDirContentMap.set(parentId, prevDirContent)
-			})
-
-			return { prevDirContentMap }
-		},
 		onSuccess: (_data, { targetDirectory, items }, _result, { client }) => {
 			const account = get(currentAccountAtom)
 			if (!account) return
@@ -276,6 +238,7 @@ export const moveDirectoryItemsMutationAtom = atom((get) =>
 				typeof targetDirectory === "string"
 					? targetDirectory
 					: targetDirectory.id
+			// Invalidate using base key (without params) to invalidate all queries for these directories
 			client.invalidateQueries({
 				queryKey: directoryContentQueryKey(account.id, dirId),
 			})
@@ -290,21 +253,6 @@ export const moveDirectoryItemsMutationAtom = atom((get) =>
 				}
 			}
 		},
-		onError: (_error, _vars, context, { client }) => {
-			if (context) {
-				const account = get(currentAccountAtom)
-				if (account) {
-					context.prevDirContentMap.forEach(
-						(prevDirContent, parentId) => {
-							client.setQueryData(
-								directoryContentQueryKey(account.id, parentId),
-								prevDirContent,
-							)
-						},
-					)
-				}
-			}
-		},
 	}),
 )
 
@@ -368,53 +316,10 @@ export const moveToTrashMutationAtom = atom((get) =>
 
 			return [...deletedFiles, ...deletedDirectories]
 		},
-		onMutate: (items, { client }) => {
-			const account = get(currentAccountAtom)
-			if (!account) {
-				return null
-			}
-
-			const trashedItems = new Map<string, Set<string>>()
-			for (const item of items) {
-				if (item.parentId) {
-					const s = trashedItems.get(item.parentId)
-					if (!s) {
-						trashedItems.set(item.parentId, new Set([item.id]))
-					} else {
-						s.add(item.id)
-					}
-				}
-			}
-
-			const prevDirContentMap = new Map<
-				string,
-				InfiniteData<typeof DirectoryContentResponse.infer> | undefined
-			>()
-			trashedItems.forEach((s, parentId) => {
-				const key = directoryContentQueryKey(account.id, parentId)
-				const prevDirContent =
-					client.getQueryData<
-						InfiniteData<typeof DirectoryContentResponse.infer>
-					>(key)
-				client.setQueryData<
-					InfiniteData<typeof DirectoryContentResponse.infer>
-				>(key, (prev) => {
-					if (!prev) return prev
-					return {
-						...prev,
-						pages: prev.pages.map((page) => ({
-							...page,
-							items: page.items.filter((it) => !s.has(it.id)),
-						})),
-					}
-				})
-				prevDirContentMap.set(parentId, prevDirContent)
-			})
-			return { prevDirContentMap }
-		},
 		onSuccess: (_data, items, _result, { client }) => {
 			const account = get(currentAccountAtom)
 			if (account) {
+				// Invalidate using base key (without params) to invalidate all queries for these directories
 				for (const item of items) {
 					if (item.parentId) {
 						client.invalidateQueries({
@@ -427,21 +332,6 @@ export const moveToTrashMutationAtom = atom((get) =>
 				}
 			}
 		},
-		onError: (_error, _items, context, { client }) => {
-			if (context) {
-				const account = get(currentAccountAtom)
-				if (account) {
-					context.prevDirContentMap.forEach(
-						(prevDirContent, parentId) => {
-							client.setQueryData(
-								directoryContentQueryKey(account.id, parentId),
-								prevDirContent,
-							)
-						},
-					)
-				}
-			}
-		},
 	}),
 )
 

apps/drive-web/src/vfs/optimistic.ts

@@ -0,0 +1,51 @@
+import type { InfiniteData, QueryClient } from "@tanstack/react-query"
+import type { DirectoryContentResponseType } from "@/vfs/api"
+import type { DirectoryItem } from "@/vfs/vfs"
+
+export type DirectoryContentOptimisticUpdate = {
+	queryKey: readonly unknown[]
+	prevDirContent: InfiniteData<DirectoryContentResponseType> | undefined
+}
+
+type OptimisticRemovalParams = {
+	queryKey: readonly unknown[]
+	items: DirectoryItem[]
+}
+
+export function optimisticallyRemoveDirectoryItems(
+	client: QueryClient,
+	{ queryKey, items }: OptimisticRemovalParams,
+): DirectoryContentOptimisticUpdate {
+	const prevDirContent =
+		client.getQueryData<InfiniteData<DirectoryContentResponseType>>(
+			queryKey,
+		)
+	const removedItemIds = new Set(items.map((item) => item.id))
+
+	client.setQueryData<InfiniteData<DirectoryContentResponseType>>(
+		queryKey,
+		(prev) => {
+			if (!prev) return prev
+			return {
+				...prev,
+				pages: prev.pages.map((page) => ({
+					...page,
+					items: page.items.filter(
+						(item) => !removedItemIds.has(item.id),
+					),
+				})),
+			}
+		},
+	)
+
+	return { queryKey, prevDirContent }
+}
+
+export function rollbackDirectoryContentOptimisticUpdate(
+	client: QueryClient,
+	update?: DirectoryContentOptimisticUpdate | null,
+) {
+	if (update) {
+		client.setQueryData(update.queryKey, update.prevDirContent)
+	}
+}