Fix issue with stale unread message counts (#10953)

src/state/queries/messages/conversation.ts

@@ -2,6 +2,7 @@ import {
   type ChatBskyActorDefs,
   type ChatBskyConvoDefs,
   type ChatBskyConvoGetConvo,
+  type ChatBskyConvoGetUnreadCounts,
 } from '@atproto/api'
 import {
   type QueryClient,
@@ -14,6 +15,11 @@ import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {STALE} from '#/state/queries'
 import {useOnMarkAsRead} from '#/state/queries/messages/list-conversations'
 import {useAgent} from '#/state/session'
+import {
+  RQKEY_PARTIAL as UNREAD_COUNTS_PARTIAL_KEY,
+  UNREAD_ACCEPTED_CAP,
+  UNREAD_REQUEST_CAP,
+} from './get-unread-counts'
 import {
   type ConvoListQueryData,
   getConvoFromQueryData,
@@ -74,11 +80,81 @@ export function useMarkAsReadMutation() {
     },
     onMutate({convoId}) {
       if (!convoId) throw new Error('No convoId provided')
+
+      // snapshot the list caches before the optimistic update so onError can
+      // restore the convo rows alongside the badge count
+      const prevListQueries = queryClient.getQueriesData<ConvoListQueryData>({
+        queryKey: [LIST_CONVOS_KEY],
+      })
+
+      // find the convo so we know which badge counter (if any) to decrement.
+      // keep scanning past a stale unreadCount === 0 cache so another cache
+      // holding the true unread state still drives the decrement
+      let unreadStatus: ChatBskyConvoDefs.ConvoView['status'] | undefined
+      for (const [, data] of prevListQueries) {
+        if (!data) continue
+        const convo = getConvoFromQueryData(convoId, data)
+        if (convo?.unreadCount) {
+          unreadStatus = convo.status
+          break
+        }
+      }
+
       optimisticUpdate(convoId)
+
+      // the badge count query is a separate server query that the list caches
+      // don't feed, so decrement it here to keep the badge in sync
+      const prevUnreadCountsQueries =
+        queryClient.getQueriesData<ChatBskyConvoGetUnreadCounts.OutputSchema>({
+          queryKey: UNREAD_COUNTS_PARTIAL_KEY,
+        })
+      if (unreadStatus) {
+        queryClient.setQueriesData<ChatBskyConvoGetUnreadCounts.OutputSchema>(
+          {queryKey: UNREAD_COUNTS_PARTIAL_KEY},
+          old => {
+            if (!old) return old
+            return {
+              ...old,
+              ...(unreadStatus === 'request'
+                ? {
+                    unreadRequestConvos:
+                      old.unreadRequestConvos >= UNREAD_REQUEST_CAP
+                        ? old.unreadRequestConvos
+                        : Math.max(0, old.unreadRequestConvos - 1),
+                  }
+                : {
+                    unreadAcceptedConvos:
+                      old.unreadAcceptedConvos >= UNREAD_ACCEPTED_CAP
+                        ? old.unreadAcceptedConvos
+                        : Math.max(0, old.unreadAcceptedConvos - 1),
+                  }),
+            }
+          },
+        )
+      }
+      return {prevListQueries, prevUnreadCountsQueries}
+    },
+    onError(_, __, context) {
+      if (context?.prevListQueries) {
+        for (const [queryKey, prevData] of context.prevListQueries) {
+          queryClient.setQueryData(queryKey, prevData)
+        }
+      }
+      if (context?.prevUnreadCountsQueries) {
+        for (const [queryKey, prevData] of context.prevUnreadCountsQueries) {
+          queryClient.setQueryData(queryKey, prevData)
+        }
+      }
     },
     onSuccess(_, {convoId}) {
       if (!convoId) return
 
+      // the optimistic badge arithmetic can drift from the server (e.g. a convo
+      // whose status differs between caches, or a sentinel-capped count). invalidate
+      // so the 15s-stale count query self-corrects on next access rather than
+      // waiting for a log event
+      void queryClient.invalidateQueries({queryKey: UNREAD_COUNTS_PARTIAL_KEY})
+
       queryClient.setQueriesData(
         {queryKey: [LIST_CONVOS_KEY]},
         (old?: ConvoListQueryData) => {