Fix issue with stale unread message counts (#10953)

src/state/queries/messages/update-all-read.ts

@@ -1,8 +1,10 @@
+import {type ChatBskyConvoGetUnreadCounts} from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
 import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {logger} from '#/logger'
 import {useAgent} from '#/state/session'
+import {RQKEY_PARTIAL as UNREAD_COUNTS_PARTIAL_KEY} from './get-unread-counts'
 import {
   type ConvoRequestListQueryData,
   markAllRead as markAllRequestsRead,
@@ -94,10 +96,36 @@ export function useUpdateAllRead(
           markAllRequestsRead,
         )
       }
+      // zero out the badge count query that actually drives the unread badge,
+      // since it's a separate server query that the list caches don't feed
+      const prevUnreadCountsQueries =
+        queryClient.getQueriesData<ChatBskyConvoGetUnreadCounts.OutputSchema>({
+          queryKey: UNREAD_COUNTS_PARTIAL_KEY,
+        })
+      queryClient.setQueriesData<ChatBskyConvoGetUnreadCounts.OutputSchema>(
+        {queryKey: UNREAD_COUNTS_PARTIAL_KEY},
+        old => {
+          if (!old) return old
+          return {
+            ...old,
+            ...(status === 'accepted'
+              ? {unreadAcceptedConvos: 0}
+              : {unreadRequestConvos: 0}),
+          }
+        },
+      )
       onMutate?.()
-      return {prevConvoListQueries, prevRequestsQueries}
+      return {
+        prevConvoListQueries,
+        prevRequestsQueries,
+        prevUnreadCountsQueries,
+      }
     },
     onSuccess: () => {
+      // the optimistic badge zeroing can drift from the server, so invalidate
+      // the count query to let it self-correct on next access rather than
+      // waiting for a log event
+      void queryClient.invalidateQueries({queryKey: UNREAD_COUNTS_PARTIAL_KEY})
       void queryClient.invalidateQueries({
         queryKey: CONVO_LIST_PARTIAL_KEY(status),
       })
@@ -121,6 +149,11 @@ export function useUpdateAllRead(
           queryClient.setQueryData(queryKey, prevData)
         }
       }
+      if (context?.prevUnreadCountsQueries) {
+        for (const [queryKey, prevData] of context.prevUnreadCountsQueries) {
+          queryClient.setQueryData(queryKey, prevData)
+        }
+      }
       void queryClient.invalidateQueries({queryKey: [CONVO_LIST_ROOT_KEY]})
       if (status === 'request') {
         void queryClient.invalidateQueries({queryKey: [REQUESTS_RQKEY_ROOT]})