ViewModel and StateFlow patterns for LLM chat apps

Building a responsive LLM chat app on Android demands a UI layer that survives configuration changes and streams tokens without leaking coroutines or blocking the main thread. A well-structured viewmodel stateflow llm chat app architecture keeps conversation state authoritative inside the ViewModel, renders incremental updates through StateFlow, and cancels in-flight requests when the user navigates away. Get this wrong and you’ll see duplicated messages, frozen spinners, or crashed streams the moment the user rotates the screen.

1. Model the conversation as immutable state

Start with a sealed interface for the UI state. The ViewModel should be the single writer; the UI only reads.

data class ChatMessage(
    val id: String,
    val role: String, // "user" | "assistant"
    val content: String,
    val isComplete: Boolean = true
)

data class ChatUiState(
    val messages: List<ChatMessage> = emptyList(),
    val isGenerating: Boolean = false,
    val error: String? = null
)

Never mutate ChatMessage.content in place. Each token arrival produces a new ChatUiState via MutableStateFlow.update { }. This makes diffing trivial in Compose or RecyclerView.

2. Expose StateFlow, not LiveData

StateFlow is built for coroutines, conflates rapidly emitted values, and requires an initial value. LiveData adds main-thread constraints you don’t need in a Kotlin-first codebase.

class ChatViewModel(
    private val repo: LLMRepository
) : ViewModel() {

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

    private var generationJob: Job? = null
}

The asStateFlow() call hides the mutable type from consumers. If you expose MutableStateFlow directly, a fragment can accidentally push state and break the unidirectional flow.

3. Drive streaming from the ViewModel with coroutines

The repository exposes a cold Flow<String> of token deltas. The ViewModel collects it inside viewModelScope, appending tokens to the in-progress assistant message.

fun sendUserMessage(text: String) {
    val userMsg = ChatMessage(UUID.randomUUID().toString(), "user", text)
    _uiState.update { it.copy(messages = it.messages + userMsg, error = null) }

    val assistantId = UUID.randomUUID().toString()
    _uiState.update {
        it.copy(
            messages = it.messages + ChatMessage(assistantId, "assistant", "", isComplete = false),
            isGenerating = true
        )
    }

    generationJob = viewModelScope.launch {
        try {
            repo.streamCompletion(_uiState.value.messages).collect { token ->
                _uiState.update { state ->
                    val updated = state.messages.map { msg ->
                        if (msg.id == assistantId) msg.copy(content = msg.content + token)
                        else msg
                    }
                    state.copy(messages = updated)
                }
            }
            _uiState.update { state ->
                state.copy(
                    isGenerating = false,
                    messages = state.messages.map { if (it.id == assistantId) it.copy(isComplete = true) else it }
                )
            }
        } catch (e: Exception) {
            _uiState.update { it.copy(isGenerating = false, error = e.message) }
        }
    }
}

The repo.streamCompletion call should return a flow built on an HTTP SSE client. Using callbackFlow with OkHttp’s EventSource is a proven pattern:

fun streamCompletion(messages: List<ChatMessage>): Flow<String> = callbackFlow {
    val request = Request.Builder()
        .url("https://api.openai.com/v1/chat/completions")
        .post(buildRequestBody(messages))
        .build()
    val client = OkHttpClient()
    val es = EventSource.create(client, request, object : EventSourceListener() {
        override fun onEvent(source: EventSource, id: String?, type: String?, data: String) {
            if (data == "[DONE]") close() else trySend(parseToken(data))
        }
        override fun onFailure(source: EventSource, t: Throwable?, response: Response?) {
            cancel(CancellationException(t?.message))
        }
    })
    awaitClose { es.cancel() }
}

This keeps the network stream scoped to the collector’s coroutine. When the ViewModel clears, viewModelScope cancels the collection and awaitClose tears down the socket.

4. Cancel, stop, and retry without orphaned jobs

A stop button is mandatory for any chat UI. Expose a function that cancels the active job and marks the message complete.

fun stopGeneration() {
    generationJob?.cancel()
    _uiState.update { state ->
        state.copy(
            isGenerating = false,
            messages = state.messages.map { if (!it.isComplete) it.copy(isComplete = true) else it }
        )
    }
}

If you relaunch sendUserMessage while a job is running, use generationJob?.cancel() before starting a new one. viewModelScope.launch does not automatically cancel previous work. For retry, surface the error in state and let the UI call sendUserMessage again with the same text—but dedupe by checking isGenerating.

5. Batch updates to avoid recomposition storms

Emitting a StateFlow on every token (often 20–100 per second) can thrash Compose or bind RecyclerView excessively. Two tradeoffs exist:

• Immediate tokens: best for perceived latency, worst for UI throughput.
• Buffered batches: use flow.buffer() and a debounce(30) before updating state, sacrificing a little freshness for frame stability.

repo.streamCompletion(msgs)
    .buffer(Channel.BUFFERED)
    .debounce(16) // ~1 frame
    .collect { token -> /* update state */ }

For a viewmodel stateflow llm chat app targeting low-end devices, batching is non-negotiable. On modern hardware, direct emission is usually fine until you attach heavy composables.

6. Persist messages off the main path

StateFlow is not a database. Write completed messages to Room inside the same viewModelScope but using Dispatchers.IO via withContext, or better, have the repository persist as it streams.

viewModelScope.launch(Dispatchers.IO) {
    messageDao.insert(userMsg)
}

Do not block _uiState.update on a DAO call. If you need restore-after-process-death, read the last session in init of the ViewModel using viewModelScope.launch { _uiState.value = loadFromDb() }.

7. Common pitfalls in a viewmodel stateflow llm chat app

• Exposing mutable state: leaking MutableStateFlow to fragments invites race conditions.
• Holding the stream in a field: store only the Job, never the Flow or EventSource.
• Ignoring cancellation: if awaitClose doesn’t cancel the EventSource, you leak sockets across screen rotations.
• Updating state with = instead of update: concurrent emits can clobber each other; update is atomic.
• Forgetting provider cache hints: some gateways forward cache_control to skip re-sending long system prompts. If your repository rebuilds the full message list every call without honoring those hints, you burn tokens and latency.
• Mixing business errors with UI state: keep error: String? for display only; log the exception separately.

8. Pointing the repository at a resilient gateway

When you replace the direct OpenAI call with an OpenAI-compatible endpoint that aggregates providers, the ViewModel logic barely changes. A gateway such as n4n.ai returns the same SSE shape, adds automatic fallback when a provider is rate-limited, and forwards per-token metering so you can surface cost in the UI without custom accounting. Your streamCompletion simply targets a different base URL and includes a route directive header; the callbackFlow code stays identical.

The key architectural win is that the viewmodel stateflow llm chat app stays blind to provider outages. The retry path in the ViewModel becomes a thin wrapper: on CancellationException from a 429, the gateway has already shifted the request upstream. You still cancel the job on user stop, but you don’t need exponential backoff logic scattered in the UI layer.

9. Testing the state machine

Use TurboTest or kotlinx-coroutines-test to drive the ViewModel.

@Test
fun `streaming appends tokens and completes`() = runTest {
    val fakeRepo = object : LLMRepository {
        override fun streamCompletion(m: List<ChatMessage>) = flowOf("Hello", " world")
    }
    val vm = ChatViewModel(fakeRepo)
    vm.sendUserMessage("hi")
    advanceUntilIdle()
    val state = vm.uiState.value
    assertEquals("Hello world", state.messages.last().content)
    assertFalse(state.isGenerating)
}

Fake the repository with a cold flow to assert exact state transitions. This catches the classic bug where the assistant message is added twice because the UI also appends on its own.

10. Final checklist

• Single MutableStateFlow inside ViewModel, exposed as StateFlow.
• All mutations via update { }.
• Stream collected in viewModelScope, cancelled on stop.
• Network socket closed in awaitClose.
• Token batching measured on target hardware.
• Persistence on IO dispatcher, never blocking state writes.
• Repository agnostic to gateway fallback.

Follow this ordered path and your viewmodel stateflow llm chat app will handle rotation, interruption, and provider flakiness without surprising the user or the garbage collector.