git.iko.soy/opencode-android
on 066cded help sign in
opencode-android/app/src/main/java/soy/iko/opencode/ui/usage/UsageViewModel.kt
100 lines 4.5 kB history raw
1package soy.iko.opencode.ui.usage
3import androidx.lifecycle.ViewModel
4import androidx.lifecycle.viewModelScope
5import kotlinx.coroutines.ExperimentalCoroutinesApi
6import kotlinx.coroutines.Dispatchers
7import kotlinx.coroutines.async
8import kotlinx.coroutines.awaitAll
9import kotlinx.coroutines.coroutineScope
10import kotlinx.coroutines.flow.MutableStateFlow
11import kotlinx.coroutines.flow.StateFlow
12import kotlinx.coroutines.flow.asStateFlow
13import kotlinx.coroutines.flow.collectLatest
14import kotlinx.coroutines.flow.merge
15import kotlinx.coroutines.launch
16import kotlinx.coroutines.sync.Semaphore
17import kotlinx.coroutines.sync.withPermit
18import kotlinx.coroutines.withContext
19import soy.iko.opencode.di.AppContainer
20import soy.iko.opencode.util.runCatchingCancellable
22/**
23 * Computes a cross-session [UsageReport] by listing the server's sessions and fetching each
24 * one's messages (bounded concurrency), then aggregating cost/tokens off the main thread.
25 * There's no server-side usage endpoint, so this reconstructs totals from the per-message
26 * cost/token fields the API already returns.
27 */
28@OptIn(ExperimentalCoroutinesApi::class)
29class UsageViewModel(private val container: AppContainer) : ViewModel() {
31 sealed interface State {
32 data object Loading : State
33 data object Disconnected : State
34 data object Error : State
35 data class Ready(val report: UsageReport) : State
36 }
38 private val _state = MutableStateFlow<State>(State.Loading)
39 val state: StateFlow<State> = _state.asStateFlow()
41 /** True while a reload runs on top of an already-loaded report, so the existing figures stay
42 * on screen and only the pull-to-refresh spinner shows. */
43 private val _refreshing = MutableStateFlow(false)
44 val refreshing: StateFlow<Boolean> = _refreshing.asStateFlow()
46 /** Bumped by [load] (refresh/retry) to re-run without waiting for a connection change. */
47 private val _reload = MutableStateFlow(0)
49 init {
50 // Observe the active connection so the report loads (or reloads) once a connection is
51 // available — including when the screen opens during a reconnect window where
52 // activeConnection.value is momentarily null (else it would latch Disconnected until a
53 // manual refresh). collectLatest cancels an in-flight load when the connection changes or
54 // a manual reload arrives, so the later-*starting* trigger always wins.
55 viewModelScope.launch {
56 merge(container.activeConnection, _reload).collectLatest { doLoad() }
57 }
58 }
60 fun load() { _reload.value++ }
62 private suspend fun doLoad() {
63 val conn = container.activeConnection.value ?: run {
64 _state.value = State.Disconnected
65 return
66 }
67 // Keep the current report visible during a manual reload; only blank to the full-screen
68 // spinner on the first load.
69 if (_state.value is State.Ready) _refreshing.value = true else _state.value = State.Loading
70 runCatchingCancellable {
71 val sessions = conn.repository.listSessions()
72 // Cap concurrent message fetches so a large history doesn't open dozens of
73 // sockets at once; a per-session fetch failure yields an empty list rather
74 // than aborting the whole report.
75 val limiter = Semaphore(MAX_CONCURRENT_FETCHES)
76 // coroutineScope so the per-session fetches are children of THIS load: when
77 // collectLatest cancels a stale load (connection changed / manual reload), the
78 // in-flight fetches are cancelled with it instead of leaking.
79 val perSession = coroutineScope {
80 sessions.map { session ->
81 async {
82 limiter.withPermit {
83 val messages = runCatchingCancellable { conn.api.listMessages(session.id) }
84 .getOrDefault(emptyList())
85 Triple(session.id, session.displayTitle, messages)
86 }
87 }
88 }.awaitAll()
89 }
90 withContext(Dispatchers.Default) { aggregateUsage(perSession) }
91 }
92 .onSuccess { _state.value = State.Ready(it) }
93 .onFailure { _state.value = State.Error }
94 _refreshing.value = false
95 }
97 private companion object {
98 const val MAX_CONCURRENT_FETCHES = 6
99 }