git.iko.soy/opencode-android
on c0621ff help sign in
opencode-android/app/src/main/java/soy/iko/opencode/data/repo/SessionRepository.kt
467 lines 25 kB history raw
1package soy.iko.opencode.data.repo
3import soy.iko.opencode.data.model.BusEvent
4import soy.iko.opencode.data.model.FilePromptPart
5import soy.iko.opencode.data.model.MessageInfo
6import soy.iko.opencode.data.model.MessagePartRemoved
7import soy.iko.opencode.data.model.MessagePartUpdated
8import soy.iko.opencode.data.model.MessageRemoved
9import soy.iko.opencode.data.model.MessageUpdated
10import soy.iko.opencode.data.model.MessageWithParts
11import soy.iko.opencode.data.model.ModelRef
12import soy.iko.opencode.data.model.Part
13import soy.iko.opencode.data.model.SessionError
14import soy.iko.opencode.data.model.SessionIdle
15import soy.iko.opencode.data.model.UnknownMessage
16import soy.iko.opencode.data.network.EventStreamClient
17import soy.iko.opencode.data.network.NetworkConfig
18import soy.iko.opencode.data.network.OpencodeApiClient
19import soy.iko.opencode.util.runCatchingCancellable
20import soy.iko.opencode.util.safeExceptionSummary
21import android.util.Log
22import kotlinx.coroutines.Dispatchers
23import kotlinx.coroutines.channels.Channel
24import kotlinx.coroutines.flow.Flow
25import kotlinx.coroutines.flow.channelFlow
26import kotlinx.coroutines.flow.consumeAsFlow
27import kotlinx.coroutines.flow.conflate
28import kotlinx.coroutines.flow.flowOn
29import kotlinx.coroutines.launch
30import kotlinx.coroutines.sync.Mutex
31import kotlinx.coroutines.sync.withLock
33/**
34 * Source of truth for a session's conversation. It seeds from `GET /session/:id/message`
35 * and then reduces [BusEvent]s from the SSE stream into an ordered list of messages,
36 * each keyed by message id with parts keyed by part id (idempotent upserts).
37 */
38open class SessionRepository(
39 private val api: OpencodeApiClient,
40 private val eventStream: EventStreamClient,
41) {
42 open suspend fun listSessions() = api.listSessions()
43 open suspend fun createSession(title: String? = null, directory: String? = null) =
44 api.createSession(title, directory)
45 open suspend fun deleteSession(id: String) = api.deleteSession(id)
46 open suspend fun abort(sessionId: String) = api.abort(sessionId)
48 open suspend fun sendPrompt(
49 sessionId: String,
50 text: String,
51 attachments: List<FilePromptPart> = emptyList(),
52 model: ModelRef?,
53 agent: String? = null,
54 ) = api.sendPrompt(sessionId, text, attachments, model, agent)
56 open suspend fun runCommand(
57 sessionId: String,
58 command: String,
59 arguments: String = "",
60 agent: String? = null,
61 ) = api.runCommand(sessionId, command, arguments = arguments, agent = agent)
63 /**
64 * A live, ordered view of [sessionId]'s messages. Begins collecting the event
65 * stream before the initial REST load so no streamed part is missed, then reduces
66 * events into the in-memory state. The flow is conflated so a fast burst of
67 * streaming tokens (each publishing a fresh snapshot) never back-pressures the
68 * event reducer — a slow collector simply sees the most recent snapshot.
69 */
70 open fun observeMessages(sessionId: String): Flow<List<MessageWithParts>> = channelFlow {
71 val store = MessageStore()
72 val lock = Mutex()
73 // Conflated dirty signal: a burst of reduce()/seed() calls coalesces into a
74 // single snapshot+send per demand cycle. Previously every changed event called
75 // snapshot() — allocating an O(N) copy of the whole message list — and ~97% of
76 // those copies were immediately discarded by the outer .conflate() during fast
77 // streaming (hundreds of tokens/sec). Marking dirty and snapshotting in one
78 // drain loop means N events in a burst allocate ~1 snapshot, not N.
79 val dirty = Channel<Unit>(Channel.CONFLATED)
81 // Mark the store dirty; the drain loop below snapshots and publishes. CONFLATED
82 // collapses repeated marks into a single pending signal, so a tight burst of
83 // events doesn't even allocate N marks.
84 fun publish() { dirty.trySend(Unit) }
86 // Subscribe to events first so we don't miss early deltas during the initial load.
87 launch {
88 eventStream.events.collect { event ->
89 if (lock.withLock { store.reduce(sessionId, event) }) publish()
90 }
91 }
93 // Seed generation counter, shared between the initial load and the reconnect
94 // re-seed. Guards against a stale fetch (from an earlier reconnect, or from a slow
95 // initial load) clobbering a newer seed. Only mutated/read under [lock].
96 var seedGeneration = 0
98 // Re-seed from REST when the SSE stream reconnects after a drop. The REST fetch
99 // is performed *outside* the lock so events arriving via SSE during the fetch are
100 // not blocked. The seed itself is under the lock; a seed generation counter ensures
101 // a stale fetch (from an earlier reconnect) doesn't clobber a newer one.
102 launch {
103 // Re-seed from REST when the SSE stream reconnects after a drop. The
104 // hasConnectedBefore flag is NOT reset on disconnect — it stays true after
105 // the first successful connection so a subsequent reconnect triggers the
106 // re-seed. Resetting it on disconnect would make the condition always false
107 // after the first cycle, turning the re-seed into dead code.
108 //
109 // A seed generation counter guards against a stale re-seed clobbering a
110 // newer one: if two reconnects fire in quick succession, the older fetch
111 // (which returns later) is discarded because its generation no longer matches.
112 var hasConnectedBefore = false
113 eventStream.state.collect { state ->
114 if (state == EventStreamClient.ConnectionState.Connected) {
115 if (hasConnectedBefore) {
116 // Pivot before the fetch: parts already in memory predate this
117 // snapshot and must yield to it, while parts that stream in during
118 // the fetch (added to streamedSincePivot via reduce) outrank it.
119 // Bump the generation under the same lock so the compare in the
120 // initial seed sees a consistent value.
121 val generation = lock.withLock {
122 store.beginReseed()
123 ++seedGeneration
124 }
125 val fresh = runCatchingCancellable { api.listMessages(sessionId) }
126 .onFailure { Log.w("SessionRepository", "Re-seed message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }
127 .getOrDefault(emptyList())
128 var seedChanged = false
129 lock.withLock {
130 if (generation == seedGeneration) {
131 // Merge without pruning: SSE events arriving during the
132 // REST fetch may have added messages not yet indexed by
133 // the REST endpoint. Pruning would silently delete them.
134 // Stale messages from deletions during the disconnect
135 // gap are eventually removed via MessageRemoved events
136 // or on the next app restart.
137 seedChanged = store.seed(fresh, prune = false)
138 }
139 }
140 if (seedChanged) publish()
141 }
142 hasConnectedBefore = true
143 }
144 }
145 }
147 val initial = runCatchingCancellable { api.listMessages(sessionId) }
148 .onFailure { Log.w("SessionRepository", "Initial message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }
149 .getOrDefault(emptyList())
150 // Apply the initial seed only if no reconnect re-seed has run yet (generation still 0).
151 // A re-seed is always a fresher full snapshot, so if the (possibly slow) initial fetch
152 // returns after one has applied, seeding this older REST data would clobber the newer
153 // state — reverting a completed run to "running", regressing cost/token/part state —
154 // until the next SSE event or reconnect healed it.
155 val seeded = lock.withLock {
156 if (seedGeneration == 0) {
157 store.seed(initial)
158 true
159 } else false
160 }
161 if (seeded) publish()
163 // Single drain: one snapshot+send per conflated dirty signal. The snapshot is
164 // taken under the lock so a reader never sees a half-reduced state; send() runs
165 // outside the lock so a slow downstream collector can't stall the drain. This
166 // loop also keeps the flow alive until the collector cancels; the launched jobs
167 // above are children of this scope and are torn down with it.
168 dirty.consumeAsFlow().collect {
169 val snapshot = lock.withLock { store.snapshot() }
170 send(snapshot)
171 }
172 // Run the whole reducer pipeline — event reduction, list rebuilds, and the O(N)
173 // snapshot copy — on Dispatchers.Default rather than the collector's context. The
174 // sole consumer collects via stateIn(viewModelScope), i.e. Main.immediate, so without
175 // this every streamed token's reduce()/snapshot() would execute on the UI thread
176 // (hundreds/sec on a fast run). The seed REST calls here only suspend, so moving them
177 // off Main is harmless. JSON decoding already runs off-Main upstream in EventStreamClient.
178 }.flowOn(Dispatchers.Default).conflate()
180 companion object {
181 /** Convenience: is this event a run-completion signal for [sessionId]?
182 * A null sessionID is NOT treated as a wildcard — an idle event with no
183 * session id must not reset the running state of every open chat. */
184 fun isIdle(event: BusEvent, sessionId: String): Boolean =
185 event is SessionIdle && event.properties.sessionID == sessionId
187 fun isError(event: BusEvent, sessionId: String): Boolean =
188 event is SessionError && event.properties.sessionID == sessionId
189 }
192/**
193 * In-memory reduction state for one observed session. Not thread-safe; guard with a Mutex.
194 * Exposed as `internal` so the reducer logic can be unit-tested directly.
195 */
196internal class MessageStore {
197 // messageId -> mutable holder (info + insertion-ordered parts); the map itself is
198 // insertion-ordered.
199 private val messages = LinkedHashMap<String, Holder>()
201 /**
202 * Mutable per-message holder. Parts live in an insertion-ordered map keyed by part id, so a
203 * streamed part update is an O(1) put instead of an O(P) list rebuild. The immutable
204 * [MessageWithParts] view is materialized lazily and cached until the next mutation — so a
205 * burst of streamed token updates rebuilds a message's part list at most once per published
206 * [snapshot], not once per token. The cached view is also reference-stable across snapshots
207 * while unchanged, so Compose skips recomposing bubbles that didn't change.
208 */
209 private class Holder(var info: MessageInfo, val parts: LinkedHashMap<String, Part>) {
210 private var cached: MessageWithParts? = null
211 fun view(): MessageWithParts =
212 cached ?: MessageWithParts(info, parts.values.toList()).also { cached = it }
213 /** Adopt an existing immutable instance as the cache — avoids an immediate rebuild and
214 * preserves its identity (e.g. when seeding an authoritative REST snapshot). */
215 fun adopt(m: MessageWithParts) { cached = m }
216 fun invalidate() { cached = null }
217 }
219 /** Build a [Holder] from an immutable message, reusing [m] itself as the cached view. */
220 private fun holderOf(m: MessageWithParts): Holder {
221 val parts = LinkedHashMap<String, Part>(m.parts.size)
222 for (p in m.parts) parts[p.id] = p
223 return Holder(m.info, parts).also { it.adopt(m) }
224 }
226 /** Part ids upserted via [reduce] since the last reseed pivot ([beginReseed]) or seed.
227 * On a reconnect re-seed these are the parts that streamed in *during* the REST fetch,
228 * so they're newer than the snapshot and must win over it; every other part is taken
229 * from the authoritative REST snapshot. Without this distinction the merge kept every
230 * overlapping in-memory part, so a part that changed server-side during the disconnect
231 * (e.g. a tool going running -> completed) stayed stale until app restart. */
232 private val streamedSincePivot = mutableSetOf<String>()
234 /** Message ids whose `info` was updated via a live SSE [reduce] since the last reseed
235 * pivot ([beginReseed]) or seed. Parallels [streamedSincePivot] but for message info:
236 * on a reconnect re-seed these are the messages whose info changed *during* the REST
237 * fetch, so the in-memory info is newer than the snapshot and must win; every other
238 * message's info is taken from the authoritative REST snapshot. Without this the merge
239 * kept the in-memory info, so a run that finished server-side during the disconnect
240 * (e.g. cost/token totals, completion time) stayed stale until app restart. */
241 private val messageInfoUpdatedSincePivot = mutableSetOf<String>()
243 /** Mark the point a reconnect re-seed's REST fetch begins: parts/info streamed from here
244 * on are newer than the fetched snapshot. Call under the same lock as [seed]/[reduce]. */
245 fun beginReseed() {
246 streamedSincePivot.clear()
247 messageInfoUpdatedSincePivot.clear()
248 }
250 /** Monotonic counter for synthetic UnknownMessage keys, avoiding nanoTime collisions. */
251 private val unknownCounter = java.util.concurrent.atomic.AtomicLong(0)
253 /** Maximum number of messages to keep in memory; oldest are evicted when exceeded. */
254 internal val maxMessages = NetworkConfig.maxInMemoryMessages
256 fun snapshot(): List<MessageWithParts> = messages.values.map { it.view() }
258 fun seed(initial: List<MessageWithParts>, prune: Boolean = false): Boolean {
259 var changed = false
260 // On re-seed (after SSE reconnect), remove messages that are no longer in the
261 // server snapshot (e.g. deleted during the disconnection gap). This keeps the
262 // in-memory state in sync with the authoritative REST snapshot rather than
263 // accumulating stale messages forever. Skipped on the initial seed to avoid
264 // racing with just-arrived SSE events whose messages may not yet be in REST.
265 if (prune) {
266 val incomingIds = initial.mapTo(mutableSetOf()) { it.info.id }
267 val before = messages.size
268 messages.keys.retainAll(incomingIds)
269 if (messages.size != before) changed = true
270 }
272 for (m in initial) {
273 val existing = messages[m.info.id]
274 if (existing == null) {
275 messages[m.info.id] = holderOf(m)
276 changed = true
277 } else {
278 // A part streamed in between subscribe and this initial load may already
279 // have populated this message (see observeMessages). Merge instead of
280 // overwriting so that newer streamed part isn't discarded: take the
281 // snapshot's part order, swap in the streamed version where ids overlap,
282 // append any streamed-only parts, and adopt the authoritative REST info.
283 // existing.parts is already keyed by part id (insertion-ordered).
284 val streamedById = existing.parts
285 val snapshotIds = m.parts.mapTo(mutableSetOf()) { it.id }
286 // Prefer the in-memory version of an overlapping part ONLY if it streamed
287 // in since the reseed pivot (i.e. during this fetch), so it's newer than
288 // the snapshot. Otherwise the REST snapshot is authoritative and wins,
289 // discarding an in-memory part that went stale during a disconnect.
290 val ordered = m.parts.map { p ->
291 if (p.id in streamedSincePivot) streamedById[p.id] ?: p else p
292 }.toMutableList()
293 for (p in existing.parts.values) {
294 if (p.id !in snapshotIds) ordered.add(p)
295 }
296 // Adopt the authoritative REST info unless the in-memory info was itself
297 // updated live since the pivot (i.e. during this fetch), in which case it's
298 // newer than the snapshot. Otherwise the REST snapshot wins, discarding
299 // in-memory info that went stale during a disconnect.
300 val info = if (existing.info is UnknownMessage || m.info.id !in messageInfoUpdatedSincePivot) m.info else existing.info
301 val merged = MessageWithParts(info = info, parts = ordered)
302 if (merged != existing.view()) {
303 messages[m.info.id] = holderOf(merged)
304 changed = true
305 }
306 }
307 }
308 // A part that streamed in *before* this initial load (observeMessages subscribes to
309 // SSE before fetching) inserts its message at the front of the insertion-ordered map,
310 // and the merge above keeps that position — pinning a running message above the older
311 // history that seed() inserts afterward. Restore chronological order by creation time
312 // FIRST, so eviction below drops the genuinely-oldest entries: evictOldMessages()
313 // removes from the map's front, which before reordering is the newest streamed message.
314 if (reorderByTime()) changed = true
315 if (evictOldMessages()) changed = true
316 // Reset the pivot: everything just merged is now the baseline, so only parts/info
317 // that stream in *after* this seed can outrank the next reconnect's snapshot.
318 streamedSincePivot.clear()
319 messageInfoUpdatedSincePivot.clear()
320 return changed
321 }
323 /** Reorder the message map by message creation time (stable). Ties and entries with no
324 * server time yet (e.g. a brand-new message still streaming and not in REST) keep their
325 * relative order, so such a message stays last. Preserves each entry's map key so the
326 * synthetic UnknownMessage keys aren't collapsed. Returns true if the order changed. */
327 private fun reorderByTime(): Boolean {
328 val entries = messages.entries.map { it.key to it.value }
329 val sorted = entries.sortedBy { it.second.info.time?.created ?: Long.MAX_VALUE }
330 // Detect a reordering by comparing keys positionally, instead of allocating two more
331 // N-sized key lists just to compare them.
332 var changed = false
333 for (i in sorted.indices) {
334 if (sorted[i].first != entries[i].first) { changed = true; break }
335 }
336 if (!changed) return false
337 messages.clear()
338 for ((k, v) in sorted) messages[k] = v
339 return true
340 }
342 /** Returns true if the state changed (and a new snapshot should be published). */
343 fun reduce(sessionId: String, event: BusEvent): Boolean {
344 return when (event) {
345 is MessageUpdated -> handleMessageUpdated(sessionId, event)
347 is MessagePartUpdated -> {
348 val part = event.properties.part
349 val messageId = part.messageID ?: event.properties.messageID
350 val partSession = part.sessionID ?: event.properties.sessionID
351 when {
352 messageId == null -> false
353 partSession != null -> if (partSession == sessionId) upsertPart(messageId, part) else false
354 // No session id anywhere on the event: a null session is NOT a wildcard
355 // (mirroring isIdle/isError). Every open session's store sees the shared
356 // event stream, so accepting it unconditionally would leak the part into
357 // unrelated conversations (e.g. both panes in two-pane mode). Attribute it
358 // only if this store already holds the parent message — MessageUpdated,
359 // which always carries a session id, will have created it here first.
360 else -> if (messages.containsKey(messageId)) upsertPart(messageId, part) else false
361 }
362 }
364 is MessagePartRemoved -> handlePartRemoved(sessionId, event)
366 is MessageRemoved -> handleMessageRemoved(sessionId, event)
368 else -> false
369 }
370 }
372 private fun handleMessageUpdated(sessionId: String, event: MessageUpdated): Boolean {
373 val info = event.properties.info
374 if (info.sessionID != sessionId) return false
375 // An UnknownMessage with an empty id (e.g. from an unrecognized server role with no id
376 // field) would collide with other such messages in the map. Generate a unique synthetic
377 // key so each unknown message gets its own entry instead of overwriting others.
378 val key = if (info.id.isEmpty() && info is UnknownMessage) {
379 "unknown-${unknownCounter.incrementAndGet()}"
380 } else {
381 info.id
382 }
383 val existing = messages[key]
384 // A server re-sending byte-identical info is a no-op: skip it so we don't publish a
385 // redundant snapshot (mirrors upsertPart's identical-part guard).
386 if (existing != null && existing.info == info) return false
387 if (existing == null) {
388 messages[key] = Holder(info, LinkedHashMap())
389 } else {
390 existing.info = info
391 existing.invalidate()
392 }
393 // Record that this message's info streamed in live, so a re-seed after a reconnect knows
394 // it's newer than the REST snapshot and keeps it (see seed()/beginReseed()).
395 messageInfoUpdatedSincePivot.add(key)
396 if (existing == null) evictOldMessages()
397 return true
398 }
400 /** Returns false if [eventSession] is non-null and doesn't match [sessionId]. */
401 private fun matchesSession(eventSession: String?, sessionId: String): Boolean =
402 eventSession == null || eventSession == sessionId
404 private fun handlePartRemoved(sessionId: String, event: MessagePartRemoved): Boolean {
405 val messageId = event.properties.messageID ?: return false
406 val partId = event.properties.partID ?: return false
407 if (!matchesSession(event.properties.sessionID, sessionId)) return false
408 return removePart(messageId, partId)
409 }
411 private fun handleMessageRemoved(sessionId: String, event: MessageRemoved): Boolean {
412 val id = event.properties.messageID ?: return false
413 if (!matchesSession(event.properties.sessionID, sessionId)) return false
414 val removed = messages.remove(id) ?: return false
415 messageInfoUpdatedSincePivot.remove(id)
416 removed.parts.keys.forEach { streamedSincePivot.remove(it) }
417 return true
418 }
420 private fun upsertPart(messageId: String, part: Part): Boolean {
421 val current = messages[messageId]
422 if (current != null) {
423 // O(1) map lookup + put, vs the former O(P) indexOfFirst scan and full-list copy
424 // on every streamed token. A LinkedHashMap re-put keeps the part's original
425 // position (replace-in-place); a new id is appended — matching the old list order.
426 if (current.parts[part.id] == part) return false
427 current.parts[part.id] = part
428 current.invalidate()
429 } else {
430 val parts = LinkedHashMap<String, Part>()
431 parts[part.id] = part
432 messages[messageId] = Holder(
433 info = UnknownMessage(id = messageId, sessionID = part.sessionID ?: ""),
434 parts = parts,
435 )
436 }
437 // Record that this part streamed in, so a re-seed after a reconnect knows it's
438 // newer than the REST snapshot and keeps it (see seed()/beginReseed()).
439 streamedSincePivot.add(part.id)
440 if (current == null) evictOldMessages()
441 return true
442 }
444 private fun removePart(messageId: String, partId: String): Boolean {
445 val current = messages[messageId] ?: return false
446 if (current.parts.remove(partId) == null) return false
447 current.invalidate()
448 return true
449 }
451 /** Evict the oldest messages when the store exceeds [maxMessages], keeping memory bounded.
452 * Returns true if any messages were evicted. */
453 private fun evictOldMessages(): Boolean {
454 var evicted = false
455 while (messages.size > maxMessages) {
456 val oldestKey = messages.keys.iterator().next()
457 val removed = messages.remove(oldestKey)
458 // Prune the evicted message's pivot bookkeeping so these sets don't grow
459 // unbounded over a long-running session that never reconnects (they're only
460 // otherwise cleared on reseed/seed).
461 messageInfoUpdatedSincePivot.remove(oldestKey)
462 removed?.parts?.keys?.forEach { streamedSincePivot.remove(it) }
463 evicted = true
464 }
465 return evicted
466 }