git.iko.soy/opencode-android
on e66c7ed help sign in
opencode-android/app/src/main/java/soy/iko/opencode/ui/usage/UsageViewModel.kt
147 lines 7.2 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
21import soy.iko.opencode.util.safeExceptionSummary
23/**
24 * Computes a cross-session [UsageReport] by listing the server's sessions and fetching each
25 * one's messages (bounded concurrency), then aggregating cost/tokens off the main thread.
26 * There's no server-side usage endpoint, so this reconstructs totals from the per-message
27 * cost/token fields the API already returns.
28 */
29@OptIn(ExperimentalCoroutinesApi::class)
30class UsageViewModel(private val container: AppContainer) : ViewModel() {
32 sealed interface State {
33 data object Loading : State
34 data object Disconnected : State
35 /** Usage failed to load. [message] carries a short, safe summary (e.g. "offline",
36 * "server error") so the user can tell an auth failure from a transient network drop. */
37 data class Error(val message: String? = null) : State
38 data class Ready(val report: UsageReport) : State
39 }
41 private val _state = MutableStateFlow<State>(State.Loading)
42 val state: StateFlow<State> = _state.asStateFlow()
44 /** True while a reload runs on top of an already-loaded report, so the existing figures stay
45 * on screen and only the pull-to-refresh spinner shows. */
46 private val _refreshing = MutableStateFlow(false)
47 val refreshing: StateFlow<Boolean> = _refreshing.asStateFlow()
49 /** Bumped by [load] (refresh/retry) to re-run without waiting for a connection change. */
50 private val _reload = MutableStateFlow(0)
52 /** Active time-range filter. Changing it re-aggregates from the cached fetch (no network)
53 * so toggling ranges is instant. Defaults to all time. */
54 private val _timeRange = MutableStateFlow(UsageTimeRange.ALL_TIME)
55 val timeRange: StateFlow<UsageTimeRange> = _timeRange.asStateFlow()
57 // Cached raw inputs from the last successful fetch, so a time-range change can re-aggregate
58 // without re-downloading every session's history.
59 @Volatile private var cachedInputs: List<UsageSessionInput>? = null
61 init {
62 // Observe the active connection so the report loads (or reloads) once a connection is
63 // available — including when the screen opens during a reconnect window where
64 // activeConnection.value is momentarily null (else it would latch Disconnected until a
65 // manual refresh). collectLatest cancels an in-flight load when the connection changes or
66 // a manual reload arrives, so the later-*starting* trigger always wins.
67 viewModelScope.launch {
68 merge(container.activeConnection, _reload).collectLatest { doLoad() }
69 }
70 // Re-aggregate from the cache when only the time range changes (no network round-trip).
71 viewModelScope.launch {
72 _timeRange.collectLatest { reaggregate() }
73 }
74 }
76 fun load() { _reload.value++ }
78 fun setTimeRange(range: UsageTimeRange) { _timeRange.value = range }
80 /** Re-aggregate the cached fetch under the current time range, if a Ready report exists. */
81 private suspend fun reaggregate() {
82 val inputs = cachedInputs ?: return
83 // Only re-aggregate into Ready when we're already Ready: a failed doLoad leaves the
84 // state as Error (with cachedInputs still populated from the prior successful fetch),
85 // and reaggregating that stale cache would flip Error back to Ready — hiding the
86 // failure and showing old totals as if the load succeeded. The user must retry (which
87 // clears cachedInputs on its own success) to see fresh data again.
88 if (_state.value !is State.Ready) return
89 val cutoff = _timeRange.value.cutoffMs()
90 withContext(Dispatchers.Default) {
91 val report = aggregateUsage(inputs, cutoff)
92 _state.value = State.Ready(report)
93 }
94 }
96 private suspend fun doLoad() {
97 val conn = container.activeConnection.value ?: run {
98 _state.value = State.Disconnected
99 return
100 }
101 // Keep the current report visible during a manual reload; only blank to the full-screen
102 // spinner on the first load.
103 if (_state.value is State.Ready) _refreshing.value = true else _state.value = State.Loading
104 try {
105 runCatchingCancellable {
106 val sessions = conn.repository.listSessions()
107 // Cap concurrent message fetches so a large history doesn't open dozens of
108 // sockets at once; a per-session fetch failure yields an empty list rather
109 // than aborting the whole report.
110 val limiter = Semaphore(MAX_CONCURRENT_FETCHES)
111 // coroutineScope so the per-session fetches are children of THIS load: when
112 // collectLatest cancels a stale load (connection changed / manual reload), the
113 // in-flight fetches are cancelled with it instead of leaking.
114 val perSession = coroutineScope {
115 sessions.map { session ->
116 async {
117 limiter.withPermit {
118 val messages = runCatchingCancellable { conn.api.listMessages(session.id) }
119 .getOrDefault(emptyList())
120 UsageSessionInput(
121 sessionId = session.id,
122 title = session.displayTitle,
123 modifiedAt = session.time?.updated ?: session.time?.created,
124 messages = messages,
125 )
126 }
127 }
128 }.awaitAll()
129 }
130 cachedInputs = perSession
131 val cutoff = _timeRange.value.cutoffMs()
132 withContext(Dispatchers.Default) { aggregateUsage(perSession, cutoff) }
133 }
134 .onSuccess { _state.value = State.Ready(it) }
135 .onFailure { _state.value = State.Error(safeExceptionSummary(it)) }
136 } finally {
137 // Reset in finally so a collectLatest cancellation (which rethrows
138 // CancellationException out of runCatchingCancellable, skipping the lines above)
139 // can't leave the pull-to-refresh spinner latched true.
140 _refreshing.value = false
141 }
142 }
144 private companion object {
145 const val MAX_CONCURRENT_FETCHES = 6
146 }