From fbca1b3dfec579eeb495940452949cfd6ca21cc7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Breno?=
 <joao.breno.rodrigues07@aluno.ifce.edu.br>
Date: Mon, 1 Sep 2025 10:07:32 -0300
Subject: [PATCH] Update state in runUpdatingState when CancellationException
 occurs (#5243)

---
 .../libraries/architecture/AsyncAction.kt     | 12 ++++--
 .../libraries/architecture/AsyncActionTest.kt | 42 +++++++++++++++++++
 .../matrix/impl/room/RoomContentForwarder.kt  |  4 --
 3 files changed, 50 insertions(+), 8 deletions(-)
 create mode 100644 libraries/architecture/src/test/kotlin/io/element/android/libraries/architecture/AsyncActionTest.kt

diff --git a/libraries/architecture/src/main/kotlin/io/element/android/libraries/architecture/AsyncAction.kt b/libraries/architecture/src/main/kotlin/io/element/android/libraries/architecture/AsyncAction.kt
index f080b64d01..c02c64135d 100644
--- a/libraries/architecture/src/main/kotlin/io/element/android/libraries/architecture/AsyncAction.kt
+++ b/libraries/architecture/src/main/kotlin/io/element/android/libraries/architecture/AsyncAction.kt
@@ -10,6 +10,7 @@ package io.element.android.libraries.architecture
 import androidx.compose.runtime.MutableState
 import androidx.compose.runtime.Stable
 import io.element.android.libraries.core.extensions.runCatchingExceptions
+import kotlinx.coroutines.TimeoutCancellationException
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.InvocationKind
 import kotlin.contracts.contract
@@ -159,16 +160,19 @@ suspend inline fun <T> runUpdatingState(
         callsInPlace(resultBlock, InvocationKind.EXACTLY_ONCE)
     }
     state.value = AsyncAction.Loading
-    return resultBlock().fold(
+    return try {
+        resultBlock()
+    } catch (e: TimeoutCancellationException) {
+        state.value = AsyncAction.Failure(errorTransform(e))
+        throw e
+    }.fold(
         onSuccess = {
             state.value = AsyncAction.Success(it)
             Result.success(it)
         },
         onFailure = {
             val error = errorTransform(it)
-            state.value = AsyncAction.Failure(
-                error = error,
-            )
+            state.value = AsyncAction.Failure(error)
             Result.failure(error)
         }
     )
diff --git a/libraries/architecture/src/test/kotlin/io/element/android/libraries/architecture/AsyncActionTest.kt b/libraries/architecture/src/test/kotlin/io/element/android/libraries/architecture/AsyncActionTest.kt
new file mode 100644
index 0000000000..fbf892a13f
--- /dev/null
+++ b/libraries/architecture/src/test/kotlin/io/element/android/libraries/architecture/AsyncActionTest.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2025 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.libraries.architecture
+
+import androidx.compose.runtime.MutableState
+import androidx.compose.runtime.mutableStateOf
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withTimeout
+import org.junit.Assert.assertSame
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Test
+import kotlin.time.Duration.Companion.milliseconds
+
+class AsyncActionTest {
+    @Test
+    fun `updates state on timeout`() = runTest {
+        val state: MutableState<AsyncAction<Int>> = mutableStateOf(AsyncAction.Uninitialized)
+        val timeoutMillis = 500L
+        val operationTimeMillis = 1000L
+
+        try {
+            runUpdatingState(state = state) {
+                withTimeout(timeoutMillis.milliseconds) {
+                    delay(operationTimeMillis)
+                }
+                Result.success(0)
+            }
+            fail("Expected TimeoutCancellationException, but nothing was thrown")
+        } catch (e: TimeoutCancellationException) {
+            assertTrue(state.value.isFailure())
+            assertSame(e, state.value.errorOrNull())
+        }
+    }
+}
diff --git a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/room/RoomContentForwarder.kt b/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/room/RoomContentForwarder.kt
index 9d7495691d..323b4b50f8 100644
--- a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/room/RoomContentForwarder.kt
+++ b/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/room/RoomContentForwarder.kt
@@ -14,7 +14,6 @@ import io.element.android.libraries.matrix.api.core.RoomId
 import io.element.android.libraries.matrix.api.room.ForwardEventException
 import io.element.android.libraries.matrix.impl.roomlist.roomOrNull
 import io.element.android.libraries.matrix.impl.timeline.runWithTimelineListenerRegistered
-import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.withTimeout
 import org.matrix.rustcomponents.sdk.MsgLikeKind
 import org.matrix.rustcomponents.sdk.RoomListService
@@ -63,9 +62,6 @@ class RoomContentForwarder(
                 }
             }.onFailure {
                 failedForwardingTo.add(RoomId(room.id()))
-                if (it is CancellationException) {
-                    throw it
-                }
             }
         }