5 files changed, 90 insertions, 32 deletions
diff --git a/src/state/queries/notifications/feed.ts b/src/state/queries/notifications/feed.ts
index 40be2ce8e..d9f019af3 100644
--- a/src/state/queries/notifications/feed.ts
+++ b/src/state/queries/notifications/feed.ts
@@ -17,7 +17,7 @@
  */
 
 import {useEffect, useRef} from 'react'
-import {AppBskyActorDefs, AppBskyFeedDefs} from '@atproto/api'
+import {AppBskyActorDefs, AppBskyFeedDefs, AtUri} from '@atproto/api'
 import {
   InfiniteData,
   QueryClient,
@@ -30,7 +30,11 @@ import {useMutedThreads} from '#/state/muted-threads'
 import {useAgent} from '#/state/session'
 import {useModerationOpts} from '../../preferences/moderation-opts'
 import {STALE} from '..'
-import {embedViewRecordToPostView, getEmbeddedPost} from '../util'
+import {
+  didOrHandleUriMatches,
+  embedViewRecordToPostView,
+  getEmbeddedPost,
+} from '../util'
 import {FeedPage} from './types'
 import {useUnreadNotificationsApi} from './unread'
 import {fetchPage} from './util'
@@ -142,6 +146,8 @@ export function* findAllPostsInQueryData(
   queryClient: QueryClient,
   uri: string,
 ): Generator<AppBskyFeedDefs.PostView, void> {
+  const atUri = new AtUri(uri)
+
   const queryDatas = queryClient.getQueriesData<InfiniteData<FeedPage>>({
     queryKey: [RQKEY_ROOT],
   })
@@ -149,14 +155,16 @@ export function* findAllPostsInQueryData(
     if (!queryData?.pages) {
       continue
     }
+
     for (const page of queryData?.pages) {
       for (const item of page.items) {
-        if (item.subject?.uri === uri) {
+        if (item.subject && didOrHandleUriMatches(atUri, item.subject)) {
           yield item.subject
         }
+
         const quotedPost = getEmbeddedPost(item.subject?.embed)
-        if (quotedPost?.uri === uri) {
-          yield embedViewRecordToPostView(quotedPost)
+        if (quotedPost && didOrHandleUriMatches(atUri, quotedPost)) {
+          yield embedViewRecordToPostView(quotedPost!)
         }
       }
     }
diff --git a/src/state/queries/post-feed.ts b/src/state/queries/post-feed.ts
index 5c483483a..2fb80de37 100644
--- a/src/state/queries/post-feed.ts
+++ b/src/state/queries/post-feed.ts
@@ -35,7 +35,11 @@ import {KnownError} from '#/view/com/posts/FeedErrorMessage'
 import {useFeedTuners} from '../preferences/feed-tuners'
 import {useModerationOpts} from '../preferences/moderation-opts'
 import {usePreferencesQuery} from './preferences'
-import {embedViewRecordToPostView, getEmbeddedPost} from './util'
+import {
+  didOrHandleUriMatches,
+  embedViewRecordToPostView,
+  getEmbeddedPost,
+} from './util'
 
 type ActorDid = string
 type AuthorFilter =
@@ -448,6 +452,8 @@ export function* findAllPostsInQueryData(
   queryClient: QueryClient,
   uri: string,
 ): Generator<AppBskyFeedDefs.PostView, undefined> {
+  const atUri = new AtUri(uri)
+
   const queryDatas = queryClient.getQueriesData<
     InfiniteData<FeedPageUnselected>
   >({
@@ -459,24 +465,38 @@ export function* findAllPostsInQueryData(
     }
     for (const page of queryData?.pages) {
       for (const item of page.feed) {
-        if (item.post.uri === uri) {
+        if (didOrHandleUriMatches(atUri, item.post)) {
           yield item.post
         }
+
         const quotedPost = getEmbeddedPost(item.post.embed)
-        if (quotedPost?.uri === uri) {
+        if (quotedPost && didOrHandleUriMatches(atUri, quotedPost)) {
           yield embedViewRecordToPostView(quotedPost)
         }
-        if (
-          AppBskyFeedDefs.isPostView(item.reply?.parent) &&
-          item.reply?.parent?.uri === uri
-        ) {
-          yield item.reply.parent
+
+        if (AppBskyFeedDefs.isPostView(item.reply?.parent)) {
+          if (didOrHandleUriMatches(atUri, item.reply.parent)) {
+            yield item.reply.parent
+          }
+
+          const parentQuotedPost = getEmbeddedPost(item.reply.parent.embed)
+          if (
+            parentQuotedPost &&
+            didOrHandleUriMatches(atUri, parentQuotedPost)
+          ) {
+            yield embedViewRecordToPostView(parentQuotedPost)
+          }
         }
-        if (
-          AppBskyFeedDefs.isPostView(item.reply?.root) &&
-          item.reply?.root?.uri === uri
-        ) {
-          yield item.reply.root
+
+        if (AppBskyFeedDefs.isPostView(item.reply?.root)) {
+          if (didOrHandleUriMatches(atUri, item.reply.root)) {
+            yield item.reply.root
+          }
+
+          const rootQuotedPost = getEmbeddedPost(item.reply.root.embed)
+          if (rootQuotedPost && didOrHandleUriMatches(atUri, rootQuotedPost)) {
+            yield embedViewRecordToPostView(rootQuotedPost)
+          }
         }
       }
     }
diff --git a/src/state/queries/post-thread.ts b/src/state/queries/post-thread.ts
index b1bff1493..f7d21a427 100644
--- a/src/state/queries/post-thread.ts
+++ b/src/state/queries/post-thread.ts
@@ -4,6 +4,7 @@ import {
   AppBskyFeedDefs,
   AppBskyFeedGetPostThread,
   AppBskyFeedPost,
+  AtUri,
   ModerationDecision,
   ModerationOpts,
 } from '@atproto/api'
@@ -24,7 +25,11 @@ import {
   findAllPostsInQueryData as findAllPostsInFeedQueryData,
   findAllProfilesInQueryData as findAllProfilesInFeedQueryData,
 } from './post-feed'
-import {embedViewRecordToPostView, getEmbeddedPost} from './util'
+import {
+  didOrHandleUriMatches,
+  embedViewRecordToPostView,
+  getEmbeddedPost,
+} from './util'
 
 const RQKEY_ROOT = 'post-thread'
 export const RQKEY = (uri: string) => [RQKEY_ROOT, uri]
@@ -91,14 +96,10 @@ export function usePostThreadQuery(uri: string | undefined) {
     },
     enabled: !!uri,
     placeholderData: () => {
-      if (!uri) {
-        return undefined
-      }
-      {
-        const post = findPostInQueryData(queryClient, uri)
-        if (post) {
-          return post
-        }
+      if (!uri) return
+      const post = findPostInQueryData(queryClient, uri)
+      if (post) {
+        return post
       }
       return undefined
     },
@@ -271,6 +272,8 @@ export function* findAllPostsInQueryData(
   queryClient: QueryClient,
   uri: string,
 ): Generator<ThreadNode, void> {
+  const atUri = new AtUri(uri)
+
   const queryDatas = queryClient.getQueriesData<ThreadNode>({
     queryKey: [RQKEY_ROOT],
   })
@@ -279,7 +282,7 @@ export function* findAllPostsInQueryData(
       continue
     }
     for (const item of traverseThread(queryData)) {
-      if (item.uri === uri) {
+      if (item.type === 'post' && didOrHandleUriMatches(atUri, item.post)) {
         const placeholder = threadNodeToPlaceholderThread(item)
         if (placeholder) {
           yield placeholder
@@ -287,7 +290,7 @@ export function* findAllPostsInQueryData(
       }
       const quotedPost =
         item.type === 'post' ? getEmbeddedPost(item.post.embed) : undefined
-      if (quotedPost?.uri === uri) {
+      if (quotedPost && didOrHandleUriMatches(atUri, quotedPost)) {
         yield embedViewRecordToPlaceholderThread(quotedPost)
       }
     }
diff --git a/src/state/queries/search-posts.ts b/src/state/queries/search-posts.ts
index f71d64255..5c50ad267 100644
--- a/src/state/queries/search-posts.ts
+++ b/src/state/queries/search-posts.ts
@@ -2,6 +2,7 @@ import {
   AppBskyActorDefs,
   AppBskyFeedDefs,
   AppBskyFeedSearchPosts,
+  AtUri,
 } from '@atproto/api'
 import {
   InfiniteData,
@@ -11,7 +12,11 @@ import {
 } from '@tanstack/react-query'
 
 import {useAgent} from '#/state/session'
-import {embedViewRecordToPostView, getEmbeddedPost} from './util'
+import {
+  didOrHandleUriMatches,
+  embedViewRecordToPostView,
+  getEmbeddedPost,
+} from './util'
 
 const searchPostsQueryKeyRoot = 'search-posts'
 const searchPostsQueryKey = ({query, sort}: {query: string; sort?: string}) => [
@@ -62,17 +67,20 @@ export function* findAllPostsInQueryData(
   >({
     queryKey: [searchPostsQueryKeyRoot],
   })
+  const atUri = new AtUri(uri)
+
   for (const [_queryKey, queryData] of queryDatas) {
     if (!queryData?.pages) {
       continue
     }
     for (const page of queryData?.pages) {
       for (const post of page.posts) {
-        if (post.uri === uri) {
+        if (didOrHandleUriMatches(atUri, post)) {
           yield post
         }
+
         const quotedPost = getEmbeddedPost(post.embed)
-        if (quotedPost?.uri === uri) {
+        if (quotedPost && didOrHandleUriMatches(atUri, quotedPost)) {
           yield embedViewRecordToPostView(quotedPost)
         }
       }
diff --git a/src/state/queries/util.ts b/src/state/queries/util.ts
index b74893fcd..f733c3788 100644
--- a/src/state/queries/util.ts
+++ b/src/state/queries/util.ts
@@ -1,8 +1,10 @@
 import {
+  AppBskyActorDefs,
   AppBskyEmbedRecord,
   AppBskyEmbedRecordWithMedia,
   AppBskyFeedDefs,
   AppBskyFeedPost,
+  AtUri,
 } from '@atproto/api'
 import {InfiniteData, QueryClient, QueryKey} from '@tanstack/react-query'
 
@@ -22,6 +24,23 @@ export function truncateAndInvalidate<T = any>(
   queryClient.invalidateQueries({queryKey})
 }
 
+// Given an AtUri, this function will check if the AtUri matches a
+// hit regardless of whether the AtUri uses a DID or handle as a host.
+//
+// AtUri should be the URI that is being searched for, while currentUri
+// is the URI that is being checked. currentAuthor is the author
+// of the currentUri that is being checked.
+export function didOrHandleUriMatches(
+  atUri: AtUri,
+  record: {uri: string; author: AppBskyActorDefs.ProfileViewBasic},
+) {
+  if (atUri.host.startsWith('did:')) {
+    return atUri.href === record.uri
+  }
+
+  return atUri.host === record.author.handle && record.uri.endsWith(atUri.rkey)
+}
+
 export function getEmbeddedPost(
   v: unknown,
 ): AppBskyEmbedRecord.ViewRecord | undefined {