Better optimistic adding/removing

src/components/PostControls/BookmarkButton.tsx

@@ -25,7 +25,7 @@ export const BookmarkButton = memo(function BookmarkButton({
   const {mutateAsync: bookmark} = useBookmarkMutation()
   const cleanError = useCleanError()
 
-  const {uri, cid, viewer} = post
+  const {viewer} = post
   const isBookmarked = !!viewer?.bookmarked
 
   const undoLabel = _(
@@ -39,8 +39,7 @@ export const BookmarkButton = memo(function BookmarkButton({
     try {
       await bookmark({
         action: 'create',
-        uri,
-        cid,
+        post,
       })
 
       toast.show(
@@ -73,7 +72,7 @@ export const BookmarkButton = memo(function BookmarkButton({
     try {
       await bookmark({
         action: 'delete',
-        uri,
+        uri: post.uri,
       })
 
       toast.show(

src/screens/Bookmarks/index.tsx

@@ -16,7 +16,6 @@ import {
   type NativeStackScreenProps,
 } from '#/lib/routes/types'
 import {isIOS} from '#/platform/detection'
-import {POST_TOMBSTONE, usePostShadow} from '#/state/cache/post-shadow'
 import {useBookmarkMutation} from '#/state/queries/bookmarks/useBookmarkMutation'
 import {useBookmarksQuery} from '#/state/queries/bookmarks/useBookmarksQuery'
 import {useSetMinimalShellMode} from '#/state/shell'
@@ -190,19 +189,6 @@ function BookmarksInner() {
   )
 }
 
-function BookmarkPost({
-  hideTopBorder,
-  post: originalPost,
-}: {
-  hideTopBorder: boolean
-  post: $Typed<AppBskyFeedDefs.PostView>
-}) {
-  const post = usePostShadow(originalPost)
-  if (post === POST_TOMBSTONE) return null
-  if (!post.viewer?.bookmarked) return null
-  return <Post post={post} hideTopBorder={hideTopBorder} />
-}
-
 function BookmarkNotFound({
   hideTopBorder,
   post,
@@ -214,14 +200,12 @@ function BookmarkNotFound({
   const {_} = useLingui()
   const {mutateAsync: bookmark} = useBookmarkMutation()
   const cleanError = useCleanError()
-  const [removed, setRemoved] = useState(false)
 
   const remove = async () => {
     try {
-      setRemoved(true)
       await bookmark({action: 'delete', uri: post.uri})
       toast.show(_(msg`Removed from saved posts`), {
-        type: 'error',
+        type: 'info',
       })
     } catch (e: any) {
       const {raw, clean} = cleanError(e)
@@ -231,8 +215,6 @@ function BookmarkNotFound({
     }
   }
 
-  if (removed) return null
-
   return (
     <View
       style={[
@@ -286,9 +268,7 @@ function renderItem({item, index}: {item: ListItem; index: number}) {
       return <EmptyState />
     }
     case 'bookmark': {
-      return (
-        <BookmarkPost post={item.bookmark.item} hideTopBorder={index === 0} />
-      )
+      return <Post post={item.bookmark.item} hideTopBorder={index === 0} />
     }
     case 'bookmarkNotFound': {
       return (

src/state/queries/bookmarks/useBookmarkMutation.ts

@@ -1,14 +1,25 @@
+import {type AppBskyFeedDefs} from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
 import {isNetworkError} from '#/lib/strings/errors'
 import {logger} from '#/logger'
 import {updatePostShadow} from '#/state/cache/post-shadow'
-import {truncateAndInvalidate as invalidateBookmarksQuery} from '#/state/queries/bookmarks/useBookmarksQuery'
+import {
+  optimisticallyDeleteBookmark,
+  optimisticallySaveBookmark,
+} from '#/state/queries/bookmarks/useBookmarksQuery'
 import {useAgent} from '#/state/session'
 
 type MutationArgs =
-  | {action: 'create'; uri: string; cid: string}
-  | {action: 'delete'; uri: string}
+  | {action: 'create'; post: AppBskyFeedDefs.PostView}
+  | {
+      action: 'delete'
+      /**
+       * For deletions, we only need to URI. Plus, in some cases we only know the
+       * URI, such as when a post was deleted by the author.
+       */
+      uri: string
+    }
 
 export function useBookmarkMutation() {
   const qc = useQueryClient()
@@ -17,10 +28,10 @@ export function useBookmarkMutation() {
   return useMutation({
     async mutationFn(args: MutationArgs) {
       if (args.action === 'create') {
-        updatePostShadow(qc, args.uri, {bookmarked: true})
+        updatePostShadow(qc, args.post.uri, {bookmarked: true})
         await agent.app.bsky.bookmark.createBookmark({
-          uri: args.uri,
-          cid: args.cid,
+          uri: args.post.uri,
+          cid: args.post.cid,
         })
       } else if (args.action === 'delete') {
         updatePostShadow(qc, args.uri, {bookmarked: false})
@@ -29,12 +40,16 @@ export function useBookmarkMutation() {
         })
       }
     },
-    onSuccess() {
-      invalidateBookmarksQuery(qc)
+    onSuccess(_, args) {
+      if (args.action === 'create') {
+        optimisticallySaveBookmark(qc, args.post)
+      } else if (args.action === 'delete') {
+        optimisticallyDeleteBookmark(qc, {uri: args.uri})
+      }
     },
     onError(e, args) {
       if (args.action === 'create') {
-        updatePostShadow(qc, args.uri, {bookmarked: false})
+        updatePostShadow(qc, args.post.uri, {bookmarked: false})
       } else if (args.action === 'delete') {
         updatePostShadow(qc, args.uri, {bookmarked: true})
       }

src/state/queries/bookmarks/useBookmarksQuery.ts

@@ -1,4 +1,8 @@
-import {type AppBskyBookmarkGetBookmarks} from '@atproto/api'
+import {
+  type $Typed,
+  type AppBskyBookmarkGetBookmarks,
+  type AppBskyFeedDefs,
+} from '@atproto/api'
 import {
   type InfiniteData,
   type QueryClient,
@@ -25,6 +29,7 @@ export function useBookmarksQuery() {
     async queryFn({pageParam}) {
       const res = await agent.app.bsky.bookmark.getBookmarks({
         cursor: pageParam,
+        limit: 5,
       })
       return res.data
     },
@@ -48,3 +53,63 @@ export async function truncateAndInvalidate(qc: QueryClient) {
   )
   return qc.invalidateQueries({queryKey: [bookmarksQueryKeyRoot]})
 }
+
+export async function optimisticallySaveBookmark(
+  qc: QueryClient,
+  post: AppBskyFeedDefs.PostView,
+) {
+  qc.setQueriesData<InfiniteData<AppBskyBookmarkGetBookmarks.OutputSchema>>(
+    {
+      queryKey: [bookmarksQueryKeyRoot],
+    },
+    data => {
+      if (!data) return data
+      return {
+        ...data,
+        pages: data.pages.map((page, index) => {
+          if (index === 0) {
+            post.$type = 'app.bsky.feed.defs#postView'
+            return {
+              ...page,
+              bookmarks: [
+                {
+                  createdAt: new Date().toISOString(),
+                  subject: {
+                    uri: post.uri,
+                    cid: post.cid,
+                  },
+                  item: post as $Typed<AppBskyFeedDefs.PostView>,
+                },
+                ...page.bookmarks,
+              ],
+            }
+          }
+          return page
+        }),
+      }
+    },
+  )
+}
+
+export async function optimisticallyDeleteBookmark(
+  qc: QueryClient,
+  {uri}: {uri: string},
+) {
+  qc.setQueriesData<InfiniteData<AppBskyBookmarkGetBookmarks.OutputSchema>>(
+    {
+      queryKey: [bookmarksQueryKeyRoot],
+    },
+    data => {
+      if (!data) return data
+      return {
+        ...data,
+        pages: data.pages.map(page => {
+          return {
+            ...page,
+            bookmarks: page.bookmarks.filter(b => b.subject.uri !== uri),
+          }
+        }),
+      }
+    },
+  )
+}