How to perform retry for loading UI state when using .stateIn() approach?

[KoTius]

Example

He we load data model from repository. Usually we have here API IO so it can always fail.
In the example we map the fail into the Error class and later into error Ui State. Understandable.

Now assuming on the Error Ui we have the "Try again button" which will call onRetryClick() callback in the ViewModel.
How we supposed to re-start all that chain or only _taskAsync chain?

private val _taskAsync = taskRepository.getTaskStream(taskId)
        .map { handleTask(it) }
        .catch { emit(Async.Error(R.string.loading_task_error)) }

   val uiState: StateFlow<TaskDetailUiState> = combine(
        _userMessage, _isLoading, _isTaskDeleted, _taskAsync
    ) { userMessage, isLoading, isTaskDeleted, taskAsync ->
        when (taskAsync) {
            Async.Loading -> {
                TaskDetailUiState(isLoading = true)
            }
            is Async.Error -> {
                TaskDetailUiState(
                    userMessage = taskAsync.errorMessage,
                    isTaskDeleted = isTaskDeleted
                )
            }
            is Async.Success -> {
                TaskDetailUiState(
                    task = taskAsync.data,
                    isLoading = isLoading,
                    userMessage = userMessage,
                    isTaskDeleted = isTaskDeleted
                )
            }
        }
    }

[vatsasiddhartha]

class YourViewModel(private val taskRepository: TaskRepository) : ViewModel() {

private val _taskAsync = MutableStateFlow<Async<Task>>(Async.Uninitialized)

val uiState: StateFlow<TaskDetailUiState> = combine(
    _userMessage, _isLoading, _isTaskDeleted, _taskAsync
) { userMessage, isLoading, isTaskDeleted, taskAsync ->
    // Your existing code for combining state
    // ...
}.stateIn(viewModelScope, SharingStarted.WhileSubscribed(), initialState)

init {
    loadTask()
}

private fun loadTask() {
    _taskAsync.value = Async.Loading
    viewModelScope.launch {
        try {
            val task = taskRepository.getTask(taskId)
            _taskAsync.value = Async.Success(task)
        } catch (e: Exception) {
            _taskAsync.value = Async.Error(R.string.loading_task_error)
        }
    }
}

fun onRetryClick() {
    loadTask()
}

}