git.iko.soy/opencode-android
on 0a26f8d help sign in
opencode-android/app/src/main/java/soy/iko/opencode/ui/usage/UsageViewModel.kt
92 lines 3.9 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 /** Bumped by [load] (refresh/retry) to re-run without waiting for a connection change. */
42 private val _reload = MutableStateFlow(0)
44 init {
45 // Observe the active connection so the report loads (or reloads) once a connection is
46 // available — including when the screen opens during a reconnect window where
47 // activeConnection.value is momentarily null (else it would latch Disconnected until a
48 // manual refresh). collectLatest cancels an in-flight load when the connection changes or
49 // a manual reload arrives, so the later-*starting* trigger always wins.
50 viewModelScope.launch {
51 merge(container.activeConnection, _reload).collectLatest { doLoad() }
52 }
53 }
55 fun load() { _reload.value++ }
57 private suspend fun doLoad() {
58 val conn = container.activeConnection.value ?: run {
59 _state.value = State.Disconnected
60 return
61 }
62 _state.value = State.Loading
63 runCatchingCancellable {
64 val sessions = conn.repository.listSessions()
65 // Cap concurrent message fetches so a large history doesn't open dozens of
66 // sockets at once; a per-session fetch failure yields an empty list rather
67 // than aborting the whole report.
68 val limiter = Semaphore(MAX_CONCURRENT_FETCHES)
69 // coroutineScope so the per-session fetches are children of THIS load: when
70 // collectLatest cancels a stale load (connection changed / manual reload), the
71 // in-flight fetches are cancelled with it instead of leaking.
72 val perSession = coroutineScope {
73 sessions.map { session ->
74 async {
75 limiter.withPermit {
76 val messages = runCatchingCancellable { conn.api.listMessages(session.id) }
77 .getOrDefault(emptyList())
78 Triple(session.id, session.displayTitle, messages)
79 }
80 }
81 }.awaitAll()
82 }
83 withContext(Dispatchers.Default) { aggregateUsage(perSession) }
84 }
85 .onSuccess { _state.value = State.Ready(it) }
86 .onFailure { _state.value = State.Error }
87 }
89 private companion object {
90 const val MAX_CONCURRENT_FETCHES = 6
91 }