4 files changed, 41 insertions, 60 deletions
diff --git a/src/state/queries/notifications/feed.ts b/src/state/queries/notifications/feed.ts
index b91db9237..405d054d4 100644
--- a/src/state/queries/notifications/feed.ts
+++ b/src/state/queries/notifications/feed.ts
@@ -133,23 +133,6 @@ export function useNotificationFeedQuery(opts?: {enabled?: boolean}) {
   return query
 }
 
-/**
- * This helper is used by the post-thread placeholder function to
- * find a post in the query-data cache
- */
-export function findPostInQueryData(
-  queryClient: QueryClient,
-  uri: string,
-): AppBskyFeedDefs.PostView | undefined {
-  const generator = findAllPostsInQueryData(queryClient, uri)
-  const result = generator.next()
-  if (result.done) {
-    return undefined
-  } else {
-    return result.value
-  }
-}
-
 export function* findAllPostsInQueryData(
   queryClient: QueryClient,
   uri: string,
diff --git a/src/state/queries/post-feed.ts b/src/state/queries/post-feed.ts
index 320009089..40399395a 100644
--- a/src/state/queries/post-feed.ts
+++ b/src/state/queries/post-feed.ts
@@ -365,23 +365,6 @@ function createApi(
   }
 }
 
-/**
- * This helper is used by the post-thread placeholder function to
- * find a post in the query-data cache
- */
-export function findPostInQueryData(
-  queryClient: QueryClient,
-  uri: string,
-): AppBskyFeedDefs.PostView | undefined {
-  const generator = findAllPostsInQueryData(queryClient, uri)
-  const result = generator.next()
-  if (result.done) {
-    return undefined
-  } else {
-    return result.value
-  }
-}
-
 export function* findAllPostsInQueryData(
   queryClient: QueryClient,
   uri: string,
diff --git a/src/state/queries/post-thread.ts b/src/state/queries/post-thread.ts
index ba4243163..26d40599c 100644
--- a/src/state/queries/post-thread.ts
+++ b/src/state/queries/post-thread.ts
@@ -8,8 +8,8 @@ import {useQuery, useQueryClient, QueryClient} from '@tanstack/react-query'
 
 import {getAgent} from '#/state/session'
 import {UsePreferencesQueryResponse} from '#/state/queries/preferences/types'
-import {findPostInQueryData as findPostInFeedQueryData} from './post-feed'
-import {findPostInQueryData as findPostInNotifsQueryData} from './notifications/feed'
+import {findAllPostsInQueryData as findAllPostsInFeedQueryData} from './post-feed'
+import {findAllPostsInQueryData as findAllPostsInNotifsQueryData} from './notifications/feed'
 import {precacheThreadPostProfiles} from './profile'
 import {getEmbeddedPost} from './util'
 
@@ -82,21 +82,9 @@ export function usePostThreadQuery(uri: string | undefined) {
         return undefined
       }
       {
-        const item = findPostInQueryData(queryClient, uri)
-        if (item) {
-          return threadNodeToPlaceholderThread(item)
-        }
-      }
-      {
-        const item = findPostInFeedQueryData(queryClient, uri)
-        if (item) {
-          return postViewToPlaceholderThread(item)
-        }
-      }
-      {
-        const item = findPostInNotifsQueryData(queryClient, uri)
-        if (item) {
-          return postViewToPlaceholderThread(item)
+        const post = findPostInQueryData(queryClient, uri)
+        if (post) {
+          return post
         }
       }
       return undefined
@@ -171,11 +159,18 @@ function responseToThreadNodes(
     AppBskyFeedPost.isRecord(node.post.record) &&
     AppBskyFeedPost.validateRecord(node.post.record).success
   ) {
+    const post = node.post
+    // These should normally be present. They're missing only for
+    // posts that were *just* created. Ideally, the backend would
+    // know to return zeros. Fill them in manually to compensate.
+    post.replyCount ??= 0
+    post.likeCount ??= 0
+    post.repostCount ??= 0
     return {
       type: 'post',
       _reactKey: node.post.uri,
       uri: node.post.uri,
-      post: node.post,
+      post: post,
       record: node.post.record,
       parent:
         node.parent && direction !== 'down'
@@ -213,14 +208,24 @@ function responseToThreadNodes(
 function findPostInQueryData(
   queryClient: QueryClient,
   uri: string,
-): ThreadNode | undefined {
-  const generator = findAllPostsInQueryData(queryClient, uri)
-  const result = generator.next()
-  if (result.done) {
-    return undefined
-  } else {
-    return result.value
+): ThreadNode | void {
+  let partial
+  for (let item of findAllPostsInQueryData(queryClient, uri)) {
+    if (item.type === 'post') {
+      // Currently, the backend doesn't send full post info in some cases
+      // (for example, for quoted posts). We use missing `likeCount`
+      // as a way to detect that. In the future, we should fix this on
+      // the backend, which will let us always stop on the first result.
+      const hasAllInfo = item.post.likeCount != null
+      if (hasAllInfo) {
+        return item
+      } else {
+        partial = item
+        // Keep searching, we might still find a full post in the cache.
+      }
+    }
   }
+  return partial
 }
 
 export function* findAllPostsInQueryData(
@@ -236,7 +241,10 @@ export function* findAllPostsInQueryData(
     }
     for (const item of traverseThread(queryData)) {
       if (item.uri === uri) {
-        yield item
+        const placeholder = threadNodeToPlaceholderThread(item)
+        if (placeholder) {
+          yield placeholder
+        }
       }
       const quotedPost =
         item.type === 'post' ? getEmbeddedPost(item.post.embed) : undefined
@@ -245,6 +253,12 @@ export function* findAllPostsInQueryData(
       }
     }
   }
+  for (let post of findAllPostsInFeedQueryData(queryClient, uri)) {
+    yield postViewToPlaceholderThread(post)
+  }
+  for (let post of findAllPostsInNotifsQueryData(queryClient, uri)) {
+    yield postViewToPlaceholderThread(post)
+  }
 }
 
 function* traverseThread(node: ThreadNode): Generator<ThreadNode, void> {
diff --git a/src/state/queries/util.ts b/src/state/queries/util.ts
index f3a87ae5d..54752b332 100644
--- a/src/state/queries/util.ts
+++ b/src/state/queries/util.ts
@@ -53,5 +53,6 @@ export function embedViewRecordToPostView(
     record: v.value,
     indexedAt: v.indexedAt,
     labels: v.labels,
+    embed: v.embeds?.[0],
   }
 }