git.iko.soy/opencode-android
on 5d56632 help sign in
opencode-android/app/src/main/java/soy/iko/opencode/data/repo/SessionRepository.kt
601 lines 34 kB history raw
1package soy.iko.opencode.data.repo
3import soy.iko.opencode.data.model.AssistantMessage
4import soy.iko.opencode.data.model.BusEvent
5import soy.iko.opencode.data.model.FilePromptPart
6import soy.iko.opencode.data.model.MessageInfo
7import soy.iko.opencode.data.model.MessagePartRemoved
8import soy.iko.opencode.data.model.MessagePartUpdated
9import soy.iko.opencode.data.model.MessageRemoved
10import soy.iko.opencode.data.model.MessageUpdated
11import soy.iko.opencode.data.model.MessageWithParts
12import soy.iko.opencode.data.model.ModelRef
13import soy.iko.opencode.data.model.Part
14import soy.iko.opencode.data.model.SessionError
15import soy.iko.opencode.data.model.SessionIdle
16import soy.iko.opencode.data.model.UnknownMessage
17import soy.iko.opencode.data.network.EventStreamClient
18import soy.iko.opencode.data.network.NetworkConfig
19import soy.iko.opencode.data.network.OpencodeApiClient
20import soy.iko.opencode.util.runCatchingCancellable
21import soy.iko.opencode.util.safeExceptionSummary
22import android.util.Log
23import kotlinx.coroutines.Dispatchers
24import kotlinx.coroutines.Job
25import kotlinx.coroutines.NonCancellable
26import kotlinx.coroutines.delay
27import kotlinx.coroutines.withContext
28import kotlinx.coroutines.channels.Channel
29import kotlinx.coroutines.flow.Flow
30import kotlinx.coroutines.flow.channelFlow
31import kotlinx.coroutines.flow.consumeAsFlow
32import kotlinx.coroutines.flow.conflate
33import kotlinx.coroutines.flow.flowOn
34import kotlinx.coroutines.launch
35import kotlinx.coroutines.sync.Mutex
36import kotlinx.coroutines.sync.withLock
38/**
39 * Source of truth for a session's conversation. It seeds from `GET /session/:id/message`
40 * and then reduces [BusEvent]s from the SSE stream into an ordered list of messages,
41 * each keyed by message id with parts keyed by part id (idempotent upserts).
42 */
43open class SessionRepository(
44 private val api: OpencodeApiClient,
45 private val eventStream: EventStreamClient,
46 // Optional on-disk cache so a conversation renders instantly on open and stays readable
47 // offline. Null in tests (and any path that doesn't wire it), where caching is skipped.
48 private val messageCache: MessageCacheStore? = null,
49) {
50 open suspend fun listSessions() = api.listSessions()
51 open suspend fun createSession(title: String? = null, directory: String? = null) =
52 api.createSession(title, directory)
53 open suspend fun deleteSession(id: String) = api.deleteSession(id)
54 open suspend fun abort(sessionId: String) = api.abort(sessionId)
56 open suspend fun sendPrompt(
57 sessionId: String,
58 text: String,
59 attachments: List<FilePromptPart> = emptyList(),
60 model: ModelRef?,
61 agent: String? = null,
62 // Stable key for retriable/re-sendable messages (the offline outbox). See
63 // [OpencodeApiClient.sendPrompt]. Null generates a fresh key per call.
64 idempotencyKey: String? = null,
65 ) = api.sendPrompt(sessionId, text, attachments, model, agent, idempotencyKey)
67 open suspend fun runCommand(
68 sessionId: String,
69 command: String,
70 arguments: String = "",
71 agent: String? = null,
72 ) = api.runCommand(sessionId, command, arguments = arguments, agent = agent)
74 /**
75 * A live, ordered view of [sessionId]'s messages. Begins collecting the event
76 * stream before the initial REST load so no streamed part is missed, then reduces
77 * events into the in-memory state. The flow is conflated so a fast burst of
78 * streaming tokens (each publishing a fresh snapshot) never back-pressures the
79 * event reducer — a slow collector simply sees the most recent snapshot.
80 */
81 open fun observeMessages(sessionId: String): Flow<List<MessageWithParts>> = channelFlow {
82 val store = MessageStore()
83 val lock = Mutex()
84 // Conflated dirty signal: a burst of reduce()/seed() calls coalesces into a
85 // single snapshot+send per demand cycle. Previously every changed event called
86 // snapshot() — allocating an O(N) copy of the whole message list — and ~97% of
87 // those copies were immediately discarded by the outer .conflate() during fast
88 // streaming (hundreds of tokens/sec). Marking dirty and snapshotting in one
89 // drain loop means N events in a burst allocate ~1 snapshot, not N.
90 val dirty = Channel<Unit>(Channel.CONFLATED)
92 // Mark the store dirty; the drain loop below snapshots and publishes. CONFLATED
93 // collapses repeated marks into a single pending signal, so a tight burst of
94 // events doesn't even allocate N marks.
95 fun publish() { dirty.trySend(Unit) }
97 // Subscribe to events first so we don't miss early deltas during the initial load.
98 launch {
99 eventStream.events.collect { event ->
100 if (lock.withLock { store.reduce(sessionId, event) }) publish()
101 }
102 }
104 // Seed generation counter, shared between the initial load and the reconnect
105 // re-seed. Guards against a stale fetch (from an earlier reconnect, or from a slow
106 // initial load) clobbering a newer seed. Only mutated/read under [lock].
107 var seedGeneration = 0
109 // Message ids introduced by the on-disk cache seed below. The initial REST seed uses
110 // these to prune ghosts: a cache-seeded message absent from the authoritative REST
111 // snapshot was deleted server-side while the app was closed. Restricting the prune to
112 // cache-seeded ids means live SSE messages added since subscribe (not yet indexed by
113 // REST) are never dropped. Written/read only in this collector coroutine.
114 var cacheSeededIds: Set<String> = emptySet()
116 // Seed from the on-disk cache first so the conversation renders instantly — and stays
117 // readable offline — before the network responds. Guarded on generation 0 so a
118 // reconnect re-seed that somehow raced ahead isn't clobbered; the initial REST seed
119 // (also generation 0) then overrides this with authoritative data when it succeeds.
120 messageCache?.let { cache ->
121 val cached = cache.load(sessionId)
122 if (cached.isNotEmpty()) {
123 val changed = lock.withLock {
124 if (seedGeneration == 0) {
125 cacheSeededIds = cached.mapTo(mutableSetOf()) { it.info.id }
126 store.seed(cached)
127 } else false
128 }
129 if (changed) publish()
130 }
131 }
133 // Re-seed from REST when the SSE stream reconnects after a drop. The REST fetch
134 // is performed *outside* the lock so events arriving via SSE during the fetch are
135 // not blocked. The seed itself is under the lock; a seed generation counter ensures
136 // a stale fetch (from an earlier reconnect) doesn't clobber a newer one.
137 launch {
138 // Re-seed from REST when the SSE stream reconnects after a drop. The
139 // hasConnectedBefore flag is NOT reset on disconnect — it stays true after
140 // the first successful connection so a subsequent reconnect triggers the
141 // re-seed. Resetting it on disconnect would make the condition always false
142 // after the first cycle, turning the re-seed into dead code.
143 //
144 // A seed generation counter guards against a stale re-seed clobbering a
145 // newer one: if two reconnects fire in quick succession, the older fetch
146 // (which returns later) is discarded because its generation no longer matches.
147 var hasConnectedBefore = false
148 eventStream.state.collect { state ->
149 if (state == EventStreamClient.ConnectionState.Connected) {
150 if (hasConnectedBefore) {
151 // Pivot before the fetch: parts already in memory predate this
152 // snapshot and must yield to it, while parts that stream in during
153 // the fetch (added to streamedSincePivot via reduce) outrank it.
154 // Bump the generation under the same lock so the compare in the
155 // initial seed sees a consistent value.
156 val generation = lock.withLock {
157 store.beginReseed()
158 ++seedGeneration
159 }
160 val fresh = runCatchingCancellable { api.listMessages(sessionId) }
161 .onFailure { Log.w("SessionRepository", "Re-seed message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }
162 .getOrDefault(emptyList())
163 var seedChanged = false
164 lock.withLock {
165 if (generation == seedGeneration) {
166 // Merge without pruning: SSE events arriving during the
167 // REST fetch may have added messages not yet indexed by
168 // the REST endpoint. Pruning would silently delete them.
169 // Stale messages from deletions during the disconnect
170 // gap are eventually removed via MessageRemoved events
171 // or on the next app restart.
172 seedChanged = store.seed(fresh, prune = false)
173 }
174 }
175 if (seedChanged) publish()
176 }
177 hasConnectedBefore = true
178 }
179 }
180 }
182 // Load the authoritative message list concurrently with the drain loop below so the
183 // cache seed and any early SSE-reduced parts are emitted immediately. Running it inline
184 // here (before the drain starts) would withhold the already-reduced cached conversation
185 // until this fetch resolved — a blank screen for the whole withRetry backoff when offline.
186 launch {
187 val initialResult = runCatchingCancellable { api.listMessages(sessionId) }
188 .onFailure { Log.w("SessionRepository", "Initial message load failed for $sessionId; relying on cache/SSE: ${safeExceptionSummary(it)}") }
189 // Apply the initial seed only if no reconnect re-seed has run yet (generation still 0).
190 // A re-seed is always a fresher full snapshot, so if the (possibly slow) initial fetch
191 // returns after one has applied, seeding this older REST data would clobber the newer
192 // state — reverting a completed run to "running", regressing cost/token/part state —
193 // until the next SSE event or reconnect healed it. Only seed on SUCCESS: on failure
194 // (e.g. offline) keep whatever's shown (the cache seed above) rather than clearing it.
195 val seeded = lock.withLock {
196 val list = initialResult.getOrNull()
197 if (seedGeneration == 0 && list != null) {
198 // Prune ghosts first: cache-seeded messages absent from this authoritative
199 // REST snapshot were deleted server-side while the app was closed. Without
200 // this they'd survive restart forever (re-seeded from cache, never pruned by
201 // the reconnect path, and re-persisted by the drain loop).
202 val pruned = store.pruneStaleCacheSeeded(cacheSeededIds, list)
203 val changed = store.seed(list)
204 pruned || changed
205 } else false
206 }
207 if (seeded) publish()
208 }
210 // Single drain: one snapshot+send per conflated dirty signal. The snapshot is
211 // taken under the lock so a reader never sees a half-reduced state; send() runs
212 // outside the lock so a slow downstream collector can't stall the drain. This
213 // loop also keeps the flow alive until the collector cancels; the launched jobs
214 // above are children of this scope and are torn down with it.
215 var lastCacheWrite = 0L
216 var pendingCacheSave: Job? = null
217 // Most recent published snapshot, and the most recent one whose on-disk write has
218 // actually completed. The teardown flush in the finally below persists a pending
219 // trailing snapshot exactly once. savedSnapshot is assigned from inside the launched
220 // saves *after* messageCache.save() returns, so a save cancelled mid-IO by teardown is
221 // NOT marked saved and the finally re-persists it under NonCancellable. It's written
222 // from the save coroutines and read from the collector/finally; the only possible race
223 // is a stale read causing one redundant, idempotent re-save (atomic temp-file+rename
224 // makes a duplicate write of identical data harmless) — never data loss.
225 var latestSnapshot: List<MessageWithParts>? = null
226 var savedSnapshot: List<MessageWithParts>? = null
227 try {
228 dirty.consumeAsFlow().collect {
229 val snapshot = lock.withLock { store.snapshot() }
230 send(snapshot)
231 // Persist the snapshot to the on-disk cache, throttled so a fast token stream
232 // doesn't hammer the disk — the cache only needs to be "recent enough" for an
233 // instant/offline first paint; the network corrects it on the next open.
234 if (messageCache != null && snapshot.isNotEmpty()) {
235 latestSnapshot = snapshot
236 val now = System.currentTimeMillis()
237 pendingCacheSave?.cancel()
238 if (now - lastCacheWrite >= NetworkConfig.messageCacheWriteThrottleMs) {
239 lastCacheWrite = now
240 pendingCacheSave = launch {
241 messageCache.save(sessionId, snapshot)
242 savedSnapshot = snapshot
243 }
244 } else {
245 // Trailing-edge write: a burst that ends inside the throttle window (a run's
246 // final tokens right after a write) would otherwise never persist its last
247 // snapshot, so an offline reopen would show the conversation missing its tail.
248 // Write the latest snapshot once the window elapses; a newer snapshot cancels
249 // and replaces this, so continuous streaming still coalesces to one write/window.
250 pendingCacheSave = launch {
251 delay(NetworkConfig.messageCacheWriteThrottleMs)
252 messageCache.save(sessionId, snapshot)
253 savedSnapshot = snapshot
254 }
255 }
256 }
257 }
258 } finally {
259 // FIX: both cache-saves above run as children of this inner flow scope, so when
260 // ChatViewModel's flatMapLatest switches (an SSE reconnect changing activeConnection)
261 // or onCleared cancels this flow, a pending trailing save is torn down before its delay
262 // elapses — a run's final tokens then never reach the on-disk cache. Flush the latest
263 // un-persisted snapshot here under NonCancellable so it survives the cancellation.
264 // savedSnapshot tracks the snapshot an immediate write already handled, so the common
265 // case (nothing pending) writes nothing extra — at most one write per teardown, never
266 // one per emission.
267 val pending = latestSnapshot
268 if (messageCache != null && pending != null && pending !== savedSnapshot) {
269 withContext(NonCancellable) { messageCache.save(sessionId, pending) }
270 }
271 }
272 // Run the whole reducer pipeline — event reduction, list rebuilds, and the O(N)
273 // snapshot copy — on Dispatchers.Default rather than the collector's context. The
274 // sole consumer collects via stateIn(viewModelScope), i.e. Main.immediate, so without
275 // this every streamed token's reduce()/snapshot() would execute on the UI thread
276 // (hundreds/sec on a fast run). The seed REST calls here only suspend, so moving them
277 // off Main is harmless. JSON decoding already runs off-Main upstream in EventStreamClient.
278 }.flowOn(Dispatchers.Default).conflate()
280 companion object {
281 /** Convenience: is this event a run-completion signal for [sessionId]?
282 * A null sessionID is NOT treated as a wildcard — an idle event with no
283 * session id must not reset the running state of every open chat. */
284 fun isIdle(event: BusEvent, sessionId: String): Boolean =
285 event is SessionIdle && event.properties.sessionID == sessionId
287 fun isError(event: BusEvent, sessionId: String): Boolean =
288 event is SessionError && event.properties.sessionID == sessionId
290 /** Does this event indicate an in-progress run for [sessionId]? A live streamed part or
291 * an assistant message that isn't yet complete/errored means the agent is still working.
292 * Used to restore the "running" indicator after an SSE reconnect: if the run is still
293 * going, such live events keep arriving; a run that finished during the outage produces
294 * none (its completion arrives via the REST re-seed, not the event stream). */
295 fun isRunActivity(event: BusEvent, sessionId: String): Boolean = when (event) {
296 // The session id rides on the part for message.part.updated; properties.sessionID
297 // is null on the wire. Mirror reduce()'s `part.sessionID ?: properties.sessionID`
298 // so a run that's only streaming parts still re-lights the indicator on reconnect.
299 is MessagePartUpdated -> (event.properties.part.sessionID ?: event.properties.sessionID) == sessionId
300 is MessageUpdated -> {
301 val info = event.properties.info
302 info.sessionID == sessionId && info is AssistantMessage && !info.isComplete && info.error == null
303 }
304 else -> false
305 }
306 }
309/**
310 * In-memory reduction state for one observed session. Not thread-safe; guard with a Mutex.
311 * Exposed as `internal` so the reducer logic can be unit-tested directly.
312 */
313internal class MessageStore {
314 // messageId -> mutable holder (info + insertion-ordered parts); the map itself is
315 // insertion-ordered.
316 private val messages = LinkedHashMap<String, Holder>()
318 /**
319 * Mutable per-message holder. Parts live in an insertion-ordered map keyed by part id, so a
320 * streamed part update is an O(1) put instead of an O(P) list rebuild. The immutable
321 * [MessageWithParts] view is materialized lazily and cached until the next mutation — so a
322 * burst of streamed token updates rebuilds a message's part list at most once per published
323 * [snapshot], not once per token. The cached view is also reference-stable across snapshots
324 * while unchanged, so Compose skips recomposing bubbles that didn't change.
325 */
326 private class Holder(var info: MessageInfo, val parts: LinkedHashMap<String, Part>) {
327 private var cached: MessageWithParts? = null
328 fun view(): MessageWithParts =
329 cached ?: MessageWithParts(info, parts.values.toList()).also { cached = it }
330 /** Adopt an existing immutable instance as the cache — avoids an immediate rebuild and
331 * preserves its identity (e.g. when seeding an authoritative REST snapshot). */
332 fun adopt(m: MessageWithParts) { cached = m }
333 fun invalidate() { cached = null }
334 }
336 /** Build a [Holder] from an immutable message, reusing [m] itself as the cached view. */
337 private fun holderOf(m: MessageWithParts): Holder {
338 val parts = LinkedHashMap<String, Part>(m.parts.size)
339 for (p in m.parts) parts[p.id] = p
340 return Holder(m.info, parts).also { it.adopt(m) }
341 }
343 /** Part ids upserted via [reduce] since the last reseed pivot ([beginReseed]) or seed.
344 * On a reconnect re-seed these are the parts that streamed in *during* the REST fetch,
345 * so they're newer than the snapshot and must win over it; every other part is taken
346 * from the authoritative REST snapshot. Without this distinction the merge kept every
347 * overlapping in-memory part, so a part that changed server-side during the disconnect
348 * (e.g. a tool going running -> completed) stayed stale until app restart. */
349 private val streamedSincePivot = mutableSetOf<String>()
351 /** Message ids whose `info` was updated via a live SSE [reduce] since the last reseed
352 * pivot ([beginReseed]) or seed. Parallels [streamedSincePivot] but for message info:
353 * on a reconnect re-seed these are the messages whose info changed *during* the REST
354 * fetch, so the in-memory info is newer than the snapshot and must win; every other
355 * message's info is taken from the authoritative REST snapshot. Without this the merge
356 * kept the in-memory info, so a run that finished server-side during the disconnect
357 * (e.g. cost/token totals, completion time) stayed stale until app restart. */
358 private val messageInfoUpdatedSincePivot = mutableSetOf<String>()
360 /** Mark the point a reconnect re-seed's REST fetch begins: parts/info streamed from here
361 * on are newer than the fetched snapshot. Call under the same lock as [seed]/[reduce]. */
362 fun beginReseed() {
363 streamedSincePivot.clear()
364 messageInfoUpdatedSincePivot.clear()
365 }
367 /** Monotonic counter for synthetic UnknownMessage keys, avoiding nanoTime collisions. */
368 private val unknownCounter = java.util.concurrent.atomic.AtomicLong(0)
370 /** Maximum number of messages to keep in memory; oldest are evicted when exceeded. */
371 internal val maxMessages = NetworkConfig.maxInMemoryMessages
373 fun snapshot(): List<MessageWithParts> = messages.values.map { it.view() }
375 fun seed(initial: List<MessageWithParts>, prune: Boolean = false): Boolean {
376 var changed = false
377 // On re-seed (after SSE reconnect), remove messages that are no longer in the
378 // server snapshot (e.g. deleted during the disconnection gap). This keeps the
379 // in-memory state in sync with the authoritative REST snapshot rather than
380 // accumulating stale messages forever. Skipped on the initial seed to avoid
381 // racing with just-arrived SSE events whose messages may not yet be in REST.
382 if (prune) {
383 val incomingIds = initial.mapTo(mutableSetOf()) { it.info.id }
384 val before = messages.size
385 messages.keys.retainAll(incomingIds)
386 if (messages.size != before) changed = true
387 }
389 for (m in initial) {
390 val existing = messages[m.info.id]
391 if (existing == null) {
392 messages[m.info.id] = holderOf(m)
393 changed = true
394 } else {
395 // A part streamed in between subscribe and this initial load may already
396 // have populated this message (see observeMessages). Merge instead of
397 // overwriting so that newer streamed part isn't discarded: take the
398 // snapshot's part order, swap in the streamed version where ids overlap,
399 // append any streamed-only parts, and adopt the authoritative REST info.
400 // existing.parts is already keyed by part id (insertion-ordered).
401 val streamedById = existing.parts
402 val snapshotIds = m.parts.mapTo(mutableSetOf()) { it.id }
403 // Prefer the in-memory version of an overlapping part ONLY if it streamed
404 // in since the reseed pivot (i.e. during this fetch), so it's newer than
405 // the snapshot. Otherwise the REST snapshot is authoritative and wins,
406 // discarding an in-memory part that went stale during a disconnect.
407 val ordered = m.parts.map { p ->
408 if (p.id in streamedSincePivot) streamedById[p.id] ?: p else p
409 }.toMutableList()
410 for (p in existing.parts.values) {
411 if (p.id !in snapshotIds) ordered.add(p)
412 }
413 // Adopt the authoritative REST info unless the in-memory info was itself
414 // updated live since the pivot (i.e. during this fetch), in which case it's
415 // newer than the snapshot. Otherwise the REST snapshot wins, discarding
416 // in-memory info that went stale during a disconnect.
417 val info = if (existing.info is UnknownMessage || m.info.id !in messageInfoUpdatedSincePivot) m.info else existing.info
418 val merged = MessageWithParts(info = info, parts = ordered)
419 if (merged != existing.view()) {
420 messages[m.info.id] = holderOf(merged)
421 changed = true
422 }
423 }
424 }
425 // A part that streamed in *before* this initial load (observeMessages subscribes to
426 // SSE before fetching) inserts its message at the front of the insertion-ordered map,
427 // and the merge above keeps that position — pinning a running message above the older
428 // history that seed() inserts afterward. Restore chronological order by creation time
429 // FIRST, so eviction below drops the genuinely-oldest entries: evictOldMessages()
430 // removes from the map's front, which before reordering is the newest streamed message.
431 if (reorderByTime()) changed = true
432 if (evictOldMessages()) changed = true
433 // Reset the pivot: everything just merged is now the baseline, so only parts/info
434 // that stream in *after* this seed can outrank the next reconnect's snapshot.
435 streamedSincePivot.clear()
436 messageInfoUpdatedSincePivot.clear()
437 return changed
438 }
440 /** Remove messages that were seeded from the on-disk cache ([cacheSeededIds]) but are
441 * absent from the authoritative REST snapshot [rest] — i.e. deleted server-side while the
442 * app was closed. Restricted to cache-seeded ids so a message added by a live SSE event
443 * since subscribe (which the REST endpoint may not have indexed yet) is never pruned.
444 * This is what actually removes ghosts on restart; the reconnect re-seed deliberately
445 * keeps prune=false to avoid racing just-arrived SSE. Call under the same lock as [seed].
446 * Returns true if anything was removed. */
447 fun pruneStaleCacheSeeded(cacheSeededIds: Set<String>, rest: List<MessageWithParts>): Boolean {
448 if (cacheSeededIds.isEmpty()) return false
449 val restIds = rest.mapTo(mutableSetOf()) { it.info.id }
450 var changed = false
451 for (id in cacheSeededIds) {
452 if (id !in restIds && messages.remove(id) != null) changed = true
453 }
454 return changed
455 }
457 /** Reorder the message map by message creation time (stable). Ties and entries with no
458 * server time yet (e.g. a brand-new message still streaming and not in REST) keep their
459 * relative order, so such a message stays last. Preserves each entry's map key so the
460 * synthetic UnknownMessage keys aren't collapsed. Returns true if the order changed. */
461 private fun reorderByTime(): Boolean {
462 val entries = messages.entries.map { it.key to it.value }
463 val sorted = entries.sortedBy { it.second.info.time?.created ?: Long.MAX_VALUE }
464 // Detect a reordering by comparing keys positionally, instead of allocating two more
465 // N-sized key lists just to compare them.
466 var changed = false
467 for (i in sorted.indices) {
468 if (sorted[i].first != entries[i].first) { changed = true; break }
469 }
470 if (!changed) return false
471 messages.clear()
472 for ((k, v) in sorted) messages[k] = v
473 return true
474 }
476 /** Returns true if the state changed (and a new snapshot should be published). */
477 fun reduce(sessionId: String, event: BusEvent): Boolean {
478 return when (event) {
479 is MessageUpdated -> handleMessageUpdated(sessionId, event)
481 is MessagePartUpdated -> {
482 val part = event.properties.part
483 val messageId = part.messageID ?: event.properties.messageID
484 val partSession = part.sessionID ?: event.properties.sessionID
485 when {
486 messageId == null -> false
487 partSession != null -> if (partSession == sessionId) upsertPart(messageId, part) else false
488 // No session id anywhere on the event: a null session is NOT a wildcard
489 // (mirroring isIdle/isError). Every open session's store sees the shared
490 // event stream, so accepting it unconditionally would leak the part into
491 // unrelated conversations (e.g. both panes in two-pane mode). Attribute it
492 // only if this store already holds the parent message — MessageUpdated,
493 // which always carries a session id, will have created it here first.
494 else -> if (messages.containsKey(messageId)) upsertPart(messageId, part) else false
495 }
496 }
498 is MessagePartRemoved -> handlePartRemoved(sessionId, event)
500 is MessageRemoved -> handleMessageRemoved(sessionId, event)
502 else -> false
503 }
504 }
506 private fun handleMessageUpdated(sessionId: String, event: MessageUpdated): Boolean {
507 val info = event.properties.info
508 if (info.sessionID != sessionId) return false
509 // An UnknownMessage with an empty id (e.g. from an unrecognized server role with no id
510 // field) would collide with other such messages in the map. Generate a unique synthetic
511 // key so each unknown message gets its own entry instead of overwriting others.
512 val key = if (info.id.isEmpty() && info is UnknownMessage) {
513 "unknown-${unknownCounter.incrementAndGet()}"
514 } else {
515 info.id
516 }
517 val existing = messages[key]
518 // A server re-sending byte-identical info is a no-op: skip it so we don't publish a
519 // redundant snapshot (mirrors upsertPart's identical-part guard).
520 if (existing != null && existing.info == info) return false
521 if (existing == null) {
522 messages[key] = Holder(info, LinkedHashMap())
523 } else {
524 existing.info = info
525 existing.invalidate()
526 }
527 // Record that this message's info streamed in live, so a re-seed after a reconnect knows
528 // it's newer than the REST snapshot and keeps it (see seed()/beginReseed()).
529 messageInfoUpdatedSincePivot.add(key)
530 if (existing == null) evictOldMessages()
531 return true
532 }
534 /** Returns false if [eventSession] is non-null and doesn't match [sessionId]. */
535 private fun matchesSession(eventSession: String?, sessionId: String): Boolean =
536 eventSession == null || eventSession == sessionId
538 private fun handlePartRemoved(sessionId: String, event: MessagePartRemoved): Boolean {
539 val messageId = event.properties.messageID ?: return false
540 val partId = event.properties.partID ?: return false
541 if (!matchesSession(event.properties.sessionID, sessionId)) return false
542 return removePart(messageId, partId)
543 }
545 private fun handleMessageRemoved(sessionId: String, event: MessageRemoved): Boolean {
546 val id = event.properties.messageID ?: return false
547 if (!matchesSession(event.properties.sessionID, sessionId)) return false
548 val removed = messages.remove(id) ?: return false
549 messageInfoUpdatedSincePivot.remove(id)
550 removed.parts.keys.forEach { streamedSincePivot.remove(it) }
551 return true
552 }
554 private fun upsertPart(messageId: String, part: Part): Boolean {
555 val current = messages[messageId]
556 if (current != null) {
557 // O(1) map lookup + put, vs the former O(P) indexOfFirst scan and full-list copy
558 // on every streamed token. A LinkedHashMap re-put keeps the part's original
559 // position (replace-in-place); a new id is appended — matching the old list order.
560 if (current.parts[part.id] == part) return false
561 current.parts[part.id] = part
562 current.invalidate()
563 } else {
564 val parts = LinkedHashMap<String, Part>()
565 parts[part.id] = part
566 messages[messageId] = Holder(
567 info = UnknownMessage(id = messageId, sessionID = part.sessionID ?: ""),
568 parts = parts,
569 )
570 }
571 // Record that this part streamed in, so a re-seed after a reconnect knows it's
572 // newer than the REST snapshot and keeps it (see seed()/beginReseed()).
573 streamedSincePivot.add(part.id)
574 if (current == null) evictOldMessages()
575 return true
576 }
578 private fun removePart(messageId: String, partId: String): Boolean {
579 val current = messages[messageId] ?: return false
580 if (current.parts.remove(partId) == null) return false
581 current.invalidate()
582 return true
583 }
585 /** Evict the oldest messages when the store exceeds [maxMessages], keeping memory bounded.
586 * Returns true if any messages were evicted. */
587 private fun evictOldMessages(): Boolean {
588 var evicted = false
589 while (messages.size > maxMessages) {
590 val oldestKey = messages.keys.iterator().next()
591 val removed = messages.remove(oldestKey)
592 // Prune the evicted message's pivot bookkeeping so these sets don't grow
593 // unbounded over a long-running session that never reconnects (they're only
594 // otherwise cleared on reseed/seed).
595 messageInfoUpdatedSincePivot.remove(oldestKey)
596 removed?.parts?.keys?.forEach { streamedSincePivot.remove(it) }
597 evicted = true
598 }
599 return evicted
600 }