Streaming AI Responses in Jetpack Compose

sealed interface ReplyState {
    data object Idle : ReplyState
    data object Waiting : ReplyState
    data class Streaming(val text: String) : ReplyState
    data class Done(val text: String) : ReplyState
    data class Stopped(val text: String) : ReplyState
    data class Failed(val partial: String, val message: String) : ReplyState
}

class StreamingViewModel(private val ai: TextStreamer) : ViewModel() {   // TextStreamer: (String) -> Flow<String>
    private val _state = MutableStateFlow<ReplyState>(ReplyState.Idle)
    val state: StateFlow<ReplyState> = _state.asStateFlow()

    private var job: Job? = null
    private var lastPrompt = ""

    fun send(prompt: String) {
        lastPrompt = prompt
        job?.cancel()
        _state.value = ReplyState.Waiting
        job = viewModelScope.launch {
            val buffer = StringBuilder()
            try {
                ai.stream(prompt).collect { chunk ->
                    buffer.append(chunk)
                    _state.value = ReplyState.Streaming(buffer.toString())
                }
                _state.value = ReplyState.Done(buffer.toString())
            } catch (e: CancellationException) {
                _state.value = ReplyState.Stopped(buffer.toString())
                throw e                                           // always rethrow cancellation
            } catch (e: Exception) {
                _state.value = ReplyState.Failed(buffer.toString(), "Response interrupted")
            }
        }
    }

    fun stop() { job?.cancel() }                // cancels the Flow, which cancels the network call
    fun retry() = send(lastPrompt)
}

@Composable
fun ReplyCard(state: ReplyState, onStop: () -> Unit, onRetry: () -> Unit) {
    val text = when (state) {
        is ReplyState.Streaming -> state.text
        is ReplyState.Done -> state.text
        is ReplyState.Stopped -> state.text
        is ReplyState.Failed -> state.partial
        else -> ""
    }
    Column(Modifier.fillMaxWidth().padding(16.dp)) {
        when (state) {
            ReplyState.Waiting -> LinearProgressIndicator(Modifier.fillMaxWidth())
            else -> Text(
                text = text,
                modifier = Modifier.semantics {
                    // Announce once when complete rather than on every chunk
                    if (state is ReplyState.Done) liveRegion = LiveRegionMode.Polite
                },
            )
        }
        Row(Modifier.padding(top = 8.dp), horizontalArrangement = Arrangement.spacedBy(8.dp)) {
            if (state is ReplyState.Waiting || state is ReplyState.Streaming)
                OutlinedButton(onClick = onStop) { Text("Stop") }
            if (state is ReplyState.Failed || state is ReplyState.Stopped)
                TextButton(onClick = onRetry) { Text("Retry") }
        }
        if (state is ReplyState.Failed) Text(state.message, color = MaterialTheme.colorScheme.error)
    }
}

val listState = rememberLazyListState()
val atBottom by remember {
    derivedStateOf {
        val last = listState.layoutInfo.visibleItemsInfo.lastOrNull()
        last == null || last.index >= listState.layoutInfo.totalItemsCount - 1
    }
}
LaunchedEffect(lastMessageText.length) {
    if (atBottom && messages.isNotEmpty()) listState.scrollToItem(messages.lastIndex, Int.MAX_VALUE / 2)
}

@Test
fun stop_keepsPartialText() = runTest {
    val chunks = Channel<String>()
    val vm = StreamingViewModel(TextStreamer { chunks.consumeAsFlow() })
    vm.send("hi")
    chunks.send("Hel"); chunks.send("lo")
    runCurrent()
    vm.stop()
    runCurrent()
    assertEquals(ReplyState.Stopped("Hello"), vm.state.value)
}