From 3247c7e28f8a7ae3135852c727fc74c21d0b0a9a Mon Sep 17 00:00:00 2001
From: Eric Bailey <git@esb.lol>
Date: Tue, 10 Jun 2025 17:38:54 -0500
Subject: [PATCH] Key post-source using URI with handle

---
 src/state/unstable-post-source.tsx          | 44 +++++++++++++++------
 src/view/com/post-thread/PostThreadItem.tsx |  9 ++++-
 src/view/com/posts/PostFeedItem.tsx         |  7 +++-
 3 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/src/state/unstable-post-source.tsx b/src/state/unstable-post-source.tsx
index 467f5dddcc..2a741e615d 100644
--- a/src/state/unstable-post-source.tsx
+++ b/src/state/unstable-post-source.tsx
@@ -1,5 +1,5 @@
 import {createContext, useCallback, useContext, useRef, useState} from 'react'
-import {type AppBskyFeedDefs} from '@atproto/api'
+import {type AppBskyFeedDefs,AtUri} from '@atproto/api'
 
 import {Logger} from '#/logger'
 import {type FeedDescriptor} from '#/state/queries/post-feed'
@@ -28,22 +28,36 @@ export function Provider({children}: {children: React.ReactNode}) {
   const sourcesRef = useRef<Map<string, Source>>(new Map())
 
   const setUnstablePostSource = useCallback((uri: string, source: Source) => {
+    if (__DEV__) {
+      const urip = new AtUri(uri)
+      if (urip.host.startsWith('did:')) {
+        throw new Error(
+          `URI passed to setUnstablePostSource should contain a handle — use buildPostSourceUri`,
+        )
+      }
+    }
+
+    logger.debug('set', {uri, source})
     sourcesRef.current.set(uri, source)
-    logger.debug('set', {
-      uri,
-      source,
-    })
   }, [])
 
   const consumeUnstablePostSource = useCallback((uri: string) => {
-    const source = sourcesRef.current.get(uri)
-    if (source) {
-      sourcesRef.current.delete(uri)
-      logger.debug('consume', {
-        uri,
-        source,
-      })
+    if (__DEV__) {
+      const urip = new AtUri(uri)
+      if (urip.host.startsWith('did:')) {
+        throw new Error(
+          `URI passed to consumeUnstablePostSource should contain a handle — use buildPostSourceUri`,
+        )
+      }
     }
+
+    const source = sourcesRef.current.get(uri)
+
+    if (source) {
+      logger.debug('consume', {uri, source})
+      sourcesRef.current.delete(uri)
+    }
+
     return source
   }, [])
 
@@ -71,3 +85,9 @@ export function useUnstablePostSource(uri: string) {
   const [source] = useState(() => consume(uri))
   return source
 }
+
+export function buildPostSourceUri(uri: string, handle: string) {
+  const urip = new AtUri(uri)
+  urip.host = handle
+  return urip.toString()
+}
diff --git a/src/view/com/post-thread/PostThreadItem.tsx b/src/view/com/post-thread/PostThreadItem.tsx
index 8b39072ba6..1c1046247e 100644
--- a/src/view/com/post-thread/PostThreadItem.tsx
+++ b/src/view/com/post-thread/PostThreadItem.tsx
@@ -40,7 +40,10 @@ import {useLanguagePrefs} from '#/state/preferences'
 import {type ThreadPost} from '#/state/queries/post-thread'
 import {useSession} from '#/state/session'
 import {useMergedThreadgateHiddenReplies} from '#/state/threadgate-hidden-replies'
-import {useUnstablePostSource} from '#/state/unstable-post-source'
+import {
+  buildPostSourceUri,
+  useUnstablePostSource,
+} from '#/state/unstable-post-source'
 import {PostThreadFollowBtn} from '#/view/com/post-thread/PostThreadFollowBtn'
 import {ErrorMessage} from '#/view/com/util/error/ErrorMessage'
 import {Link, TextLink} from '#/view/com/util/Link'
@@ -204,7 +207,9 @@ let PostThreadItemLoaded = ({
   threadgateRecord?: AppBskyFeedThreadgate.Record
 }): React.ReactNode => {
   const {currentAccount, hasSession} = useSession()
-  const source = useUnstablePostSource(post.uri)
+  const source = useUnstablePostSource(
+    buildPostSourceUri(post.uri, post.author.handle),
+  )
   const feedFeedback = useFeedFeedback(source?.feed, hasSession)
 
   const t = useTheme()
diff --git a/src/view/com/posts/PostFeedItem.tsx b/src/view/com/posts/PostFeedItem.tsx
index b9aa676737..88a63fee3b 100644
--- a/src/view/com/posts/PostFeedItem.tsx
+++ b/src/view/com/posts/PostFeedItem.tsx
@@ -36,7 +36,10 @@ import {useFeedFeedbackContext} from '#/state/feed-feedback'
 import {unstableCacheProfileView} from '#/state/queries/profile'
 import {useSession} from '#/state/session'
 import {useMergedThreadgateHiddenReplies} from '#/state/threadgate-hidden-replies'
-import {useSetUnstablePostSource} from '#/state/unstable-post-source'
+import {
+  buildPostSourceUri,
+  useSetUnstablePostSource,
+} from '#/state/unstable-post-source'
 import {FeedNameText} from '#/view/com/util/FeedInfoText'
 import {Link, TextLink, TextLinkOnWebOnly} from '#/view/com/util/Link'
 import {PostEmbeds, PostEmbedViewContext} from '#/view/com/util/post-embeds'
@@ -232,7 +235,7 @@ let FeedItemInner = ({
       reqId,
     })
     unstableCacheProfileView(queryClient, post.author)
-    unstableSetPostSource(post.uri, {
+    unstableSetPostSource(buildPostSourceUri(post.uri, post.author.handle), {
       feed: feedDescriptor,
       post: {
         post,