Send beta user state with AppView requests (#11474)

src/state/preferences/beta-user-cache.ts

@@ -0,0 +1,64 @@
+import {account} from '#/storage'
+
+const values = new Map<string, boolean | undefined>()
+const listeners = new Map<string, Set<() => void>>()
+
+/**
+ * Returns the last known beta-user preference for an account.
+ *
+ * The first read for a DID in this process is hydrated synchronously from
+ * persistent storage so cold starts retain the last value fetched from the
+ * PDS. Subsequent reads, including request-header reads, stay in memory.
+ */
+export function getCachedIsBetaUser(did: string): boolean | undefined {
+  if (!values.has(did)) {
+    try {
+      values.set(did, account.get([did, 'isBetaUser']))
+    } catch {
+      values.set(did, undefined)
+    }
+  }
+  return values.get(did)
+}
+
+/**
+ * Updates the runtime cache and its persistent cold-start snapshot.
+ *
+ * Call this only with a value confirmed by the PDS, which remains the source
+ * of truth for the preference.
+ */
+export function setCachedIsBetaUser(did: string, value: boolean): void {
+  if (getCachedIsBetaUser(did) === value) return
+  account.set([did, 'isBetaUser'], value)
+  values.set(did, value)
+  listeners.get(did)?.forEach(listener => listener())
+}
+
+/**
+ * Subscribe to runtime cache changes for one account.
+ */
+export function subscribeToCachedIsBetaUser(
+  did: string,
+  listener: () => void,
+): () => void {
+  let didListeners = listeners.get(did)
+  if (!didListeners) {
+    didListeners = new Set()
+    listeners.set(did, didListeners)
+  }
+  didListeners.add(listener)
+
+  return () => {
+    didListeners.delete(listener)
+    if (didListeners.size === 0) listeners.delete(did)
+  }
+}
+
+/**
+ * Drops the in-memory value so the next read rehydrates from persistence.
+ * This is useful when persistence is changed outside this cache.
+ */
+export function invalidateCachedIsBetaUser(did: string): void {
+  if (!values.delete(did)) return
+  listeners.get(did)?.forEach(listener => listener())
+}

src/state/preferences/beta-user-sync.tsx

@@ -2,12 +2,12 @@ import {useEffect} from 'react'
 
 import {usePreferencesQuery} from '#/state/queries/preferences'
 import {useSession} from '#/state/session'
-import {account} from '#/storage'
+import {getCachedIsBetaUser, setCachedIsBetaUser} from './beta-user-cache'
 
 /**
- * Caches `bskyAppState.isBetaUser` from preferences into synchronous device
- * storage so analytics can read it at init (before beta-gated features are
- * evaluated). Scoped per account, since `isBetaUser` is account-specific:
+ * Caches `bskyAppState.isBetaUser` from preferences in memory and synchronous
+ * device storage so analytics can read it at init (before beta-gated features
+ * are evaluated). Scoped per account, since `isBetaUser` is account-specific:
  * a global cache would let one account's value leak into another after a
  * switch, until the new account's preferences loaded. Must be mounted below
  * `QueryProvider`, since the analytics providers that consume the cached value
@@ -27,10 +27,10 @@ export function BetaUserStorageSync() {
     if (isBetaUser === undefined) return
     /*
      * Guard against a redundant write on every warm boot. Writing triggers the
-     * storage change listener, which re-renders the analytics subtree.
+     * cache change listener, which re-renders the analytics subtree.
      */
-    if (account.get([did, 'isBetaUser']) === isBetaUser) return
-    account.set([did, 'isBetaUser'], isBetaUser)
+    if (getCachedIsBetaUser(did) === isBetaUser) return
+    setCachedIsBetaUser(did, isBetaUser)
   }, [did, isBetaUser])
 
   return null

src/state/session/__tests__/clients-test.ts

@@ -14,7 +14,12 @@ jest.mock('jwt-decode', () => ({
 }))
 
 import {BLUESKY_PROXY_HEADER, CHAT_PROXY_SERVICE} from '#/lib/constants'
+import {
+  invalidateCachedIsBetaUser,
+  setCachedIsBetaUser,
+} from '#/state/preferences/beta-user-cache'
 import {app, chat, com} from '#/lexicons'
+import {account} from '#/storage'
 import {configureGlobalAppLabelers} from '../additional-moderation-authorities'
 import {
   buildAppviewClient,
@@ -84,6 +89,8 @@ describe('buildAppviewClient', () => {
   beforeEach(() => {
     fetchMock = makeProfileFetch()
     configureGlobalAppLabelers([])
+    account.remove([DID, 'isBetaUser'])
+    invalidateCachedIsBetaUser(DID)
   })
 
   it('passes through the session did', () => {
@@ -111,6 +118,46 @@ describe('buildAppviewClient', () => {
     ).toBe(BLUESKY_PROXY_HEADER.get())
   })
 
+  it.each([true, false])(
+    'emits the current beta user header when the cached value is %s',
+    async isBetaUser => {
+      const client = buildAppviewClient(makeSession(fetchMock))
+      setCachedIsBetaUser(DID, isBetaUser)
+
+      await client.call(app.bsky.actor.getProfile, {actor: HANDLE})
+
+      expect(
+        headersFor(fetchMock, 'app.bsky.actor.getProfile').get(
+          'x-bsky-is-beta-user',
+        ),
+      ).toBe(String(isBetaUser))
+    },
+  )
+
+  it('omits the beta user header when the preference is not cached', async () => {
+    const client = buildAppviewClient(makeSession(fetchMock))
+
+    await client.call(app.bsky.actor.getProfile, {actor: HANDLE})
+
+    expect(
+      headersFor(fetchMock, 'app.bsky.actor.getProfile').get(
+        'x-bsky-is-beta-user',
+      ),
+    ).toBeNull()
+  })
+
+  it('reads the persisted beta preference only once', async () => {
+    account.set([DID, 'isBetaUser'], true)
+    const getSpy = jest.spyOn(account, 'get')
+    const client = buildAppviewClient(makeSession(fetchMock))
+
+    await client.call(app.bsky.actor.getProfile, {actor: HANDLE})
+    await client.call(app.bsky.actor.getProfile, {actor: HANDLE})
+
+    expect(getSpy).toHaveBeenCalledTimes(1)
+    getSpy.mockRestore()
+  })
+
   it('emits an account subscription exactly once', async () => {
     const client = buildAppviewClient(makeSession(fetchMock))
     client.setLabelers(['did:plc:labeler'])

src/state/session/clients.ts

@@ -7,8 +7,33 @@ import {
   PUBLIC_BSKY_SERVICE,
 } from '#/lib/constants'
 import {createLexClient} from '#/lib/lexClient'
+import {getCachedIsBetaUser} from '#/state/preferences/beta-user-cache'
 import {networkAwareFetch} from './network'
 
+const IS_BETA_USER_HEADER = 'X-Bsky-Is-Beta-User'
+
+/**
+ * Add account-scoped headers to appview requests.
+ *
+ * Values are read from memory per request so preference changes are reflected
+ * immediately without rebuilding the session bundle.
+ */
+function withAppviewRequestHeaders(agent: Agent): Agent {
+  return {
+    get did() {
+      return agent.did
+    },
+    fetchHandler(path, init) {
+      const headers = new Headers(init?.headers)
+      const isBetaUser = agent.did ? getCachedIsBetaUser(agent.did) : undefined
+      if (isBetaUser !== undefined) {
+        headers.set(IS_BETA_USER_HEADER, String(isBetaUser))
+      }
+      return agent.fetchHandler(path, {...init, headers})
+    },
+  }
+}
+
 /**
  * Build the signed-in appview {@link Client}.
  *
@@ -28,7 +53,9 @@ import {networkAwareFetch} from './network'
  * fetch, which is `networkAwareFetch` wrapped in the disposal kill switch.
  */
 export function buildAppviewClient(agent: Agent): Client {
-  return createLexClient(agent, {service: BLUESKY_PROXY_HEADER.get()})
+  return createLexClient(withAppviewRequestHeaders(agent), {
+    service: BLUESKY_PROXY_HEADER.get(),
+  })
 }
 
 /**