phase 4: migrate ageAssurance unit and chat wrapper signatures to lex clients

src/ageAssurance/components/RedirectOverlay.tsx

@@ -16,7 +16,7 @@ import {Trans} from '@lingui/react/macro'
 import {retry} from '#/lib/async/retry'
 import {wait} from '#/lib/async/wait'
 import {parseLinkingUrl} from '#/lib/parseLinkingUrl'
-import {useAgent, useSession} from '#/state/session'
+import {useAppviewClient, useSession} from '#/state/session'
 import {atoms as a, platform, useBreakpoints, useTheme} from '#/alf'
 import {AgeAssuranceBadge} from '#/components/ageAssurance/AgeAssuranceBadge'
 import {Button, ButtonText} from '#/components/Button'
@@ -176,7 +176,8 @@ function Inner() {
   const t = useTheme()
   const ax = useAnalytics()
   const {_} = useLingui()
-  const agent = useAgent()
+  const {hasSession} = useSession()
+  const appviewClient = useAppviewClient()
   const polling = useRef(false)
   const unmounted = useRef(false)
   const [error, setError] = useState(false)
@@ -196,10 +197,10 @@ function Inner() {
         5,
         () => true,
         async () => {
-          if (!agent.session) return
+          if (!hasSession) return
           if (unmounted.current) return
 
-          const data = await refetchAgeAssuranceServerState({agent})
+          const data = await refetchAgeAssuranceServerState({appviewClient})
 
           if (data?.state.status !== 'assured') {
             throw new Error(
@@ -214,7 +215,7 @@ function Inner() {
     )
       .then(async data => {
         if (!data) return
-        if (!agent.session) return
+        if (!hasSession) return
         if (unmounted.current) return
 
         setSuccess(true)
@@ -230,7 +231,7 @@ function Inner() {
     return () => {
       unmounted.current = true
     }
-  }, [ax, agent])
+  }, [ax, hasSession, appviewClient])
 
   if (success) {
     return (

src/ageAssurance/data.tsx

@@ -1,6 +1,7 @@
 import {createContext, useCallback, useContext, useEffect, useMemo} from 'react'
 import * as AgeRange from 'expo-age-range'
 import {Client} from '@atproto/lex-client'
+import {getPreferences} from '@bsky.app/sdk'
 import {createAsyncStoragePersister} from '@tanstack/query-async-storage-persister'
 import {focusManager, QueryClient, useQuery} from '@tanstack/react-query'
 import {persistQueryClient} from '@tanstack/react-query-persist-client'
@@ -15,7 +16,7 @@ import {
   snoozeBirthdateUpdateAllowedForDid,
 } from '#/state/birthdate'
 import {fetchActorDeclarationRecord} from '#/state/queries/messages/actor-declaration'
-import {type SessionAgent, useAgent, useSession} from '#/state/session'
+import {useAppviewClient, usePdsClient, useSession} from '#/state/session'
 import {DEVICE_SIGNALS_SUPPORTED} from '#/ageAssurance/const'
 import * as debug from '#/ageAssurance/debug'
 import {logger} from '#/ageAssurance/logger'
@@ -33,7 +34,6 @@ import {IS_DEV} from '#/env'
 import {useGeolocation} from '#/geolocation'
 import {app, type chat} from '#/lexicons'
 import {device} from '#/storage'
-import {toLex} from '#/types/bsky'
 
 /**
  * Special query client for age assurance data so we can prefetch on app
@@ -59,10 +59,12 @@ const [, cacheHydrationPromise] = persistQueryClient({
   persister,
 })
 
-export function getDidFromAgentSession(agent: SessionAgent) {
+/**
-  const sessionManager = agent.sessionManager
+ * Reads the active account did from a lex {@link Client}. Logged-out clients
-  if (!sessionManager || !sessionManager.did) return
+ * expose `did: undefined`, so this returns undefined in that case.
-  return sessionManager.did
+ */
+export function getDidFromClient(client: Client) {
+  return client.did
 }
 
 /*
@@ -183,9 +185,9 @@ export function createServerStateQueryKey({did}: {did: string}) {
   return ['serverState', did]
 }
 export async function getServerState({
-  agent,
+  appviewClient,
 }: {
-  agent: SessionAgent
+  appviewClient: Client
 }): Promise<app.bsky.ageassurance.getState.$OutputBody | null> {
   if (debug.enabled && debug.serverState)
     return debug.resolve(debug.serverState)
@@ -194,24 +196,23 @@ export async function getServerState({
     logger.error(`getServerState: missing geolocation countryCode`)
     return null
   }
-  const {data} = await agent.app.bsky.ageassurance.getState({
+  const data = await appviewClient.call(app.bsky.ageassurance.getState, {
     countryCode: geolocation.countryCode,
     regionCode: geolocation.regionCode,
   })
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(appviewClient)
   if (data && did && createdAtCache.has(did)) {
     /*
      * If account was just created, just use the local cache if available. On
-     * subsequent reloads, the server should have the correct value.
+     * subsequent reloads, the server should have the correct value. The cache
+     * holds ISO datetime strings (written from `new Date().toISOString()`), so
+     * assert the branded DatetimeString at this boundary.
      */
-    data.metadata.accountCreatedAt = createdAtCache.get(did)
+    data.metadata.accountCreatedAt = createdAtCache.get(
+      did,
+    ) as typeof data.metadata.accountCreatedAt
   }
-  /*
+  return data ?? null
-   * TODO(phase4): the bridge agent still returns the old `@atproto/api`
-   * getState output; `toLex` reconciles it with the `#/lexicons` shape at this
-   * single boundary until the agent itself is migrated.
-   */
-  return data ? toLex<app.bsky.ageassurance.getState.$OutputBody>(data) : null
 }
 export function getServerStateFromCache({
   did,
@@ -222,8 +223,12 @@ export function getServerStateFromCache({
     createServerStateQueryKey({did}),
   )
 }
-export async function prefetchServerState({agent}: {agent: SessionAgent}) {
+export async function prefetchServerState({
-  const did = getDidFromAgentSession(agent)
+  appviewClient,
+}: {
+  appviewClient: Client
+}) {
+  const did = getDidFromClient(appviewClient)
 
   if (!did) return
 
@@ -238,7 +243,7 @@ export async function prefetchServerState({agent}: {agent: SessionAgent}) {
 
   try {
     logger.debug(`prefetchServerState: resolving...`)
-    const res = await networkRetry(3, () => getServerState({agent}))
+    const res = await networkRetry(3, () => getServerState({appviewClient}))
     if (res) {
       qc.setQueryData<app.bsky.ageassurance.getState.$OutputBody>(qk, res)
     }
@@ -249,11 +254,15 @@ export async function prefetchServerState({agent}: {agent: SessionAgent}) {
     })
   }
 }
-export async function refetchServerState({agent}: {agent: SessionAgent}) {
+export async function refetchServerState({
-  const did = getDidFromAgentSession(agent)
+  appviewClient,
+}: {
+  appviewClient: Client
+}) {
+  const did = getDidFromClient(appviewClient)
   if (!did) return
   logger.debug(`refetchServerState: fetching...`)
-  const res = await networkRetry(3, () => getServerState({agent}))
+  const res = await networkRetry(3, () => getServerState({appviewClient}))
   if (res) {
     qc.setQueryData<app.bsky.ageassurance.getState.$OutputBody>(
       createServerStateQueryKey({did}),
@@ -283,8 +292,8 @@ export function usePatchServerState() {
   )
 }
 export function useServerStateQuery() {
-  const agent = useAgent()
+  const appviewClient = useAppviewClient()
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(appviewClient)
   const query = useQuery(
     {
       enabled: !!did,
@@ -294,7 +303,7 @@ export function useServerStateQuery() {
       },
       queryKey: createServerStateQueryKey({did: did!}),
       async queryFn() {
-        return getServerState({agent})
+        return getServerState({appviewClient})
       },
     },
     qc,
@@ -346,15 +355,15 @@ export function createOtherRequiredDataQueryKey({did}: {did: string}) {
   return ['otherRequiredData', did]
 }
 async function getOtherRequiredData({
-  agent,
+  accountClient,
 }: {
-  agent: SessionAgent
+  accountClient: Client
 }): Promise<OtherRequiredData> {
   if (debug.enabled) return debug.resolve(debug.otherRequiredData)
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(accountClient)
   const [prefs, actorDeclaration] = await Promise.all([
-    agent.getPreferences(),
+    accountClient.call(getPreferences),
-    fetchActorDeclarationRecord({did, agent}),
+    fetchActorDeclarationRecord({did, client: accountClient}),
   ])
   const data: OtherRequiredData = {
     birthdate: prefs.birthDate ? prefs.birthDate.toISOString() : undefined,
@@ -431,11 +440,11 @@ export function setOtherRequiredDataActorDeclarationCache({
   )
 }
 export async function prefetchOtherRequiredData({
-  agent,
+  accountClient,
 }: {
-  agent: SessionAgent
+  accountClient: Client
 }) {
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(accountClient)
 
   if (!did) return
 
@@ -450,7 +459,9 @@ export async function prefetchOtherRequiredData({
 
   try {
     logger.debug(`prefetchOtherRequiredData: resolving...`)
-    const res = await networkRetry(3, () => getOtherRequiredData({agent}))
+    const res = await networkRetry(3, () =>
+      getOtherRequiredData({accountClient}),
+    )
     qc.setQueryData<OtherRequiredData>(qk, res)
   } catch (err) {
     const e = err as Error
@@ -479,8 +490,8 @@ export function usePatchOtherRequiredData() {
   )
 }
 export function useOtherRequiredDataQuery() {
-  const agent = useAgent()
+  const accountClient = usePdsClient()
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(accountClient)
   return useQuery(
     {
       enabled: !!did,
@@ -490,7 +501,7 @@ export function useOtherRequiredDataQuery() {
       },
       queryKey: createOtherRequiredDataQueryKey({did: did!}),
       async queryFn() {
-        return getOtherRequiredData({agent})
+        return getOtherRequiredData({accountClient})
       },
     },
     qc,
@@ -585,8 +596,12 @@ export function setDeviceSignalsForRegion({
     prev => ({...prev, [regionKey]: signals}),
   )
 }
-export async function prefetchDeviceSignals({agent}: {agent: SessionAgent}) {
+export async function prefetchDeviceSignals({
-  const did = getDidFromAgentSession(agent)
+  appviewClient,
+}: {
+  appviewClient: Client
+}) {
+  const did = getDidFromClient(appviewClient)
   if (!did) return
 
   /**
@@ -621,8 +636,8 @@ export async function prefetchDeviceSignals({agent}: {agent: SessionAgent}) {
    */
 }
 export function useDeviceSignalsQuery() {
-  const agent = useAgent()
+  const appviewClient = useAppviewClient()
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(appviewClient)
   const {data: config} = useConfigQuery()
   const geolocation = useGeolocation()
   /*
@@ -665,15 +680,23 @@ export function useDeviceSignalsQuery() {
 }
 
 /**
- * Helper to prefetch all age assurance data from the server.
+ * Helper to prefetch all age assurance data from the server. Reads that hit
+ * the appview (`getState`, device signals) take the appview client; the
+ * preferences/actor-declaration read hits the PDS via the account client.
  */
-export function prefetchAgeAssuranceServerData({agent}: {agent: SessionAgent}) {
+export function prefetchAgeAssuranceServerData({
+  appviewClient,
+  accountClient,
+}: {
+  appviewClient: Client
+  accountClient: Client
+}) {
   return Promise.allSettled([
     // config fetch initiated at the top of the App.platform.tsx files, awaited here
     configPrefetchPromise,
-    prefetchServerState({agent}),
+    prefetchServerState({appviewClient}),
-    prefetchOtherRequiredData({agent}),
+    prefetchOtherRequiredData({accountClient}),
-    prefetchDeviceSignals({agent}),
+    prefetchDeviceSignals({appviewClient}),
   ])
 }
 

src/ageAssurance/useBeginAgeAssurance.ts

@@ -9,20 +9,20 @@ import {
   PUBLIC_APPVIEW_DID,
 } from '#/lib/constants'
 import {isNetworkError} from '#/lib/hooks/useCleanError'
-import {useAgent} from '#/state/session'
+import {usePdsClient} from '#/state/session'
 import {usePatchAgeAssuranceServerState} from '#/ageAssurance'
 import {logger} from '#/ageAssurance/logger'
 import {useAnalytics} from '#/analytics'
 import {BLUESKY_PROXY_DID} from '#/env'
 import {useGeolocation} from '#/geolocation'
-import {app} from '#/lexicons'
+import {app, com} from '#/lexicons'
 
 const IS_DEV_ENV = BLUESKY_PROXY_DID !== PUBLIC_APPVIEW_DID
 const APPVIEW = IS_DEV_ENV ? DEV_ENV_APPVIEW : PUBLIC_APPVIEW
 
 export function useBeginAgeAssurance() {
   const ax = useAnalytics()
-  const agent = useAgent()
+  const pdsClient = usePdsClient()
   const geolocation = useGeolocation()
   const patchAgeAssuranceStateResponse = usePatchAgeAssuranceServerState()
 
@@ -39,9 +39,7 @@ export function useBeginAgeAssurance() {
         throw new Error(`Geolocation not available, cannot init age assurance.`)
       }
 
-      const {
+      const {token} = await pdsClient.call(com.atproto.server.getServiceAuth, {
-        data: {token},
-      } = await agent.com.atproto.server.getServiceAuth({
         aud: BLUESKY_PROXY_DID,
         lxm: `app.bsky.ageassurance.begin`,
       })

src/components/ageAssurance/AgeAssuranceRedirectDialog.tsx

@@ -6,7 +6,7 @@ import {Trans} from '@lingui/react/macro'
 
 import {retry} from '#/lib/async/retry'
 import {wait} from '#/lib/async/wait'
-import {useAgent} from '#/state/session'
+import {useAppviewClient, useSession} from '#/state/session'
 import {atoms as a, useTheme, web} from '#/alf'
 import {AgeAssuranceBadge} from '#/components/ageAssurance/AgeAssuranceBadge'
 import {Button, ButtonText} from '#/components/Button'
@@ -84,7 +84,8 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
   const t = useTheme()
   const ax = useAnalytics()
   const {_} = useLingui()
-  const agent = useAgent()
+  const {hasSession} = useSession()
+  const appviewClient = useAppviewClient()
   const polling = useRef(false)
   const unmounted = useRef(false)
   const control = useAgeAssuranceRedirectDialogControl()
@@ -104,10 +105,10 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
         5,
         () => true,
         async () => {
-          if (!agent.session) return
+          if (!hasSession) return
           if (unmounted.current) return
 
-          const data = await refetchAgeAssuranceServerState({agent})
+          const data = await refetchAgeAssuranceServerState({appviewClient})
 
           if (data?.state.status !== 'assured') {
             throw new Error(
@@ -122,7 +123,7 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
     )
       .then(async data => {
         if (!data) return
-        if (!agent.session) return
+        if (!hasSession) return
         if (unmounted.current) return
 
         setSuccess(true)
@@ -138,7 +139,7 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
     return () => {
       unmounted.current = true
     }
-  }, [ax, agent, control])
+  }, [ax, hasSession, appviewClient, control])
 
   if (success) {
     return (

src/state/birthdate.ts

@@ -3,7 +3,7 @@ import {useMutation, useQueryClient} from '@tanstack/react-query'
 
 import {restrictChatSettings} from '#/state/queries/messages/restrictChatSettings'
 import {preferencesQueryKey} from '#/state/queries/preferences'
-import {useAgent, useSession} from '#/state/session'
+import {useAgent, usePdsClient, useSession} from '#/state/session'
 import {usePatchAgeAssuranceOtherRequiredData} from '#/ageAssurance'
 import {isUnderAge} from '#/ageAssurance/util'
 import {IS_DEV} from '#/env'
@@ -55,6 +55,7 @@ export function useIsBirthdateUpdateAllowed() {
 export function useBirthdateMutation() {
   const queryClient = useQueryClient()
   const agent = useAgent()
+  const pdsClient = usePdsClient()
   const patchOtherRequiredData = usePatchAgeAssuranceOtherRequiredData()
 
   return useMutation<void, unknown, {birthDate: Date}>({
@@ -68,7 +69,7 @@ export function useBirthdateMutation() {
 
       if (isUnderAge(birthDate.toISOString(), 18)) {
         await restrictChatSettings({
-          agent,
+          client: pdsClient,
           restrictIncoming: true,
           restrictGroupInvites: true,
         })

src/state/queries/messages/actor-declaration.ts

@@ -1,10 +1,9 @@
+import {type Client} from '@atproto/lex-client'
 import {type DidString} from '@atproto/syntax'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
 import {logger} from '#/logger'
 import {usePdsClient, useSession} from '#/state/session'
-import {agentToLexClient} from '#/state/session/clients'
-import {type SessionAgent} from '#/state/session/session-core'
 import {resolveAllowGroupInvites} from '#/components/dms/util'
 import {type app, chat, com} from '#/lexicons'
 import {RQKEY as PROFILE_RKEY} from '../profile'
@@ -116,25 +115,13 @@ export function useDeleteActorDeclaration() {
 }
 
 export async function fetchActorDeclarationRecord({
-  agent,
+  client,
   did,
 }: {
-  agent: SessionAgent
+  client: Client
   did?: string
 }) {
   if (!did) return
-  /*
-   * This helper is called with a bridge `SessionAgent` threaded from
-   * `#/ageAssurance/data`. Wrap it as an account lex `Client` so the record
-   * read goes through the same path as the migrated hooks; the caller keeps
-   * passing the agent until the bridge is removed (Phase 4). The cast is safe:
-   * `agentToLexClient` only reads `did` and `fetchHandler`, both of which the
-   * base `Agent` (and thus `SessionAgent`) provides - its `AtpAgent` parameter
-   * type is just narrower than it needs. TODO(phase4): drop with the bridge.
-   */
-  const client = agentToLexClient(
-    agent as unknown as Parameters<typeof agentToLexClient>[0],
-  )
   const res = await client
     .get(chat.bsky.actor.declaration, {repo: did as DidString, rkey: 'self'})
     .catch(_e => undefined)

src/state/queries/messages/restrictChatSettings.ts

@@ -1,11 +1,9 @@
-import {type DidString} from '@atproto/syntax'
+import {type Client} from '@atproto/lex-client'
 
 import {networkRetry} from '#/lib/async/retry'
 import {logger} from '#/logger'
-import {type SessionAgent} from '#/state/session'
-import {agentToLexClient} from '#/state/session/clients'
 import {
-  getDidFromAgentSession,
+  getDidFromClient,
   getOtherRequiredDataFromCache,
   setOtherRequiredDataActorDeclarationCache,
 } from '#/ageAssurance/data'
@@ -26,15 +24,15 @@ import {chat} from '#/lexicons'
  * back to the lexicon defaults when the cache is empty.
  */
 export async function restrictChatSettings({
-  agent,
+  client,
   restrictIncoming = false,
   restrictGroupInvites = false,
 }: {
-  agent: SessionAgent
+  client: Client
   restrictIncoming?: boolean
   restrictGroupInvites?: boolean
 }): Promise<void> {
-  const did = getDidFromAgentSession(agent)
+  const did = getDidFromClient(client)
   if (!did) return
 
   const cached = getOtherRequiredDataFromCache({did})?.actorDeclaration
@@ -69,22 +67,10 @@ export async function restrictChatSettings({
     return
   }
 
-  /*
-   * Callers thread a bridge `SessionAgent` (session-core / birthdate); wrap it
-   * as an account lex `Client` so the record write goes through the lex path.
-   * The cast is safe: `agentToLexClient` only reads `did` and `fetchHandler`,
-   * both of which the base `Agent` provides - its `AtpAgent` parameter type is
-   * just narrower than it needs. TODO(phase4): take a Client directly once the
-   * bridge is removed.
-   */
-  const client = agentToLexClient(
-    agent as unknown as Parameters<typeof agentToLexClient>[0],
-  )
-
   try {
     await networkRetry(3, () =>
       client.put(chat.bsky.actor.declaration, record, {
-        repo: did as DidString,
+        repo: did,
         rkey: 'self',
       }),
     )

src/state/session/session-core.ts

@@ -619,7 +619,8 @@ export async function createSessionBundleAndResume(
 
   const moderation = configureModerationForAccount(bundle, account)
   const aa = prefetchAgeAssuranceServerData({
-    agent: bundle.agent,
+    appviewClient: bundle.appviewClient,
+    accountClient: bundle.accountClient,
   })
   await Promise.all([gates, moderation, aa])
   hooks.arm()
@@ -670,7 +671,8 @@ export async function createSessionBundleAndLogin(
   const gates = features.refresh({strategy: 'prefer-fresh-gates'})
   const moderation = configureModerationForAccount(bundle, account)
   const aa = prefetchAgeAssuranceServerData({
-    agent: bundle.agent,
+    appviewClient: bundle.appviewClient,
+    accountClient: bundle.accountClient,
   })
   await Promise.all([gates, moderation, aa])
   hooks.arm()
@@ -750,7 +752,10 @@ export async function createSessionBundleAndCreateAccount(
   setBirthdateForDid({did: account.did, birthdate})
   snoozeBirthdateUpdateAllowedForDid(account.did)
   // do this last
-  const aa = prefetchAgeAssuranceServerData({agent})
+  const aa = prefetchAgeAssuranceServerData({
+    appviewClient: bundle.appviewClient,
+    accountClient: bundle.accountClient,
+  })
 
   // Not awaited so that we can still get into onboarding.
   // This is OK because we won't let you toggle adult stuff until you set the date.
@@ -801,7 +806,7 @@ export async function createSessionBundleAndCreateAccount(
         const {flags} = unsafeGetAndComputeAgeAssurance({did: account.did})
         if (flags?.chatDisabled || flags?.groupChatDisabled) {
           void restrictChatSettings({
-            agent,
+            client: bundle.accountClient,
             restrictIncoming: flags.chatDisabled,
             restrictGroupInvites: flags.groupChatDisabled,
           })