[APP-1398] Fix post shadow/like shadow state when replying (#8866)

* Add optimisticReplyCount to post shadow

* Add special util

* Fix subtle bug in query cache

2 files changed, 40 insertions, 4 deletions

diff --git a/src/state/cache/post-shadow.ts b/src/state/cache/post-shadow.ts
index d7f1eb8b9..8cc3dca1a 100644
--- a/src/state/cache/post-shadow.ts
+++ b/src/state/cache/post-shadow.ts
@@ -24,6 +24,7 @@ export interface PostShadow {
   isDeleted: boolean
   embed: AppBskyEmbedRecord.View | AppBskyEmbedRecordWithMedia.View | undefined
   pinned: boolean
+  optimisticReplyCount: number | undefined
 }
 
 export const POST_TOMBSTONE = Symbol('PostTombstone')
@@ -34,6 +35,14 @@ const shadows: WeakMap<
   Partial<PostShadow>
 > = new WeakMap()
 
+/**
+ * Use with caution! This function returns the raw shadow data for a post.
+ * Prefer using `usePostShadow`.
+ */
+export function dangerousGetPostShadow(post: AppBskyFeedDefs.PostView) {
+  return shadows.get(post)
+}
+
 export function usePostShadow(
   post: AppBskyFeedDefs.PostView,
 ): Shadow<AppBskyFeedDefs.PostView> | typeof POST_TOMBSTONE {
@@ -95,6 +104,11 @@ function mergeShadow(
     repostCount = Math.max(0, repostCount)
   }
 
+  let replyCount = post.replyCount ?? 0
+  if ('optimisticReplyCount' in shadow) {
+    replyCount = shadow.optimisticReplyCount ?? replyCount
+  }
+
   let embed: typeof post.embed
   if ('embed' in shadow) {
     if (
@@ -112,6 +126,7 @@ function mergeShadow(
     embed: embed || post.embed,
     likeCount: likeCount,
     repostCount: repostCount,
+    replyCount: replyCount,
     viewer: {
       ...(post.viewer || {}),
       like: 'likeUri' in shadow ? shadow.likeUri : post.viewer?.like,
diff --git a/src/state/queries/usePostThread/queryCache.ts b/src/state/queries/usePostThread/queryCache.ts
index 826932349..5e27ebb87 100644
--- a/src/state/queries/usePostThread/queryCache.ts
+++ b/src/state/queries/usePostThread/queryCache.ts
@@ -9,6 +9,10 @@ import {
 } from '@atproto/api'
 import {type QueryClient} from '@tanstack/react-query'
 
+import {
+  dangerousGetPostShadow,
+  updatePostShadow,
+} from '#/state/cache/post-shadow'
 import {findAllPostsInQueryData as findAllPostsInExploreFeedPreviewsQueryData} from '#/state/queries/explore-feed-previews'
 import {findAllPostsInQueryData as findAllPostsInNotifsQueryData} from '#/state/queries/notifications/feed'
 import {findAllPostsInQueryData as findAllPostsInFeedQueryData} from '#/state/queries/post-feed'
@@ -85,10 +89,27 @@ export function createCacheMutator({
           /*
            * Update parent data
            */
-          parent.value.post = {
-            ...parent.value.post,
-            replyCount: (parent.value.post.replyCount || 0) + 1,
-          }
+          const shadow = dangerousGetPostShadow(parent.value.post)
+          const prevOptimisticCount = shadow?.optimisticReplyCount
+          const prevReplyCount = parent.value.post.replyCount
+          // prefer optimistic count, if we already have some
+          const currentReplyCount =
+            (prevOptimisticCount ?? prevReplyCount ?? 0) + 1
+
+          /*
+           * We must update the value in the query cache in order for thread
+           * traversal to properly compute required metadata.
+           */
+          parent.value.post.replyCount = currentReplyCount
+
+          /**
+           * Additionally, we need to update the post shadow to keep track of
+           * these new values, since mutating the post object above does not
+           * cause a re-render.
+           */
+          updatePostShadow(queryClient, parent.value.post.uri, {
+            optimisticReplyCount: currentReplyCount,
+          })
 
           const opDid = getRootPostAtUri(parent.value.post)?.host
           const nextPreexistingItem = thread.at(i + 1)