git.iko.soy/opencode-android
on 5d56632 help sign in
opencode-android/app/src/main/java/soy/iko/opencode/ui/usage/UsageViewModel.kt
78 lines 3.1 kB history raw
1package soy.iko.opencode.ui.usage
3import androidx.lifecycle.ViewModel
4import androidx.lifecycle.viewModelScope
5import kotlinx.coroutines.Dispatchers
6import kotlinx.coroutines.Job
7import kotlinx.coroutines.async
8import kotlinx.coroutines.awaitAll
9import kotlinx.coroutines.flow.MutableStateFlow
10import kotlinx.coroutines.flow.StateFlow
11import kotlinx.coroutines.flow.asStateFlow
12import kotlinx.coroutines.launch
13import kotlinx.coroutines.sync.Semaphore
14import kotlinx.coroutines.sync.withPermit
15import kotlinx.coroutines.withContext
16import soy.iko.opencode.di.AppContainer
17import soy.iko.opencode.util.runCatchingCancellable
19/**
20 * Computes a cross-session [UsageReport] by listing the server's sessions and fetching each
21 * one's messages (bounded concurrency), then aggregating cost/tokens off the main thread.
22 * There's no server-side usage endpoint, so this reconstructs totals from the per-message
23 * cost/token fields the API already returns.
24 */
25class UsageViewModel(private val container: AppContainer) : ViewModel() {
27 sealed interface State {
28 data object Loading : State
29 data object Disconnected : State
30 data object Error : State
31 data class Ready(val report: UsageReport) : State
32 }
34 private val _state = MutableStateFlow<State>(State.Loading)
35 val state: StateFlow<State> = _state.asStateFlow()
37 // Tracks the in-flight load so a second trigger (refresh + retry are both wired to
38 // load()) cancels the first instead of racing it — otherwise the later-*completing* load
39 // wins regardless of which started first, e.g. a stale failure clobbering fresh data.
40 private var loadJob: Job? = null
42 init { load() }
44 fun load() {
45 val conn = container.activeConnection.value ?: run {
46 loadJob?.cancel()
47 _state.value = State.Disconnected
48 return
49 }
50 _state.value = State.Loading
51 loadJob?.cancel()
52 loadJob = viewModelScope.launch {
53 runCatchingCancellable {
54 val sessions = conn.repository.listSessions()
55 // Cap concurrent message fetches so a large history doesn't open dozens of
56 // sockets at once; a per-session fetch failure yields an empty list rather
57 // than aborting the whole report.
58 val limiter = Semaphore(MAX_CONCURRENT_FETCHES)
59 val perSession = sessions.map { session ->
60 async {
61 limiter.withPermit {
62 val messages = runCatchingCancellable { conn.api.listMessages(session.id) }
63 .getOrDefault(emptyList())
64 Triple(session.id, session.displayTitle, messages)
65 }
66 }
67 }.awaitAll()
68 withContext(Dispatchers.Default) { aggregateUsage(perSession) }
69 }
70 .onSuccess { _state.value = State.Ready(it) }
71 .onFailure { _state.value = State.Error }
72 }
73 }
75 private companion object {
76 const val MAX_CONCURRENT_FETCHES = 6
77 }