From c0714be0ec416a496ae08b880ca1ddeed8a92a3b Mon Sep 17 00:00:00 2001
From: Eric Bailey <git@esb.lol>
Date: Tue, 30 Jun 2026 10:22:26 -0500
Subject: [PATCH] Fix region key symmetry

---
 src/ageAssurance/data.tsx | 54 ++++++++++++++++++---------------------
 src/ageAssurance/state.ts | 17 ++++++++++--
 src/ageAssurance/util.ts  |  4 +--
 3 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/src/ageAssurance/data.tsx b/src/ageAssurance/data.tsx
index fcd9bce7a1..03c724e322 100644
--- a/src/ageAssurance/data.tsx
+++ b/src/ageAssurance/data.tsx
@@ -35,6 +35,7 @@ import {
   isLegacyBirthdateBug,
 } from '#/ageAssurance/util'
 import {IS_DEV, IS_NATIVE} from '#/env'
+import {useGeolocation} from '#/geolocation'
 import {device} from '#/storage'
 
 /**
@@ -523,7 +524,7 @@ export async function getDeviceSignals(): Promise<
 /**
  * The raw region-keyed map of device signals (all regions). Used internally by
  * the query + writer, which operate on the full map. Most consumers want
- * {@link getDeviceSignalsFromCacheForCurrentRegion}, which resolves to the
+ * {@link getDeviceSignalsFromCacheForRegion}, which resolves to the
  * current region.
  */
 export function getDeviceSignalsMapFromCache({
@@ -535,39 +536,19 @@ export function getDeviceSignalsMapFromCache({
     createDeviceSignalsQueryKey({did}),
   )
 }
-/**
- * Resolves a region-keyed signals map down to the signals for the region the
- * user is currently in (per `mergedGeolocation`). Returns undefined when we
- * have no map, no geolocation, or no stored signals for that region.
- *
- * Device assurance is region-bound, so we only ever surface the signals
- * captured in the user's current region.
- */
-function selectDeviceSignalsForCurrentRegion(
-  map: AgeAssuranceDeviceSignals | undefined,
-): AgeRange.AgeRangeResponse | undefined {
-  if (!map) return undefined
-  const geolocation = device.get(['mergedGeolocation'])
-  if (!geolocation?.countryCode) return undefined
-  return map[
-    createRegionKey({
-      countryCode: geolocation.countryCode,
-      regionCode: geolocation.regionCode,
-    })
-  ]
-}
 /**
  * Returns the device signals for the region the user is currently in, or
- * undefined. See {@link selectDeviceSignalsForCurrentRegion}.
+ * undefined.
  */
-export function getDeviceSignalsFromCacheForCurrentRegion({
+export function getDeviceSignalsFromCacheForRegion({
   did,
+  region,
 }: {
   did: string
+  region: AppBskyAgeassuranceDefs.ConfigRegion
 }): AgeRange.AgeRangeResponse | undefined {
-  return selectDeviceSignalsForCurrentRegion(
-    getDeviceSignalsMapFromCache({did}),
-  )
+  const regionKey = createRegionKey(region)
+  return getDeviceSignalsMapFromCache({did})?.[regionKey]
 }
 /**
  * Stores freshly granted device signals into the (persisted) cache under the
@@ -632,6 +613,21 @@ export async function prefetchDeviceSignals({agent}: {agent: AtpAgent}) {
 export function useDeviceSignalsQuery() {
   const agent = useAgent()
   const did = getDidFromAgentSession(agent)
+  const {data: config} = useConfigQuery()
+  const geolocation = useGeolocation()
+  /*
+   * Resolve the matched config region (no fallback) and key off it, so the read
+   * stays symmetric with the write (see `setDeviceSignalsForRegion`). When
+   * geolocation matches no AA region there's no device grant to surface.
+   */
+  const regionConfig = config
+    ? getAgeAssuranceRegionConfig(config, {
+        countryCode: geolocation.countryCode ?? '',
+        regionCode: geolocation.regionCode,
+      })
+    : undefined
+  const regionKey = regionConfig ? createRegionKey(regionConfig) : undefined
+
   return useQuery(
     {
       /**
@@ -655,7 +651,7 @@ export function useDeviceSignalsQuery() {
       // The cache holds the full region-keyed map (the writer merges into it);
       // `select` resolves it to the current region for consumers without
       // mutating the cached value.
-      select: selectDeviceSignalsForCurrentRegion,
+      select: map => (map && regionKey ? map[regionKey] : undefined),
     },
     qc,
   )
@@ -705,7 +701,7 @@ export type AgeAssuranceServerData = {
   /**
    * The native on-device age signals for the region the user is currently in,
    * if they've granted access there. Already resolved from the region-keyed
-   * cache (see `getDeviceSignalsFromCacheForCurrentRegion`), so a grant from
+   * cache (see `getDeviceSignalsFromCacheForRegion`), so a grant from
    * another region won't appear here. Only consumed for regions that permit
    * device verification.
    */
diff --git a/src/ageAssurance/state.ts b/src/ageAssurance/state.ts
index 677937b781..457ff8275a 100644
--- a/src/ageAssurance/state.ts
+++ b/src/ageAssurance/state.ts
@@ -3,13 +3,14 @@ import type * as AgeRange from 'expo-age-range'
 import {
   type AppBskyAgeassuranceDefs,
   computeAgeAssuranceRegionAccess,
+  getAgeAssuranceRegionConfig,
 } from '@atproto/api'
 
 import {getAge} from '#/lib/strings/time'
 import {useSession} from '#/state/session'
 import {
   getConfigFromCache,
-  getDeviceSignalsFromCacheForCurrentRegion,
+  getDeviceSignalsFromCacheForRegion,
   getOtherRequiredDataFromCache,
   getServerStateFromCache,
   useAgeAssuranceServerDataContext,
@@ -140,7 +141,6 @@ export function unsafeGetAndComputeAgeAssurance({did}: {did: string}) {
   const config = getConfigFromCache()
   const state = getServerStateFromCache({did})
   const requiredData = getOtherRequiredDataFromCache({did})
-  const deviceSignals = getDeviceSignalsFromCacheForCurrentRegion({did})
   const geolocation = device.get(['mergedGeolocation'])
 
   if (!geolocation || !config || !state || !requiredData) {
@@ -153,6 +153,19 @@ export function unsafeGetAndComputeAgeAssurance({did}: {did: string}) {
   }
 
   const region = getAgeAssuranceRegionConfigWithFallback(config, geolocation)
+  /*
+   * Device signals are keyed off the matched config region (no fallback): if
+   * geolocation matches no AA region there's no device grant to read, so we
+   * skip the lookup rather than keying off FALLBACK_REGION_CONFIG. This keeps
+   * the read key symmetric with the write (see `setDeviceSignalsForRegion`).
+   */
+  const deviceRegion = getAgeAssuranceRegionConfig(config, {
+    countryCode: geolocation.countryCode ?? '',
+    regionCode: geolocation.regionCode,
+  })
+  const deviceSignals = deviceRegion
+    ? getDeviceSignalsFromCacheForRegion({did, region: deviceRegion})
+    : undefined
   const metadata: AgeAssuranceMetadata = {
     accountCreatedAt: state.metadata?.accountCreatedAt,
     declaredAge: requiredData?.birthdate
diff --git a/src/ageAssurance/util.ts b/src/ageAssurance/util.ts
index 2c5ff860ba..f2768860c2 100644
--- a/src/ageAssurance/util.ts
+++ b/src/ageAssurance/util.ts
@@ -79,8 +79,8 @@ export function createRegionKey(region: {
  * Derives age assurance data from native device signals for the given region,
  * but only when the region permits device verification. The signals are
  * expected to already be resolved to the user's current region (see
- * `getDeviceSignalsFromCacheForCurrentRegion`), so a grant captured in another
- * region won't reach here.
+ * `getDeviceSignalsFromCacheForRegion`), so a grant captured in another region
+ * won't reach here.
  *
  * The OS-provided `lowerBound` is the minimum age the platform will attest to,
  * which maps onto the `assuredAge` input of the rule engine (i.e.