fix: directory table optimistic update

fix optimistic update not working for directory table and trash table

apps/drive-web/src/components/with-atom.tsx

@@ -1,15 +1,64 @@
-import { type PrimitiveAtom, useAtom } from "jotai"
+import {
+	type Atom,
+	type ExtractAtomArgs,
+	type ExtractAtomResult,
+	type ExtractAtomValue,
+	type PrimitiveAtom,
+	type SetStateAction,
+	useAtom,
+	type WritableAtom,
+} from "jotai"
+import type * as React from "react"
 
-export function WithAtom<Value>({
+type SetAtom<Args extends unknown[], Result> = (...args: Args) => Result
+
+export function WithAtom<Value, Args extends unknown[], Result>(props: {
+	atom: WritableAtom<Value, Args, Result>
+	children: (
+		value: Awaited<Value>,
+		setAtom: SetAtom<Args, Result>,
+	) => React.ReactNode
+}): React.ReactNode
+export function WithAtom<Value>(props: {
+	atom: PrimitiveAtom<Value>
+	children: (
+		value: Awaited<Value>,
+		setAtom: SetAtom<[SetStateAction<Value>], void>,
+	) => React.ReactNode
+}): React.ReactNode
+export function WithAtom<Value>(props: {
+	atom: Atom<Value>
+	children: (value: Awaited<Value>, setAtom: never) => React.ReactNode
+}): React.ReactNode
+export function WithAtom<
+	AtomType extends WritableAtom<unknown, never[], unknown>,
+>(props: {
+	atom: AtomType
+	children: (
+		value: Awaited<ExtractAtomValue<AtomType>>,
+		setAtom: SetAtom<
+			ExtractAtomArgs<AtomType>,
+			ExtractAtomResult<AtomType>
+		>,
+	) => React.ReactNode
+}): React.ReactNode
+export function WithAtom<AtomType extends Atom<unknown>>(props: {
+	atom: AtomType
+	children: (
+		value: Awaited<ExtractAtomValue<AtomType>>,
+		setAtom: never,
+	) => React.ReactNode
+}): React.ReactNode
+export function WithAtom<Value, Args extends unknown[], Result>({
 	atom,
 	children,
 }: {
-	atom: PrimitiveAtom<Value>
+	atom: Atom<Value> | WritableAtom<Value, Args, Result>
 	children: (
-		value: Value,
-		setValue: (value: Value | ((current: Value) => Value)) => void,
+		value: Awaited<Value>,
+		setAtom: SetAtom<Args, Result> | never,
 	) => React.ReactNode
 }) {
-	const [value, setValue] = useAtom(atom)
-	return children(value, setValue)
+	const [value, setAtom] = useAtom(atom as WritableAtom<Value, Args, Result>)
+	return children(value, setAtom)
 }