Add PostThreadContext, cache mutator for threadgates on threads, pipe it through

src/components/dialogs/PostInteractionSettingsDialog.tsx

@@ -31,6 +31,10 @@ import {
   useSetThreadgateAllowMutation,
   useThreadgateViewQuery,
 } from '#/state/queries/threadgate'
+import {
+  PostThreadContextProvider,
+  usePostThreadContext,
+} from '#/state/queries/usePostThread'
 import {useAgent, useSession} from '#/state/session'
 import * as Toast from '#/view/com/util/Toast'
 import {atoms as a, useTheme} from '#/alf'
@@ -133,10 +137,13 @@ export type PostInteractionSettingsDialogProps = {
 export function PostInteractionSettingsDialog(
   props: PostInteractionSettingsDialogProps,
 ) {
+  const postThreadContext = usePostThreadContext()
   return (
     <Dialog.Outer control={props.control}>
       <Dialog.Handle />
-      <PostInteractionSettingsDialogControlledInner {...props} />
+      <PostThreadContextProvider context={postThreadContext}>
+        <PostInteractionSettingsDialogControlledInner {...props} />
+      </PostThreadContextProvider>
     </Dialog.Outer>
   )
 }

src/screens/PostThread/index.tsx

@@ -7,7 +7,11 @@ import {useInitialNumToRender} from '#/lib/hooks/useInitialNumToRender'
 import {useOpenComposer} from '#/lib/hooks/useOpenComposer'
 import {useFeedFeedback} from '#/state/feed-feedback'
 import {type ThreadViewOption} from '#/state/queries/preferences/useThreadPreferences'
-import {type ThreadItem, usePostThread} from '#/state/queries/usePostThread'
+import {
+  PostThreadContextProvider,
+  type ThreadItem,
+  usePostThread,
+} from '#/state/queries/usePostThread'
 import {useSession} from '#/state/session'
 import {type OnPostSuccessData} from '#/state/shell/composer'
 import {useShellLayout} from '#/state/shell/shell-layout'
@@ -495,7 +499,7 @@ export function PostThread({uri}: {uri: string}) {
   const defaultListFooterHeight = hasParents ? windowHeight - 200 : undefined
 
   return (
-    <>
+    <PostThreadContextProvider context={thread.context}>
       <Layout.Header.Outer headerRef={headerRef}>
         <Layout.Header.BackButton />
         <Layout.Header.Content>
@@ -578,7 +582,7 @@ export function PostThread({uri}: {uri: string}) {
       {!gtMobile && canReply && hasSession && (
         <MobileComposePrompt onPressReply={onReplyToAnchor} />
       )}
-    </>
+    </PostThreadContextProvider>
   )
 }
 

src/state/queries/threadgate/index.ts

@@ -1,6 +1,5 @@
 import {
   AppBskyFeedDefs,
-  type AppBskyFeedGetPostThread,
   AppBskyFeedThreadgate,
   AtUri,
   type BskyAgent,
@@ -8,8 +7,8 @@ import {
 import {useMutation, useQuery, useQueryClient} from '@tanstack/react-query'
 
 import {networkRetry, retry} from '#/lib/async/retry'
-import {until} from '#/lib/async/until'
 import {STALE} from '#/state/queries'
+import {useGetPost} from '#/state/queries/post'
 import {type ThreadgateAllowUISetting} from '#/state/queries/threadgate/types'
 import {
   createThreadgateRecord,
@@ -17,6 +16,7 @@ import {
   threadgateAllowUISettingToAllowRecordValue,
   threadgateViewToAllowUISetting,
 } from '#/state/queries/threadgate/util'
+import {useUpdatePostThreadThreadgateQueryCache} from '#/state/queries/usePostThread'
 import {useAgent} from '#/state/session'
 import {useThreadgateHiddenReplyUrisAPI} from '#/state/threadgate-hidden-replies'
 import * as bsky from '#/types/bsky'
@@ -247,6 +247,8 @@ export async function updateThreadgateAllow({
 export function useSetThreadgateAllowMutation() {
   const agent = useAgent()
   const queryClient = useQueryClient()
+  const getPost = useGetPost()
+  const updatePostThreadThreadgate = useUpdatePostThreadThreadgateQueryCache()
 
   return useMutation({
     mutationFn: async ({
@@ -271,26 +273,31 @@ export function useSetThreadgateAllowMutation() {
       })
     },
     async onSuccess(_, {postUri, allow}) {
-      await until(
+      const data = await retry<AppBskyFeedDefs.ThreadgateView | undefined>(
         5, // 5 tries
-        1e3, // 1s delay between tries
-        (res: AppBskyFeedGetPostThread.Response) => {
-          const thread = res.data.thread
-          if (AppBskyFeedDefs.isThreadViewPost(thread)) {
-            const fetchedSettings = threadgateViewToAllowUISetting(
-              thread.post.threadgate,
+        _e => true,
+        async () => {
+          const post = await getPost({uri: postUri})
+          const threadgate = post.threadgate
+          if (!threadgate) {
+            throw new Error(
+              `useSetThreadgateAllowMutation: could not fetch threadgate, appview may not be ready yet`,
             )
-            return JSON.stringify(fetchedSettings) === JSON.stringify(allow)
           }
-          return false
+          const fetchedSettings = threadgateViewToAllowUISetting(threadgate)
+          const isReady =
+            JSON.stringify(fetchedSettings) === JSON.stringify(allow)
+          if (!isReady) {
+            throw new Error(
+              `useSetThreadgateAllowMutation: appview isn't ready yet`,
+            ) // try again
+          }
+          return threadgate
         },
-        () => {
-          return agent.app.bsky.feed.getPostThread({
-            uri: postUri,
-            depth: 0,
-          })
-        },
-      )
+        1e3, // 1s delay between tries
+      ).catch(() => {})
+
+      if (data) updatePostThreadThreadgate(data)
 
       queryClient.invalidateQueries({
         queryKey: [threadgateRecordQueryKeyRoot],

src/state/queries/usePostThread/context.tsx

@@ -0,0 +1,43 @@
+import {createContext, useContext} from 'react'
+
+import {
+  type createPostThreadOtherQueryKey,
+  type createPostThreadQueryKey,
+} from '#/state/queries/usePostThread/types'
+
+/**
+ * Contains static metadata about the post thread query, suitable for
+ * context e.g. query keys and other things that don't update frequently.
+ *
+ * Be careful adding things here, as it could cause unnecessary re-renders.
+ */
+export type PostThreadContextType = {
+  postThreadQueryKey: ReturnType<typeof createPostThreadQueryKey>
+  postThreadOtherQueryKey: ReturnType<typeof createPostThreadOtherQueryKey>
+}
+
+const PostThreadContext = createContext<PostThreadContextType | undefined>(
+  undefined,
+)
+
+/**
+ * Use the current {@link PostThreadContext}, if one is available. If not,
+ * returns `undefined`.
+ */
+export function usePostThreadContext() {
+  return useContext(PostThreadContext)
+}
+
+export function PostThreadContextProvider({
+  children,
+  context,
+}: {
+  children: React.ReactNode
+  context?: PostThreadContextType
+}) {
+  return (
+    <PostThreadContext.Provider value={context}>
+      {children}
+    </PostThreadContext.Provider>
+  )
+}

src/state/queries/usePostThread/index.tsx

@@ -31,6 +31,8 @@ import {useAgent, useSession} from '#/state/session'
 import {useMergeThreadgateHiddenReplies} from '#/state/threadgate-hidden-replies'
 import {useBreakpoints} from '#/alf'
 
+export * from '#/state/queries/usePostThread/context'
+export {useUpdatePostThreadThreadgateQueryCache} from '#/state/queries/usePostThread/queryCache'
 export * from '#/state/queries/usePostThread/types'
 
 export function usePostThread({anchor}: {anchor?: string}) {
@@ -277,8 +279,13 @@ export function usePostThread({anchor}: {anchor?: string}) {
     setOtherItemsVisible,
   ])
 
-  return useMemo(
-    () => ({
+  return useMemo(() => {
+    const context: PostThreadContextType = {
+      postThreadQueryKey,
+      postThreadOtherQueryKey,
+    }
+    return {
+      context,
       state: {
         /*
          * Copy in any query state that is useful
@@ -309,17 +316,18 @@ export function usePostThread({anchor}: {anchor?: string}) {
         setSort,
         setView,
       },
-    }),
-    [
-      query,
-      mutator.insertReplies,
-      otherItemsVisible,
-      sort,
-      view,
-      setSort,
-      setView,
-      threadgate,
-      items,
-    ],
-  )
+    }
+  }, [
+    query,
+    mutator.insertReplies,
+    otherItemsVisible,
+    sort,
+    view,
+    setSort,
+    setView,
+    threadgate,
+    items,
+    postThreadQueryKey,
+    postThreadOtherQueryKey,
+  ])
 }

src/state/queries/usePostThread/queryCache.ts

@@ -1,3 +1,4 @@
+import {useCallback} from 'react'
 import {
   type $Typed,
   type AppBskyActorDefs,
@@ -7,7 +8,7 @@ import {
   type AppBskyUnspeccedGetPostThreadV2,
   AtUri,
 } from '@atproto/api'
-import {type QueryClient} from '@tanstack/react-query'
+import {type QueryClient, useQueryClient} from '@tanstack/react-query'
 
 import {
   dangerousGetPostShadow,
@@ -18,6 +19,7 @@ import {findAllPostsInQueryData as findAllPostsInNotifsQueryData} from '#/state/
 import {findAllPostsInQueryData as findAllPostsInFeedQueryData} from '#/state/queries/post-feed'
 import {findAllPostsInQueryData as findAllPostsInQuoteQueryData} from '#/state/queries/post-quotes'
 import {findAllPostsInQueryData as findAllPostsInSearchQueryData} from '#/state/queries/search-posts'
+import {usePostThreadContext} from '#/state/queries/usePostThread'
 import {getBranch} from '#/state/queries/usePostThread/traversal'
 import {
   type ApiThreadItem,
@@ -322,3 +324,51 @@ export function* findAllProfilesInQueryData(
     }
   }
 }
+
+export function useUpdatePostThreadThreadgateQueryCache() {
+  const qc = useQueryClient()
+  const context = usePostThreadContext()
+
+  return useCallback(
+    (threadgate: AppBskyFeedDefs.ThreadgateView) => {
+      if (!context) return
+
+      function mutator<T>(thread: ApiThreadItem[]): T[] {
+        for (let i = 0; i < thread.length; i++) {
+          const item = thread[i]
+
+          if (!AppBskyUnspeccedDefs.isThreadItemPost(item.value)) continue
+
+          if (item.depth === 0) {
+            thread.splice(i, 1, {
+              ...item,
+              value: {
+                ...item.value,
+                post: {
+                  ...item.value.post,
+                  threadgate,
+                },
+              },
+            })
+          }
+        }
+
+        return thread as T[]
+      }
+
+      qc.setQueryData<AppBskyUnspeccedGetPostThreadV2.OutputSchema>(
+        context.postThreadQueryKey,
+        data => {
+          if (!data) return
+          return {
+            ...data,
+            thread: mutator<AppBskyUnspeccedGetPostThreadV2.ThreadItem>([
+              ...data.thread,
+            ]),
+          }
+        },
+      )
+    },
+    [qc, context],
+  )
+}