Ensure background work is flagged as such to Rust SDK

ET-5803

mail-notifications/src/main/kotlin/ch/protonmail/android/mailnotifications/data/usecase/ExecutePushNotificationAction.kt

@@ -25,6 +25,7 @@ import ch.protonmail.android.mailnotifications.data.model.QuickActionPayloadData
 import ch.protonmail.android.mailnotifications.domain.model.LocalNotificationAction
 import ch.protonmail.android.mailsession.data.mapper.toLocalUserId
 import ch.protonmail.android.mailsession.data.repository.MailSessionRepository
+import ch.protonmail.android.mailsession.data.repository.runInRustBackground
 import me.proton.core.domain.entity.UserId
 import timber.log.Timber
 import uniffi.mail_uniffi.MailSession
@@ -39,32 +40,36 @@ internal class ExecutePushNotificationAction @Inject constructor(
     private val mailSessionRepository: MailSessionRepository
 ) {
 
-    suspend operator fun invoke(payload: QuickActionPayloadData) = either {
-        val mailSession = mailSessionRepository.getMailSession().getRustMailSession()
+    suspend operator fun invoke(payload: QuickActionPayloadData) =
+        mailSessionRepository.runInRustBackground { mailSession ->
+            either {
+                val mailSession = mailSession.getRustMailSession()
 
-        val remoteId = RemoteId(payload.remoteId)
+                val remoteId = RemoteId(payload.remoteId)
 
-        val action = when (payload.action) {
-            LocalNotificationAction.MarkAsRead -> PushNotificationQuickAction.MarkAsRead(remoteId)
-            LocalNotificationAction.MoveTo.Archive -> PushNotificationQuickAction.MoveToArchive(remoteId)
-            LocalNotificationAction.MoveTo.Trash -> PushNotificationQuickAction.MoveToTrash(remoteId)
+                val action = when (payload.action) {
+                    LocalNotificationAction.MarkAsRead -> PushNotificationQuickAction.MarkAsRead(remoteId)
+                    LocalNotificationAction.MoveTo.Archive -> PushNotificationQuickAction.MoveToArchive(remoteId)
+                    LocalNotificationAction.MoveTo.Trash -> PushNotificationQuickAction.MoveToTrash(remoteId)
+                }
+
+                val session = getStoredSessionForUser(mailSession, payload.userId)
+                    ?: raise(QuickActionPushError.NoMailSession)
+
+                when (val result = mailSession.executeNotificationQuickAction(session, action, timeLeftMs = null)) {
+                    is VoidActionResult.Error -> raise(QuickActionPushError.Error(result.v1.toDataError()))
+                    VoidActionResult.Ok -> Unit
+                }
+            }
         }
 
-        val session = getStoredSessionForUser(mailSession, payload.userId)
-            ?: raise(QuickActionPushError.NoMailSession)
-
-        when (val result = mailSession.executeNotificationQuickAction(session, action, timeLeftMs = null)) {
-            is VoidActionResult.Error -> raise(QuickActionPushError.Error(result.v1.toDataError()))
-            VoidActionResult.Ok -> Unit
-        }
-    }
-
     private suspend fun getStoredSessionForUser(mailSession: MailSession, userId: UserId): StoredSession? =
         when (val result = mailSession.getSessions()) {
             is MailSessionGetSessionsResult.Error -> {
                 Timber.e("execute-push-action failed due to no user session. Error: $result")
                 null
             }
+
             is MailSessionGetSessionsResult.Ok -> result.v1.firstOrNull { it.userId() == userId.toLocalUserId() }
         }
 }

mail-notifications/src/main/kotlin/ch/protonmail/android/mailnotifications/domain/usecase/content/HandlePushNotification.kt

@@ -19,43 +19,55 @@
 package ch.protonmail.android.mailnotifications.domain.usecase.content
 
 import ch.protonmail.android.mailcommon.data.worker.Enqueuer
-import ch.protonmail.android.mailnotifications.data.local.ProcessPushNotificationDataWorker
 import ch.protonmail.android.mailfeatureflags.domain.annotation.IsPushProcessingWithoutWorkerEnabled
 import ch.protonmail.android.mailfeatureflags.domain.model.FeatureFlag
+import ch.protonmail.android.mailnotifications.data.local.ProcessPushNotificationDataWorker
 import ch.protonmail.android.mailnotifications.domain.usecase.ProcessPushNotification
-import ch.protonmail.android.mailsession.domain.repository.UserSessionRepository
+import ch.protonmail.android.mailsession.data.mapper.toUserId
+import ch.protonmail.android.mailsession.data.repository.MailSessionRepository
+import ch.protonmail.android.mailsession.data.repository.runInRustBackground
+import ch.protonmail.android.mailsession.data.wrapper.MailSessionWrapper
+import me.proton.core.domain.entity.UserId
 import me.proton.core.network.domain.session.SessionId
 import timber.log.Timber
 import javax.inject.Inject
 
 internal class HandlePushNotification @Inject constructor(
     private val enqueuer: Enqueuer,
-    private val userSessionRepository: UserSessionRepository,
     private val processPushNotification: ProcessPushNotification,
+    private val mailSessionRepository: MailSessionRepository,
     @IsPushProcessingWithoutWorkerEnabled private val processingWithoutWorkerEnabled: FeatureFlag<Boolean>
 ) {
 
     suspend operator fun invoke(sessionId: SessionId, encryptedMessage: String) {
-        val userId = userSessionRepository.getUserId(sessionId) ?: run {
-            Timber.w(
-                "Notification: Push received but userId could not be resolved for sessionId=%s",
-                sessionId.id
-            )
-            return
-        }
+        mailSessionRepository.runInRustBackground { mailSession ->
+            val userId = resolveUserId(mailSession, sessionId) ?: run {
+                Timber.w(
+                    "Notification: Push received but userId could not be resolved for sessionId=%s",
+                    sessionId.id
+                )
+                return@runInRustBackground
+            }
 
-        if (processingWithoutWorkerEnabled.get()) {
-            processPushNotification(
-                userId,
-                sessionId,
-                encryptedMessage
-            )
-        } else {
-            enqueuer.enqueue<ProcessPushNotificationDataWorker>(
-                userId,
-                ProcessPushNotificationDataWorker.params(userId.id, sessionId.id, encryptedMessage)
-            )
+            Timber.d("Notification: Resolved userId for the given session")
 
+            if (processingWithoutWorkerEnabled.get()) {
+                processPushNotification(
+                    userId,
+                    sessionId,
+                    encryptedMessage
+                )
+            } else {
+                enqueuer.enqueue<ProcessPushNotificationDataWorker>(
+                    userId,
+                    ProcessPushNotificationDataWorker.params(userId.id, sessionId.id, encryptedMessage)
+                )
+            }
         }
     }
+
+    private suspend fun resolveUserId(mailSession: MailSessionWrapper, sessionId: SessionId): UserId? =
+        mailSession.getSessionById(sessionId).getOrNull()
+            ?.userId()
+            ?.toUserId()
 }

mail-notifications/src/test/kotlin/ch/protonmail/android/mailnotifications/data/usecase/ExecutePushNotificationActionTest.kt

@@ -22,15 +22,19 @@ import ch.protonmail.android.mailnotifications.data.model.QuickActionPayloadData
 import ch.protonmail.android.mailnotifications.domain.model.LocalNotificationAction
 import ch.protonmail.android.mailsession.data.mapper.toLocalUserId
 import ch.protonmail.android.mailsession.data.repository.MailSessionRepository
+import ch.protonmail.android.mailsession.data.wrapper.MailSessionWrapper
 import io.mockk.clearAllMocks
 import io.mockk.coEvery
 import io.mockk.coVerify
 import io.mockk.confirmVerified
+import io.mockk.every
+import io.mockk.justRun
 import io.mockk.mockk
 import kotlinx.coroutines.test.runTest
 import me.proton.core.domain.entity.UserId
 import uniffi.mail_uniffi.ActionError
 import uniffi.mail_uniffi.ActionErrorReason
+import uniffi.mail_uniffi.MailBackgroundExecScope
 import uniffi.mail_uniffi.MailSession
 import uniffi.mail_uniffi.MailSessionGetSessionsResult
 import uniffi.mail_uniffi.PushNotificationQuickAction
@@ -54,8 +58,15 @@ internal class ExecutePushNotificationActionTest {
             MailSessionGetSessionsResult.Ok(listOf(storedSession))
     }
 
+    private val backgroundScope = mockk<MailBackgroundExecScope> {
+        justRun { finsihed() }
+    }
+    private val mailSessionWrapper = mockk<MailSessionWrapper> {
+        every { getRustMailSession() } returns mailSession
+        every { newBackgroundExecutionScope() } returns backgroundScope
+    }
     private val mailSessionRepo = mockk<MailSessionRepository> {
-        coEvery { this@mockk.getMailSession().getRustMailSession() } returns mailSession
+        every { getMailSession() } returns mailSessionWrapper
     }
 
     private lateinit var executePushNotificationAction: ExecutePushNotificationAction

mail-session/data/src/main/kotlin/ch/protonmail/android/mailsession/data/repository/MailSessionRepositoryExtensions.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2022 Proton Technologies AG
+ * This file is part of Proton Technologies AG and Proton Mail.
+ *
+ * Proton Mail is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Proton Mail is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Proton Mail. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package ch.protonmail.android.mailsession.data.repository
+
+import ch.protonmail.android.mailsession.data.wrapper.MailSessionWrapper
+
+/**
+ * Runs the block within a Rust background execution scope to ensure that calls are able to execute while
+ * the app is in the background without stalling. The Rust scope teardown is handled internally.
+ */
+suspend inline fun <T> MailSessionRepository.runInRustBackground(block: suspend (session: MailSessionWrapper) -> T): T {
+    val session = getMailSession()
+    val scope = session.newBackgroundExecutionScope()
+    return try {
+        block(session)
+    } finally {
+        scope.finsihed()
+    }
+}

mail-session/data/src/main/kotlin/ch/protonmail/android/mailsession/data/wrapper/MailSessionWrapper.kt

@@ -27,12 +27,14 @@ import ch.protonmail.android.mailcommon.data.mapper.toDataError
 import ch.protonmail.android.mailcommon.domain.model.DataError
 import ch.protonmail.android.mailsession.data.mapper.toAutoLockPinError
 import ch.protonmail.android.mailsession.domain.wrapper.MailUserSessionWrapper
+import me.proton.core.network.domain.session.SessionId
 import uniffi.mail_uniffi.BackgroundExecutionCallback
 import uniffi.mail_uniffi.MailSession
 import uniffi.mail_uniffi.MailSessionDeletePinCodeResult
 import uniffi.mail_uniffi.MailSessionGetAccountResult
 import uniffi.mail_uniffi.MailSessionGetAccountSessionsResult
 import uniffi.mail_uniffi.MailSessionGetPrimaryAccountResult
+import uniffi.mail_uniffi.MailSessionGetSessionResult
 import uniffi.mail_uniffi.MailSessionGetSessionsResult
 import uniffi.mail_uniffi.MailSessionInitializedUserSessionFromStoredSessionResult
 import uniffi.mail_uniffi.MailSessionNewLoginFlowResult
@@ -47,6 +49,7 @@ import uniffi.mail_uniffi.MeasurementValue
 import uniffi.mail_uniffi.StoredAccount
 import uniffi.mail_uniffi.StoredSession
 
+@Suppress("TooManyFunctions")
 class MailSessionWrapper(private val mailSession: MailSession) {
 
     fun getRustMailSession() = mailSession
@@ -84,6 +87,16 @@ class MailSessionWrapper(private val mailSession: MailSession) {
         is MailSessionGetSessionsResult.Ok -> result.v1.right()
     }
 
+    suspend fun getSessionById(sessionId: SessionId): Either<DataError, StoredSession> = when (
+        val result = mailSession.getSession(sessionId.id)
+    ) {
+        is MailSessionGetSessionResult.Error -> result.v1.toDataError().left()
+        is MailSessionGetSessionResult.Ok -> when (val data = result.v1) {
+            null -> DataError.Local.NotFound.left()
+            else -> data.right()
+        }
+    }
+
     suspend fun userContextFromSession(session: StoredSession): Either<DataError, MailUserSessionWrapper> =
         when (val result = mailSession.userSessionFromStoredSession(session)) {
             is MailSessionUserSessionFromStoredSessionResult.Error -> result.v1.toDataError().left()
@@ -169,4 +182,6 @@ class MailSessionWrapper(private val mailSession: MailSession) {
             is MailSessionNewLoginFlowResult.Error -> result.v1.toDataError().left()
         }
     }
+
+    fun newBackgroundExecutionScope() = mailSession.newBackgroundExecutionScope()
 }

mail-session/data/src/test/kotlin/ch/protonmail/android/mailsession/data/repository/MailSessionRepositoryExtensionsTest.kt

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2026 Proton Technologies AG
+ * This file is part of Proton Technologies AG and Proton Mail.
+ *
+ * Proton Mail is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Proton Mail is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Proton Mail. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package ch.protonmail.android.mailsession.data.repository
+
+import ch.protonmail.android.mailsession.data.wrapper.MailSessionWrapper
+import io.mockk.every
+import io.mockk.justRun
+import io.mockk.mockk
+import io.mockk.verify
+import kotlinx.coroutines.test.runTest
+import uniffi.mail_uniffi.MailBackgroundExecScope
+import uniffi.mail_uniffi.MailSession
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+
+internal class MailSessionRepositoryExtensionsTest {
+
+    private val scope = mockk<MailBackgroundExecScope>()
+    private val rustMailSession = mockk<MailSession> {
+        every { newBackgroundExecutionScope() } returns scope
+    }
+    private val wrapper = mockk<MailSessionWrapper> {
+        every { getRustMailSession() } returns rustMailSession
+        every { newBackgroundExecutionScope() } returns scope
+    }
+    private val mailSessionRepository = mockk<MailSessionRepository> {
+        every { getMailSession() } returns wrapper
+    }
+
+    @Test
+    fun `calls finishes after block completes successfully`() = runTest {
+        // Given
+        justRun { scope.finsihed() }
+
+        // When
+        val result = mailSessionRepository.runInRustBackground { "ok" }
+
+        // Then
+        assertEquals("ok", result)
+        verify(exactly = 1) { scope.finsihed() }
+    }
+
+    @Test
+    fun `calls finishes when block throws`() = runTest {
+        // Given
+        justRun { scope.finsihed() }
+        val boom = IllegalStateException("boom")
+
+        // When
+        val thrown = assertFailsWith<IllegalStateException> {
+            mailSessionRepository.runInRustBackground { throw boom }
+        }
+
+        // Then
+        assertEquals(boom, thrown)
+        verify(exactly = 1) { scope.finsihed() }
+    }
+
+    @Test
+    fun `calls finishes exactly once even when block does early return`() = runTest {
+        // Given
+        justRun { scope.finsihed() }
+
+        // When
+        val result = earlyReturning()
+
+        // Then
+        assertEquals(Unit, result)
+        verify(exactly = 1) { scope.finsihed() }
+    }
+
+    private suspend fun earlyReturning() = mailSessionRepository.runInRustBackground {
+        if (System.currentTimeMillis() > 0) return@runInRustBackground
+        error("unreachable")
+    }
+}