A little bit of inheritance and generics to have all events processed at the same place

EventDrivenViewModel.kt

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.hobblies.libraries.architecture.extensions.forContext
import com.hobblies.libraries.architecture.extensions.eventDrivenState
import kotlinx.coroutines.channels.ProducerScope
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.launch

abstract class EventDrivenViewModel<State, Event> : ViewModel() {

    private val _events = MutableSharedFlow<Event>()

    val uiState: StateFlow<State> by eventDrivenState(
        events = _events,
        initialEvent = initialEvent(),
        initialState = initialState(),
        onEvent = { event ->
            forContext(this, uiState.value) {
                onEvent(event)
            }
        }
    )

    open fun initialEvent(): Event? = null
    abstract fun initialState(): State

    context(collector: ProducerScope<State>, currentState: State)
    abstract suspend fun onEvent(event: Event)

    fun handleEvent(event: Event) {
        viewModelScope.launch {
            _events.emit(event)
        }
    }
}

forContext.kt

inline fun <A, B, Result> forContext(a: A, b: B, block: context(A, B) () -> Result): Result {
    return block(a, b)
}

ViewModelExtensions.kt

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.ProducerScope
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.flatMapLatest
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.merge
import kotlinx.coroutines.flow.stateIn
import kotlin.properties.ReadOnlyProperty
import kotlin.reflect.KProperty

@OptIn(ExperimentalCoroutinesApi::class)
fun <State, Event> ViewModel.eventDrivenState(
    initialState: State,
    initialEvent: Event? = null,
    events: Flow<Event>,
    onEvent: suspend ProducerScope<State>.(Event) -> Unit
): ReadOnlyProperty<Any?, StateFlow<State>> = object : ReadOnlyProperty<Any?, StateFlow<State>> {
    private val stateFlow: StateFlow<State> by lazy {

        val initialFlow: Flow<Event> = initialEvent
            ?.let { merge(flow { emit(initialEvent) }, events) }
            ?: events

        initialFlow
            .flatMapLatest { event ->
                channelFlow {
                    onEvent(event)
                    awaitClose()
                }
            }.stateIn(
                scope = viewModelScope,
                started = SharingStarted.WhileSubscribed(5000),
                initialValue = initialState
            )
    }
    override fun getValue(thisRef: Any?, property: KProperty<*>): StateFlow<State> = stateFlow
}