I have the following screen flow:

Initial → StateA → Loading → StateB

The ViewModel first fetches some initial data and then renders StateA. The user enters values manually and presses Next. I then show Loading, call another API, and display StateB. When the user presses Back from StateB, I want to restore StateA with the same manual input values.

sealed interface UiState {
    data object Initial : UiState
    data class StateA(
        val availableTimes: List<String>,
        val hour: String = "",
        val minute: String = "",
        val isPm: Boolean = false
    ) : UiState
    data object Loading : UiState
    data class StateB(
        val result: String
    ) : UiState
}

class MyViewModel : ViewModel() {

    private val _uiState = MutableStateFlow<UiState>(UiState.Initial)
    val uiState: StateFlow<UiState> = _uiState.asStateFlow()

    init {
        fetchInitialData()
    }

    private fun fetchInitialData() {
        viewModelScope.launch {
            delay(500)

            _uiState.value = UiState.StateA(
                availableTimes = listOf(
                    "10:00 AM",
                    "11:00 AM",
                    "12:00 PM"
                )
            )
        }
    }

    fun onInputChanged(
        hour: String,
        minute: String,
        isPm: Boolean
    ) {
        val currentState = _uiState.value as? UiState.StateA ?: return

        _uiState.value = currentState.copy(
            hour = hour,
            minute = minute,
            isPm = isPm
        )
    }

    fun onNextClicked() {
        val screenAState = _uiState.value as? UiState.StateA ?: return

        viewModelScope.launch {
            _uiState.value = UiState.Loading

            delay(1_000)

            _uiState.value = UiState.StateB(
                result = "${screenAState.hour}:${screenAState.minute}"
            )
        }
    }

    fun onBackClicked() {
        // How should ScreenA be restored with its previous input?
    }
}

Since the StateFlow holds only one sealed state at a time, the previous StateA state is no longer available after changing the state to Loading and then StateB.

My team follows this single-state pattern, so using one large state object containing data for every screen is not an option.

What is the recommended way to restore the previous StateA state when the user presses Back?

The state is inside a Dialog fragment with Compose ui and that's why I use a single viewmodel for the state.