fix: preserve paykit cleanup retries

Author: ben-kaufman

app/src/main/java/to/bitkit/repositories/PrivatePaykitRepo.kt

@@ -112,6 +112,16 @@ class PrivatePaykitRepo @Inject constructor(
         val firstError: Throwable?,
     )
 
+    private data class NormalizedPublicKeyBatch(
+        val publicKeys: List<String>,
+        val hadInvalidKey: Boolean,
+    )
+
+    private data class LinkedReceiverPathInspection(
+        val receiverPaths: Set<String>,
+        val error: Throwable?,
+    )
+
     private data class PrivateMessageDrainRetryKey(
         val publicKey: String,
         val receiverPath: String,
@@ -699,7 +709,7 @@ class PrivatePaykitRepo @Inject constructor(
         var firstError: Throwable? = null
         val updates = mutableListOf<PrivatePaymentListReservationUpdateInput>()
         val linkRetryKeys = mutableListOf<PrivateMessageDrainRetryKey>()
-        val linkedReceiverPaths = linkedReceiverPathsByPublicKey()
+        val linkedReceiverPathsSnapshot = linkedReceiverPathsSnapshotOrNull(reason)
 
         for (publicKey in publicKeys) {
             val receiverPaths = runSuspendCatching { receiverPathsForSavedContact(publicKey) }
@@ -718,22 +728,19 @@ class PrivatePaykitRepo @Inject constructor(
                 firstError = firstError ?: it
                 logPrivateReceiverPathSelectionFailure(publicKey, reason, it)
             }
+            val linkedReceiverPathInspection = inspectLinkedReceiverPaths(
+                publicKey,
+                linkedReceiverPathsSnapshot,
+                reason,
+            )
+            firstError = firstError ?: linkedReceiverPathInspection.error
             val cleanupReceiverPaths = receiverPathsForPrivateEndpointCleanup(
                 publicKey = publicKey,
                 excludedReceiverPaths = publicationReceiverPaths + receiverPathSelection.cleanupProtectedReceiverPaths,
-                linkedReceiverPaths = linkedReceiverPaths[publicKey].orEmpty(),
+                linkedReceiverPaths = linkedReceiverPathInspection.receiverPaths,
             )
 
-            (linkableReceiverPaths + cleanupReceiverPaths).distinct().forEach { receiverPath ->
-                linkRetryKeys += PrivateMessageDrainRetryKey(publicKey, receiverPath)
-                runSuspendCatching { paykitSdkService.ensureLinkWithPeer(publicKey, receiverPath) }.onFailure {
-                    Logger.warn(
-                        "Failed to prepare private Paykit link for '${redacted(publicKey)}' during '$reason'",
-                        it,
-                        context = TAG,
-                    )
-                }
-            }
+            linkRetryKeys += preparePrivateLinks(publicKey, linkableReceiverPaths + cleanupReceiverPaths, reason)
 
             cleanupReceiverPaths.forEach { receiverPath ->
                 updates += PrivatePaymentListReservationUpdateInput(
@@ -785,6 +792,21 @@ class PrivatePaykitRepo @Inject constructor(
         drainAndSchedulePrivateLinkRetries(reason, retryKeys.distinct())
     }
 
+    private suspend fun preparePrivateLinks(
+        publicKey: String,
+        receiverPaths: Collection<String>,
+        reason: String,
+    ): List<PrivateMessageDrainRetryKey> = receiverPaths.distinct().map { receiverPath ->
+        runSuspendCatching { paykitSdkService.ensureLinkWithPeer(publicKey, receiverPath) }.onFailure {
+            Logger.warn(
+                "Failed to prepare private Paykit link for '${redacted(publicKey)}' during '$reason'",
+                it,
+                context = TAG,
+            )
+        }
+        PrivateMessageDrainRetryKey(publicKey, receiverPath)
+    }
+
     private suspend fun drainAndSchedulePrivateLinkRetries(
         reason: String,
         retryKeys: Collection<PrivateMessageDrainRetryKey>,
@@ -1207,13 +1229,20 @@ class PrivatePaykitRepo @Inject constructor(
     private suspend fun removePublishedEndpoints(publicKeys: Collection<String>): Result<Unit> =
         withContext(serializedDispatcher) {
             runSuspendCatching {
-                val publicKeys = publicKeys.mapNotNull(::normalizedPublicKey).distinct()
-                if (publicKeys.isEmpty()) return@runSuspendCatching Unit
+                val normalizedBatch = normalizedPublicKeyBatch(publicKeys)
+                val publicKeys = normalizedBatch.publicKeys
+                var firstError: Throwable? = PrivatePaykitError.InvalidPublicKey.takeIf {
+                    normalizedBatch.hadInvalidKey
+                }
+                if (publicKeys.isEmpty()) {
+                    firstError?.let { throw it }
+                    return@runSuspendCatching Unit
+                }
 
-                val linkedReceiverPaths = linkedReceiverPathsByPublicKey()
-                val preparation = clearPrivatePaymentLists(publicKeys, linkedReceiverPaths)
+                val linkedReceiverPathsSnapshot = linkedReceiverPathsSnapshotOrNull("private endpoint cleanup")
+                val preparation = clearPrivatePaymentLists(publicKeys, linkedReceiverPathsSnapshot)
                 val failedPublicKeys = preparation.failedPublicKeys.toMutableSet()
-                var firstError = preparation.firstError
+                firstError = firstError ?: preparation.firstError
 
                 if (preparation.clearedRetryKeys.isNotEmpty()) {
                     drainPendingPrivateMessages(
@@ -1237,16 +1266,28 @@ class PrivatePaykitRepo @Inject constructor(
 
     private suspend fun clearPrivatePaymentLists(
         publicKeys: Collection<String>,
-        linkedReceiverPaths: Map<String, Set<String>>,
+        linkedReceiverPathsSnapshot: Map<String, Set<String>>?,
     ): PrivateEndpointCleanupPreparation {
         val failedPublicKeys = mutableSetOf<String>()
         val clearedRetryKeys = mutableListOf<PrivateMessageDrainRetryKey>()
         var firstError: Throwable? = null
 
         publicKeys.forEach { publicKey ->
+            val contactLinkedReceiverPaths = linkedReceiverPaths(
+                publicKey = publicKey,
+                snapshot = linkedReceiverPathsSnapshot,
+            ).onFailure {
+                failedPublicKeys += publicKey
+                firstError = firstError ?: it
+                Logger.warn(
+                    "Failed to inspect private Paykit links for '${redacted(publicKey)}' during cleanup",
+                    it,
+                    context = TAG,
+                )
+            }.getOrDefault(emptySet())
             receiverPathsForCleanup(
                 publicKey = publicKey,
-                linkedReceiverPaths = linkedReceiverPaths[publicKey].orEmpty(),
+                linkedReceiverPaths = contactLinkedReceiverPaths,
             ).forEach { receiverPath ->
                 runSuspendCatching {
                     val report = paykitSdkService.clearPrivatePaymentList(publicKey, receiverPath)
@@ -1325,6 +1366,53 @@ class PrivatePaykitRepo @Inject constructor(
         return linkedPaths
     }
 
+    private suspend fun linkedReceiverPathsSnapshotOrNull(reason: String): Map<String, Set<String>>? =
+        runSuspendCatching { linkedReceiverPathsByPublicKey() }
+            .onFailure {
+                Logger.warn(
+                    "Failed to inspect private Paykit links during '$reason'; retrying per contact",
+                    it,
+                    context = TAG,
+                )
+            }.getOrNull()
+
+    private suspend fun linkedReceiverPaths(
+        publicKey: String,
+        snapshot: Map<String, Set<String>>?,
+    ): Result<Set<String>> = if (snapshot != null) {
+        Result.success(snapshot[publicKey].orEmpty())
+    } else {
+        runSuspendCatching { linkedReceiverPathsByPublicKey()[publicKey].orEmpty() }
+    }
+
+    private suspend fun inspectLinkedReceiverPaths(
+        publicKey: String,
+        snapshot: Map<String, Set<String>>?,
+        reason: String,
+    ): LinkedReceiverPathInspection {
+        val result = linkedReceiverPaths(publicKey, snapshot)
+        val error = result.exceptionOrNull()
+        if (error != null) {
+            Logger.warn(
+                "Failed to inspect private Paykit links for '${redacted(publicKey)}' during '$reason'",
+                error,
+                context = TAG,
+            )
+        }
+        return LinkedReceiverPathInspection(result.getOrDefault(emptySet()), error)
+    }
+
+    private fun normalizedPublicKeyBatch(publicKeys: Collection<String>): NormalizedPublicKeyBatch {
+        var hadInvalidKey = false
+        val normalizedKeys = publicKeys.mapNotNull { publicKey ->
+            normalizedPublicKey(publicKey) ?: run {
+                hadInvalidKey = true
+                null
+            }
+        }.distinct()
+        return NormalizedPublicKeyBatch(normalizedKeys, hadInvalidKey)
+    }
+
     private fun publishedPrivatePaymentReceiverPaths(publicKey: String): List<String> {
         val contactState = state?.contacts?.get(publicKey) ?: return emptyList()
         return contactState.publishedPrivatePaymentReceiverPaths.toList()

app/src/test/java/to/bitkit/repositories/PrivatePaykitRepoTest.kt

@@ -632,6 +632,51 @@ class PrivatePaykitRepoTest : BaseUnitTest(StandardTestDispatcher()) {
 
         assertTrue(result.isFailure)
         assertTrue(cacheData.value.cleanupPending)
+        verifyBlocking(paykitSdkService) { clearPrivatePaymentList(CONTACT_KEY, WALLET_RECEIVER_PATH) }
+        assertEquals(
+            setOf(WALLET_RECEIVER_PATH),
+            cacheData.value.contacts.getValue(CONTACT_KEY).publishedPrivatePaymentReceiverPaths,
+        )
+    }
+
+    @Test
+    fun `cleanup falls back per contact when shared linked receiver inspection fails`() = test {
+        cacheData.value = PrivatePaykitCacheData(
+            contacts = mapOf(
+                CONTACT_KEY to cachedPublishedContact(WALLET_RECEIVER_PATH),
+                OTHER_CONTACT_KEY to cachedPublishedContact(SERVER_RECEIVER_PATH),
+            ),
+        )
+        sut = createSut()
+        var linkedPeerReads = 0
+        whenever { paykitSdkService.linkedPeers() }.thenAnswer {
+            linkedPeerReads += 1
+            if (linkedPeerReads <= 2) error("link inspection failed")
+            emptyList<LinkedPeerRecord>()
+        }
+
+        val result = sut.removePublishedEndpointsForCleanup("test")
+
+        assertTrue(result.isFailure)
+        verifyBlocking(paykitSdkService) { clearPrivatePaymentList(CONTACT_KEY, WALLET_RECEIVER_PATH) }
+        verifyBlocking(paykitSdkService) { clearPrivatePaymentList(OTHER_CONTACT_KEY, SERVER_RECEIVER_PATH) }
+        assertTrue(CONTACT_KEY in cacheData.value.contacts)
+        assertTrue(OTHER_CONTACT_KEY !in cacheData.value.contacts)
+        assertTrue(cacheData.value.cleanupPending)
+    }
+
+    @Test
+    fun `invalid deleted contact key remains pending for cleanup`() = test {
+        val invalidPublicKey = "not-a-pubky"
+        cacheData.value = PrivatePaykitCacheData(
+            deletedContactCleanupPendingPublicKeys = setOf(invalidPublicKey),
+        )
+        sut = createSut()
+
+        val result = sut.retryPendingEndpointRemoval(emptyList())
+
+        assertTrue(result.isFailure)
+        assertEquals(setOf(invalidPublicKey), cacheData.value.deletedContactCleanupPendingPublicKeys)
         verifyBlocking(paykitSdkService, never()) { clearPrivatePaymentList(any(), any()) }
     }