Firebase – UI wont update after using StateFlows and .collectAsState

Expected Behavior: When the button is clicked, it should immediately update the text and color to reflect the new reward value.

Actual Behavior: When the button is clicked, it calls the functions to update the data and save it to Firestore. However, the UI is not updating until the screen is reloaded.

Tried Things: I’ve switched from LiveData to StateFlow but that did not solve it.

Repository:

var currentUser = MutableStateFlow<User?>(null)
        private set
fun updateCurrentUserFirestore() {
        val updatedUser = hashMapOf(
            "skipCoins" to currentUser.value?.skipCoins,
            "tasksDone" to currentUser.value?.tasksDone,
            "tasksSkipped" to currentUser.value?.tasksSkipped,
            "gamesWon" to currentUser.value?.gamesWon,
            "gamesLost" to currentUser.value?.gamesLost,
            "reward" to currentUser.value?.reward
        )
        firestore.collection("user").document(currentUser.value?.userId.toString())
            .set(updatedUser)
            .addOnSuccessListener {
                Log.d(TAG, "updateCurrentUserFirebase:success")
            }
            .addOnFailureListener {
                Log.w(TAG, "updateCurrentUserFirebase:failure", it)
            }
    }

    fun changeRewardCurrentUser() {
        if (currentUser.value?.reward == "Joke") {
            currentUser.value.let {
                if (it != null) {
                    currentUser.value = User(
                        userId = it.userId,
                        skipCoins = it.skipCoins,
                        tasksDone = it.tasksDone,
                        tasksSkipped = it.tasksSkipped,
                        gamesWon = it.gamesWon,
                        gamesLost = it.gamesLost,
                        reward = if (it.reward == "Joke") "Bored" else "Joke"
                    )
                }
            }
        }
    }
fun updateCurrentUser(uid: String) {
        firestore.collection("user").document(uid)
            .get()
            .addOnSuccessListener { result ->
                currentUser.value = User(
                    userId = result.id,
                    skipCoins = result.data?.get("skipCoins").toString().toLong(),
                    tasksDone = result.data?.get("tasksDone").toString().toLong(),
                    tasksSkipped = result.data?.get("tasksSkipped").toString().toLong(),
                    gamesWon = result.data?.get("gamesWon").toString().toLong(),
                    gamesLost = result.data?.get("gamesLost").toString().toLong(),
                    reward = result.data?.get("reward").toString()
                )
                Log.d(TAG, "updateCurrentUserLocal:success")
            }
            .addOnFailureListener {
                Log.w(TAG, "updateCurrentUserLocal:failure", it)
            }
    }

ViewModel:

val currentUser: StateFlow<User?> = repository.currentUser
fun updateCurrentUserFirestore() {
        repository.updateCurrentUserFirestore()
    }

    fun changeRewardCurrentUser() {
        repository.changeRewardCurrentUser()
    }

Composable:

val user by viewModel.currentUser.collectAsState()
Button(
                        shape = ShapeDefaults.ExtraSmall,
                        onClick = {
                            viewModel.changeRewardCurrentUser()
                            viewModel.updateCurrentUserFirestore()
                        },
                        colors = ButtonDefaults.buttonColors(
                            containerColor = if (user?.reward == "Joke") Purple40 else Orange80,
                            contentColor = if (user?.reward == "Joke") Orange80 else Purple40
                        ),
                        modifier = Modifier.width(
                            calcDp(
                                percentage = 0.25f,
                                dimension = Dimension.Width
                            )
                        )
                    ) {
                        user?.let { Text(text = it.reward) }
                    }