[Chat] Fix optimistic updates by using partial keys (bluesky-social#10838)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: mozzius

src/state/queries/messages/accept-conversation.ts

@@ -1,6 +1,6 @@
 import {
   type ChatBskyConvoAcceptConvo,
-  type ChatBskyConvoListConvos,
+  type ChatBskyConvoDefs,
 } from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
@@ -13,7 +13,11 @@ import {
   RQKEY_ROOT as REQUESTS_RQKEY_ROOT,
 } from './list-conversation-requests'
 import {
-  RQKEY as CONVO_LIST_KEY,
+  type ConvoListQueryData,
+  convoListQueryPredicate,
+  getConvoFromQueryData,
+  optimisticDelete,
+  RQKEY_PARTIAL as CONVO_LIST_PARTIAL_KEY,
   RQKEY_ROOT as CONVO_LIST_ROOT_KEY,
 } from './list-conversations'
 
@@ -42,65 +46,49 @@ export function useAcceptConversation(
       return data
     },
     onMutate: () => {
-      let prevAcceptedPages: ChatBskyConvoListConvos.OutputSchema[] = []
-      let prevInboxPages: ChatBskyConvoListConvos.OutputSchema[] = []
-      let convoBeingAccepted:
-        | ChatBskyConvoListConvos.OutputSchema['convos'][number]
-        | undefined
-      queryClient.setQueryData(
-        CONVO_LIST_KEY('request'),
-        (old?: {
-          pageParams: Array<string | undefined>
-          pages: Array<ChatBskyConvoListConvos.OutputSchema>
-        }) => {
-          if (!old) return old
-          prevInboxPages = old.pages
-          return {
-            ...old,
-            pages: old.pages.map(page => {
-              const found = page.convos.find(convo => convo.id === convoId)
-              if (found) {
-                convoBeingAccepted = found
-                return {
-                  ...page,
-                  convos: page.convos.filter(convo => convo.id !== convoId),
-                }
-              }
-              return page
-            }),
-          }
-        },
+      // snapshot every convo-list cache up front so onError can restore them
+      // all by their exact keys
+      const prevConvoListQueries =
+        queryClient.getQueriesData<ConvoListQueryData>({
+          queryKey: [CONVO_LIST_ROOT_KEY],
+        })
+      let convoBeingAccepted: ChatBskyConvoDefs.ConvoView | null = null
+      for (const [_key, data] of queryClient.getQueriesData<ConvoListQueryData>(
+        {queryKey: CONVO_LIST_PARTIAL_KEY('request')},
+      )) {
+        if (!data) continue
+        convoBeingAccepted = getConvoFromQueryData(convoId, data)
+        if (convoBeingAccepted) break
+      }
+      queryClient.setQueriesData(
+        {queryKey: CONVO_LIST_PARTIAL_KEY('request')},
+        (old?: ConvoListQueryData) => optimisticDelete(convoId, old),
       )
-      queryClient.setQueryData(
-        CONVO_LIST_KEY('accepted'),
-        (old?: {
-          pageParams: Array<string | undefined>
-          pages: Array<ChatBskyConvoListConvos.OutputSchema>
-        }) => {
-          if (!old) return old
-          prevAcceptedPages = old.pages
-          if (convoBeingAccepted) {
+      if (convoBeingAccepted) {
+        const acceptedConvo: ChatBskyConvoDefs.ConvoView = {
+          ...convoBeingAccepted,
+          status: 'accepted',
+        }
+        queryClient.setQueriesData(
+          {
+            queryKey: CONVO_LIST_PARTIAL_KEY('accepted'),
+            predicate: convoListQueryPredicate(acceptedConvo),
+          },
+          (old?: ConvoListQueryData) => {
+            if (!old) return old
             return {
               ...old,
-              pages: [
-                {
-                  ...old.pages[0],
-                  convos: [
-                    {
-                      ...convoBeingAccepted,
-                      status: 'accepted',
-                    },
-                    ...old.pages[0].convos,
-                  ],
-                },
-                ...old.pages.slice(1),
-              ],
+              pages: old.pages.map((page, i) => {
+                const convos = page.convos.filter(c => c.id !== convoId)
+                if (i === 0) {
+                  return {...page, convos: [acceptedConvo, ...convos]}
+                }
+                return {...page, convos}
+              }),
             }
-          } else {
-            return old
-          }
-        },
-      )
+          },
+        )
+      }
       const prevRequestsQueries =
         queryClient.getQueriesData<ConvoRequestListQueryData>({
           queryKey: [REQUESTS_RQKEY_ROOT],
@@ -110,41 +98,20 @@ export function useAcceptConversation(
         old => optimisticDeleteRequest(convoId, old),
       )
       onMutate?.()
-      return {prevAcceptedPages, prevInboxPages, prevRequestsQueries}
+      return {prevConvoListQueries, prevRequestsQueries}
     },
     onSuccess: data => {
-      void queryClient.invalidateQueries({queryKey: [CONVO_LIST_KEY]})
+      void queryClient.invalidateQueries({queryKey: [CONVO_LIST_ROOT_KEY]})
       void queryClient.invalidateQueries({queryKey: [REQUESTS_RQKEY_ROOT]})
       onSuccess?.(data)
     },
     onError: (error, _, context) => {
       logger.error(error)
-      queryClient.setQueryData(
-        CONVO_LIST_KEY('accepted'),
-        (old?: {
-          pageParams: Array<string | undefined>
-          pages: Array<ChatBskyConvoListConvos.OutputSchema>
-        }) => {
-          if (!old) return old
-          return {
-            ...old,
-            pages: context?.prevAcceptedPages || old.pages,
-          }
-        },
-      )
-      queryClient.setQueryData(
-        CONVO_LIST_KEY('request'),
-        (old?: {
-          pageParams: Array<string | undefined>
-          pages: Array<ChatBskyConvoListConvos.OutputSchema>
-        }) => {
-          if (!old) return old
-          return {
-            ...old,
-            pages: context?.prevInboxPages || old.pages,
-          }
-        },
-      )
+      if (context?.prevConvoListQueries) {
+        for (const [queryKey, prevData] of context.prevConvoListQueries) {
+          queryClient.setQueryData(queryKey, prevData)
+        }
+      }
       if (context?.prevRequestsQueries) {
         for (const [queryKey, prevData] of context.prevRequestsQueries) {
           queryClient.setQueryData(queryKey, prevData)

src/state/queries/messages/list-conversations.tsx

@@ -8,7 +8,9 @@ import {
 } from '@atproto/api'
 import {
   type InfiniteData,
+  type Query,
   type QueryClient,
+  type QueryKey,
   useInfiniteQuery,
   useQueryClient,
 } from '@tanstack/react-query'
@@ -48,7 +50,60 @@ export const RQKEY = (
     | 'locked-permanently'
     | undefined = undefined,
   limit?: number,
-) => [RQKEY_ROOT, status, readState, kind, lockStatus, limit]
+) => [RQKEY_ROOT, status, readState, kind, lockStatus, limit] as const
+
+/**
+ * Prefix key matching every convo-list query with the given status (and
+ * optionally readState), regardless of the remaining params (kind,
+ * lockStatus, limit). Only valid with prefix-matching APIs (setQueriesData,
+ * getQueriesData, invalidateQueries) - exact-match APIs (getQueryData,
+ * setQueryData) hash the full key and will never match a prefix.
+ */
+export const RQKEY_PARTIAL = (
+  status: 'accepted' | 'request' | 'all',
+  readState?: 'all' | 'unread',
+) => (readState ? [RQKEY_ROOT, status, readState] : [RQKEY_ROOT, status])
+
+/**
+ * Whether a convo satisfies the filters encoded in a convo-list query key.
+ * Caches are server-filtered, so optimistic inserts must apply the same
+ * filters client-side or convos leak into lists that should exclude them.
+ */
+export function convoMatchesQueryKey(
+  convo: ChatBskyConvoDefs.ConvoView,
+  queryKey: QueryKey,
+): boolean {
+  const [, status, readState, kind, lockStatus] = queryKey as ReturnType<
+    typeof RQKEY
+  >
+  if (status !== 'all' && status !== convo.status) return false
+  if (readState === 'unread' && convo.unreadCount === 0) return false
+  if (ChatBskyConvoDefs.isGroupConvo(convo.kind)) {
+    if (kind === 'direct') return false
+    if (lockStatus && convo.kind.lockStatus !== lockStatus) return false
+  } else {
+    if (kind === 'group') return false
+    // direct convos are never locked
+    if (lockStatus && lockStatus !== 'unlocked') return false
+  }
+  return true
+}
+
+/**
+ * Query predicate for optimistically upserting a convo into convo-list
+ * caches. Targets caches whose filters the convo satisfies, plus caches the
+ * convo is already in - those get updated in place even if the convo no
+ * longer matches (e.g. unreadCount dropped to 0), mirroring how read/mute
+ * log events update convos in place everywhere.
+ */
+export function convoListQueryPredicate(convo: ChatBskyConvoDefs.ConvoView) {
+  return (query: Query): boolean => {
+    const data = query.state.data as ConvoListQueryData | undefined
+    if (data && getConvoFromQueryData(convo.id, data)) return true
+    return convoMatchesQueryKey(convo, query.queryKey)
+  }
+}
+
 type RQPageParam = string | undefined
 
 export function useListConvosQuery({
@@ -347,9 +402,12 @@ export function ListConvosProviderInner({
                 }),
               }
             }
-            // always update the unread one
+            // always update the unread ones, where the convo qualifies
             queryClient.setQueriesData(
-              {queryKey: RQKEY('all', 'unread')},
+              {
+                queryKey: RQKEY_PARTIAL('all', 'unread'),
+                predicate: convoListQueryPredicate(updatedConvo),
+              },
               (old?: ConvoListQueryData) =>
                 old
                   ? updateFn(old)
@@ -361,11 +419,20 @@ export function ListConvosProviderInner({
             // update the other ones based on status of the incoming message
             if (updatedConvo.status === 'accepted') {
               queryClient.setQueriesData(
-                {queryKey: RQKEY('accepted')},
+                {
+                  queryKey: RQKEY_PARTIAL('accepted'),
+                  predicate: convoListQueryPredicate(updatedConvo),
+                },
                 updateFn,
               )
             } else if (updatedConvo.status === 'request') {
-              queryClient.setQueriesData({queryKey: RQKEY('request')}, updateFn)
+              queryClient.setQueriesData(
+                {
+                  queryKey: RQKEY_PARTIAL('request'),
+                  predicate: convoListQueryPredicate(updatedConvo),
+                },
+                updateFn,
+              )
               // also move-to-top in the new requests cache
               queryClient.setQueriesData<ConvoRequestListQueryData>(
                 {queryKey: [REQUESTS_RQKEY_ROOT]},
@@ -385,20 +452,26 @@ export function ListConvosProviderInner({
               rev: log.rev,
             }))
           } else if (ChatBskyConvoDefs.isLogAcceptConvo(log)) {
-            const requests = queryClient.getQueryData<ConvoListQueryData>(
-              RQKEY('request'),
-            )
-            if (!requests) {
-              debouncedRefetch()
-              return
+            const requestQueries =
+              queryClient.getQueriesData<ConvoListQueryData>({
+                queryKey: RQKEY_PARTIAL('request'),
+              })
+            let foundConvo: ChatBskyConvoDefs.ConvoView | null = null
+            for (const [_key, data] of requestQueries) {
+              if (!data) continue
+              foundConvo = getConvoFromQueryData(log.convoId, data)
+              if (foundConvo) break
             }
-            const acceptedConvo = getConvoFromQueryData(log.convoId, requests)
-            if (!acceptedConvo) {
+            if (!foundConvo) {
               debouncedRefetch()
               return
             }
-            queryClient.setQueryData(
-              RQKEY('request'),
+            const acceptedConvo: ChatBskyConvoDefs.ConvoView = {
+              ...foundConvo,
+              status: 'accepted',
+            }
+            queryClient.setQueriesData(
+              {queryKey: RQKEY_PARTIAL('request')},
               (old?: ConvoListQueryData) => optimisticDelete(log.convoId, old),
             )
             // also remove from the new requests cache
@@ -407,7 +480,10 @@ export function ListConvosProviderInner({
               old => optimisticDeleteRequest(log.convoId, old),
             )
             queryClient.setQueriesData(
-              {queryKey: RQKEY('accepted')},
+              {
+                queryKey: RQKEY_PARTIAL('accepted'),
+                predicate: convoListQueryPredicate(acceptedConvo),
+              },
               (old?: ConvoListQueryData) => {
                 if (!old) {
                   debouncedRefetch()
@@ -420,12 +496,15 @@ export function ListConvosProviderInner({
                       return {
                         ...page,
                         convos: [
-                          {...acceptedConvo, status: 'accepted'},
-                          ...page.convos,
+                          acceptedConvo,
+                          ...page.convos.filter(c => c.id !== log.convoId),
                         ],
                       }
                     }
-                    return page
+                    return {
+                      ...page,
+                      convos: page.convos.filter(c => c.id !== log.convoId),
+                    }
                   }),
                 }
               },
@@ -723,7 +802,14 @@ export function useUnreadMessageCount() {
         hasNew: false,
       }
     }
-  }, [accepted, request, currentAccount?.did, currentConvoId, moderationOpts])
+  }, [
+    accepted,
+    request,
+    currentAccount?.did,
+    currentConvoId,
+    moderationOpts,
+    aa.flags,
+  ])
 }
 
 function calculateCount(
@@ -911,7 +997,7 @@ function addMemberToConvoView(
   }
 }
 
-function optimisticDelete(chatId: string, old?: ConvoListQueryData) {
+export function optimisticDelete(chatId: string, old?: ConvoListQueryData) {
   if (!old) return old
 
   return {