3 files changed, 33 insertions, 67 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 5035065a0..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,
@@ -391,9 +374,6 @@ export function* findAllPostsInQueryData(
   >({
     queryKey: ['post-feed'],
   })
-
-  let foundEmbed: AppBskyFeedDefs.PostView | undefined
-
   for (const [_queryKey, queryData] of queryDatas) {
     if (!queryData?.pages) {
       continue
@@ -405,7 +385,7 @@ export function* findAllPostsInQueryData(
         }
         const quotedPost = getEmbeddedPost(item.post.embed)
         if (quotedPost?.uri === uri) {
-          foundEmbed = embedViewRecordToPostView(quotedPost)
+          yield embedViewRecordToPostView(quotedPost)
         }
         if (
           AppBskyFeedDefs.isPostView(item.reply?.parent) &&
@@ -422,10 +402,6 @@ export function* findAllPostsInQueryData(
       }
     }
   }
-
-  if (foundEmbed) {
-    yield foundEmbed
-  }
 }
 
 function assertSomePostsPassModeration(feed: AppBskyFeedDefs.FeedViewPost[]) {
diff --git a/src/state/queries/post-thread.ts b/src/state/queries/post-thread.ts
index ba4243163..56173207d 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
@@ -213,14 +201,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 +234,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 +246,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> {