feat: Dedupe and invalidate in-flight wallet address fetches

src/hooks/useWalletManager.ts

@@ -14,6 +14,7 @@
 
 import { produce } from 'immer'
 import { useMemo, useCallback } from 'react'
+import { AddressService } from '../services/addressService'
 import { WalletSetupService } from '../services/walletSetupService'
 import { WorkletLifecycleService } from '../services/workletLifecycleService'
 import {
@@ -363,6 +364,10 @@ export function useWalletManager(): UseWalletManagerResult {
       }
 
       try {
+        // Invalidate before clearing store so in-flight getAddress calls
+        // skip stale writes under this walletId.
+        AddressService.invalidateWallet(walletId)
+
         await WalletSetupService.deleteWallet(walletId)
 
         walletStore.setState((prev) =>

src/services/addressService.ts

@@ -38,6 +38,71 @@ import { AddressInfoResult } from '../types'
  * Provides methods for retrieving and caching wallet addresses.
  */
 export class AddressService {
+  /**
+   * Dedupes concurrent getAddress calls: walletId:network:accountIndex -> Promise
+   */
+  private static inflight = new Map<string, Promise<string>>()
+
+  /**
+   * Per-wallet epoch; bumped on invalidate so in-flight writes are skipped
+   */
+  private static walletEpochs = new Map<string, number>()
+
+  /**
+   * Global epoch; bumped on invalidateAll (e.g. lock/reset)
+   */
+  private static globalEpoch = 0
+
+  /**
+   * Generate a cache key for a specific wallet, network, and account index
+   */
+  private static cacheKey(
+    walletId: string,
+    network: string,
+    accountIndex: number,
+  ): string {
+    return `${walletId}:${network}:${accountIndex}`
+  }
+
+  /**
+   * Get the epoch for a specific wallet
+   */
+  private static getEpoch(walletId: string): number {
+    return this.globalEpoch + (this.walletEpochs.get(walletId) ?? 0)
+  }
+
+  /**
+   * Clear the inflight map for a specific wallet
+   */
+  private static clearInflightForWallet(walletId: string): void {
+    const prefix = `${walletId}:`
+    for (const key of [...this.inflight.keys()]) {
+      if (key.startsWith(prefix)) {
+        this.inflight.delete(key)
+      }
+    }
+  }
+
+  /**
+   * Invalidate in-flight address fetches for a wallet.
+   * Clears the promise cache for that wallet and bumps its epoch so any
+   * still-running fetch will not write to the store when it completes.
+   */
+  static invalidateWallet(walletId: string): void {
+    this.clearInflightForWallet(walletId)
+    this.walletEpochs.set(walletId, (this.walletEpochs.get(walletId) ?? 0) + 1)
+    log('[AddressService] Invalidated wallet address fetches', { walletId })
+  }
+
+  /**
+   * Invalidate all in-flight address fetches (e.g. on lock/reset).
+   */
+  static invalidateAll(): void {
+    this.inflight.clear()
+    this.globalEpoch += 1
+    log('[AddressService] Invalidated all address fetches')
+  }
+
   /**
    * Get address for a specific network and account index
    * Caches the address in walletStore for future use
@@ -64,17 +129,48 @@ export class AddressService {
       return cachedAddress
     }
 
+    const cacheKey = this.cacheKey(targetWalletId, network, accountIndex)
+    const existing = this.inflight.get(cacheKey)
+    if (existing) {
+      return existing
+    }
+
+    const promise = this.fetchAndCacheAddress(
+      network,
+      accountIndex,
+      targetWalletId,
+    )
+    this.inflight.set(cacheKey, promise)
+
+    try {
+      return await promise
+    } finally {
+      if (this.inflight.get(cacheKey) === promise) {
+        this.inflight.delete(cacheKey)
+      }
+    }
+  }
+
+  private static async fetchAndCacheAddress(
+    network: string,
+    accountIndex: number,
+    targetWalletId: string,
+  ): Promise<string> {
+    const walletStore = getWalletStore()
+    const epoch = this.getEpoch(targetWalletId)
     const hrpc = await requireInitialized()
 
     const loadingKey = `${network}-${accountIndex}`
 
     try {
-      walletStore.setState((prev) =>
-        produce(prev, (state) => {
-          state.walletLoading[targetWalletId] ??= {}
-          state.walletLoading[targetWalletId][loadingKey] = true
-        }),
-      )
+      if (this.getEpoch(targetWalletId) === epoch) {
+        walletStore.setState((prev) =>
+          produce(prev, (state) => {
+            state.walletLoading[targetWalletId] ??= {}
+            state.walletLoading[targetWalletId][loadingKey] = true
+          }),
+        )
+      }
 
       const response = await hrpc.callMethod({
         methodName: 'getAddress',
@@ -101,6 +197,15 @@ export class AddressService {
         )
       }
 
+      // Skip store writes if wallet was invalidated while this fetch was in flight
+      if (this.getEpoch(targetWalletId) !== epoch) {
+        log(
+          '[AddressService] Skipping stale address write after invalidation',
+          { walletId: targetWalletId, network, accountIndex },
+        )
+        return address
+      }
+
       // Cache the address using helper (per-wallet)
       walletStore.setState((prev) =>
         produce(
@@ -120,12 +225,14 @@ export class AddressService {
 
       return address
     } catch (error) {
-      walletStore.setState((prev) =>
-        produce(prev, (state) => {
-          state.walletLoading[targetWalletId] ??= {}
-          state.walletLoading[targetWalletId][loadingKey] = false
-        }),
-      )
+      if (this.getEpoch(targetWalletId) === epoch) {
+        walletStore.setState((prev) =>
+          produce(prev, (state) => {
+            state.walletLoading[targetWalletId] ??= {}
+            state.walletLoading[targetWalletId][loadingKey] = false
+          }),
+        )
+      }
 
       handleServiceError(error, 'AddressService', 'getAddress', {
         network,

src/services/workletLifecycleService.ts

@@ -21,6 +21,7 @@
 
 import { Worklet } from 'react-native-bare-kit'
 
+import { AddressService } from './addressService'
 import { getWalletStore } from '../store/walletStore'
 import { getWorkletStore } from '../store/workletStore'
 import { DEFAULT_MNEMONIC_WORD_COUNT } from '../utils/constants'
@@ -532,6 +533,9 @@ export class WorkletLifecycleService {
 
     workletStore.setState({ isWorkletInitializedPromise: createResolvablePromise<boolean>() })
 
+    // Drop in-flight address fetches so late responses cannot repopulate the store
+    AddressService.invalidateAll()
+
     // Clear addresses from wallet store
     walletStore.setState({
       addresses: {},

tests/services/addressService.test.ts

@@ -39,6 +39,9 @@ describe('AddressService', () => {
   beforeEach(() => {
     jest.clearAllMocks()
 
+    // Clear static inflight map / bump global epoch so tests don't leak state
+    AddressService.invalidateAll()
+
     // Setup mock HRPC
     mockHRPC = {
       callMethod: jest.fn(),
@@ -323,6 +326,81 @@ describe('AddressService', () => {
       expect(loadingWasSetToTrue).toBe(true)
       expect(loadingWasSetToFalse).toBe(true)
     })
+
+    it('should dedupe concurrent calls for the same key', async () => {
+      const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'
+      mockHRPC.callMethod.mockResolvedValue({
+        result: JSON.stringify(address),
+      })
+
+      const p1 = AddressService.getAddress('ethereum', 0)
+      const p2 = AddressService.getAddress('ethereum', 0)
+      const results = await Promise.all([p1, p2])
+
+      expect(mockHRPC.callMethod).toHaveBeenCalledTimes(1)
+      expect(results).toEqual([address, address])
+    })
+
+    it('should not dedupe different keys', async () => {
+      mockHRPC.callMethod.mockResolvedValue({
+        result: JSON.stringify('0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'),
+      })
+
+      await Promise.all([
+        AddressService.getAddress('ethereum', 0),
+        AddressService.getAddress('ethereum', 1),
+        AddressService.getAddress('polygon', 0),
+      ])
+
+      expect(mockHRPC.callMethod).toHaveBeenCalledTimes(3)
+    })
+
+    it('should skip store writes after invalidateWallet', async () => {
+      const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'
+      mockHRPC.callMethod.mockResolvedValue({
+        result: JSON.stringify(address),
+      })
+
+      const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1')
+      AddressService.invalidateWallet('test-wallet-1')
+      const result = await pending
+
+      expect(result).toBe(address)
+      expect(mockWalletStore.setState).not.toHaveBeenCalled()
+    })
+
+    it('should skip store writes after invalidateAll', async () => {
+      const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'
+      mockHRPC.callMethod.mockResolvedValue({
+        result: JSON.stringify(address),
+      })
+
+      const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1')
+      AddressService.invalidateAll()
+      const result = await pending
+
+      expect(result).toBe(address)
+      expect(mockWalletStore.setState).not.toHaveBeenCalled()
+    })
+
+    it('should start a fresh fetch after invalidate clears inflight', async () => {
+      const staleAddress = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'
+      const freshAddress = '0x842d35Cc6634C0532925a3b844Bc9e7595f0bEb0'
+
+      mockHRPC.callMethod
+        .mockResolvedValueOnce({ result: JSON.stringify(staleAddress) })
+        .mockResolvedValueOnce({ result: JSON.stringify(freshAddress) })
+
+      const stale = AddressService.getAddress('ethereum', 0, 'test-wallet-1')
+      AddressService.invalidateWallet('test-wallet-1')
+
+      const fresh = AddressService.getAddress('ethereum', 0, 'test-wallet-1')
+      const [staleResult, freshResult] = await Promise.all([stale, fresh])
+
+      expect(mockHRPC.callMethod).toHaveBeenCalledTimes(2)
+      expect(staleResult).toBe(staleAddress)
+      expect(freshResult).toBe(freshAddress)  
+    })
   })
 })