@@ -20,6 +20,7 @@ import soy.iko.opencode.data.network.OpencodeApiClient
2020import soy.iko.opencode.util.runCatchingCancellable
2121import soy.iko.opencode.util.safeExceptionSummary
2222import android.util.Log
23+import kotlinx.coroutines.CoroutineScope
2324import kotlinx.coroutines.Dispatchers
2425import kotlinx.coroutines.Job
2526import kotlinx.coroutines.NonCancellable
@@ -46,6 +47,9 @@ open class SessionRepository(
4647 // Optional on-disk cache so a conversation renders instantly on open and stays readable
4748 // offline. Null in tests (and any path that doesn't wire it), where caching is skipped.
4849 private val messageCache: MessageCacheStore? = null,
50+ // Profile the cache is namespaced under, so two servers with a colliding session id don't
51+ // share a cache file. Empty in tests / when no cache is wired.
52+ private val profileId: String = "",
4953) {
5054 open suspend fun listSessions() = api.listSessions()
5155 open suspend fun createSession(title: String? = null, directory: String? = null) =
@@ -78,6 +82,58 @@ open class SessionRepository(
7882 * streaming tokens (each publishing a fresh snapshot) never back-pressures the
7983 * event reducer — a slow collector simply sees the most recent snapshot.
8084 */
85+ /** Per-subscription helper that persists reduced snapshots to [messageCache], throttled so a
86+ * fast token stream doesn't hammer the disk, plus a NonCancellable teardown flush. Extracted
87+ * from [observeMessages] so this cache-write branching lives in its own methods. [scope] is
88+ * the channelFlow scope, so scheduled saves are torn down with the flow — the teardown flush
89+ * re-persists a trailing write the cancellation dropped. */
90+ private inner class CacheWriter(private val scope: CoroutineScope, private val sessionId: String) {
91+ private var lastWrite = 0L
92+ private var pending: Job? = null
93+ // Most recent snapshot handed to [schedule], and the most recent one whose on-disk write
94+ // completed. `saved` is assigned from inside the launched save *after* it returns, so a
95+ // save cancelled mid-IO by teardown is NOT marked saved and [flushOnTeardown] re-persists
96+ // it. Written from the save coroutines, read from the collector/teardown; the only race is
97+ // a stale read causing one redundant, idempotent re-save (atomic temp+rename makes a
98+ // duplicate write of identical data harmless) — never data loss.
99+ private var latest: List<MessageWithParts>? = null
100+ private var saved: List<MessageWithParts>? = null
101+
102+ fun schedule(snapshot: List<MessageWithParts>) {
103+ val cache = messageCache ?: return
104+ if (snapshot.isEmpty()) return
105+ latest = snapshot
106+ val now = System.currentTimeMillis()
107+ pending?.cancel()
108+ // Immediate write once the throttle window has elapsed; otherwise a trailing-edge
109+ // write so a burst that ends inside the window (a run's final tokens right after a
110+ // write) still persists its last snapshot. A newer snapshot cancels and replaces a
111+ // pending trailing write, so continuous streaming coalesces to one write per window.
112+ val delayMs = if (now - lastWrite >= NetworkConfig.messageCacheWriteThrottleMs) {
113+ lastWrite = now
114+ 0L
115+ } else {
116+ NetworkConfig.messageCacheWriteThrottleMs
117+ }
118+ pending = scope.launch {
119+ if (delayMs > 0) delay(delayMs)
120+ cache.save(profileId, sessionId, snapshot)
121+ saved = snapshot
122+ }
123+ }
124+
125+ suspend fun flushOnTeardown() {
126+ // Scheduled saves are children of [scope], so teardown cancels a pending trailing
127+ // write before its delay elapses — a run's final tokens would never reach disk. Flush
128+ // the latest un-saved snapshot under NonCancellable so it survives the cancellation.
129+ val cache = messageCache ?: return
130+ val pendingSnapshot = latest ?: return
131+ if (pendingSnapshot !== saved) {
132+ withContext(NonCancellable) { cache.save(profileId, sessionId, pendingSnapshot) }
133+ }
134+ }
135+ }
136+
81137 open fun observeMessages(sessionId: String): Flow<List<MessageWithParts>> = channelFlow {
82138 val store = MessageStore()
83139 val lock = Mutex()
@@ -118,12 +174,12 @@ open class SessionRepository(
118174 // reconnect re-seed that somehow raced ahead isn't clobbered; the initial REST seed
119175 // (also generation 0) then overrides this with authoritative data when it succeeds.
120176 messageCache?.let { cache ->
121− val cached = cache.load(sessionId)
177+ val cached = cache.load(profileId, sessionId)
122178 if (cached.isNotEmpty()) {
123179 val changed = lock.withLock {
124180 if (seedGeneration == 0) {
125181 cacheSeededIds = cached.mapTo(mutableSetOf()) { it.info.id }
126− store.seed(cached)
182+ store.seed(cached, clearPivot = false)
127183 } else false
128184 }
129185 if (changed) publish()
@@ -212,18 +268,10 @@ open class SessionRepository(
212268 // outside the lock so a slow downstream collector can't stall the drain. This
213269 // loop also keeps the flow alive until the collector cancels; the launched jobs
214270 // 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
271+ // Per-subscription cache writer: throttled snapshot persistence + a NonCancellable
272+ // teardown flush. `this` is the channelFlow scope, so its scheduled saves are torn down
273+ // with the flow (the teardown flush re-persists a cancelled trailing write).
274+ val cacheWriter = CacheWriter(this, sessionId)
227275 // Guarantee at least one emission even when nothing seeds the store: with no on-disk
228276 // cache, a failing initial REST load (swallowed above), and a quiet SSE stream, none of
229277 // the seed paths publish, so the flow would emit nothing at all. The sole collector
@@ -236,46 +284,13 @@ open class SessionRepository(
236284 dirty.consumeAsFlow().collect {
237285 val snapshot = lock.withLock { store.snapshot() }
238286 send(snapshot)
239− // Persist the snapshot to the on-disk cache, throttled so a fast token stream
240− // doesn't hammer the disk — the cache only needs to be "recent enough" for an
241− // instant/offline first paint; the network corrects it on the next open.
242− if (messageCache != null && snapshot.isNotEmpty()) {
243− latestSnapshot = snapshot
244− val now = System.currentTimeMillis()
245− pendingCacheSave?.cancel()
246− if (now - lastCacheWrite >= NetworkConfig.messageCacheWriteThrottleMs) {
247− lastCacheWrite = now
248− pendingCacheSave = launch {
249− messageCache.save(sessionId, snapshot)
250− savedSnapshot = snapshot
251− }
252− } else {
253− // Trailing-edge write: a burst that ends inside the throttle window (a run's
254− // final tokens right after a write) would otherwise never persist its last
255− // snapshot, so an offline reopen would show the conversation missing its tail.
256− // Write the latest snapshot once the window elapses; a newer snapshot cancels
257− // and replaces this, so continuous streaming still coalesces to one write/window.
258− pendingCacheSave = launch {
259− delay(NetworkConfig.messageCacheWriteThrottleMs)
260− messageCache.save(sessionId, snapshot)
261− savedSnapshot = snapshot
262− }
263− }
264− }
287+ // Persist to the on-disk cache, throttled so a fast token stream doesn't hammer
288+ // the disk — the cache only needs to be "recent enough" for an instant/offline
289+ // first paint; the network corrects it on the next open.
290+ cacheWriter.schedule(snapshot)
265291 }
266292 } finally {
267− // FIX: both cache-saves above run as children of this inner flow scope, so when
268− // ChatViewModel's flatMapLatest switches (an SSE reconnect changing activeConnection)
269− // or onCleared cancels this flow, a pending trailing save is torn down before its delay
270− // elapses — a run's final tokens then never reach the on-disk cache. Flush the latest
271− // un-persisted snapshot here under NonCancellable so it survives the cancellation.
272− // savedSnapshot tracks the snapshot an immediate write already handled, so the common
273− // case (nothing pending) writes nothing extra — at most one write per teardown, never
274− // one per emission.
275− val pending = latestSnapshot
276− if (messageCache != null && pending != null && pending !== savedSnapshot) {
277− withContext(NonCancellable) { messageCache.save(sessionId, pending) }
278− }
293+ cacheWriter.flushOnTeardown()
279294 }
280295 // Run the whole reducer pipeline — event reduction, list rebuilds, and the O(N)
281296 // snapshot copy — on Dispatchers.Default rather than the collector's context. The
@@ -380,7 +395,7 @@ internal class MessageStore {
380395
381396 fun snapshot(): List<MessageWithParts> = messages.values.map { it.view() }
382397
383− fun seed(initial: List<MessageWithParts>, prune: Boolean = false): Boolean {
398+ fun seed(initial: List<MessageWithParts>, prune: Boolean = false, clearPivot: Boolean = true): Boolean {
384399 var changed = false
385400 // On re-seed (after SSE reconnect), remove messages that are no longer in the
386401 // server snapshot (e.g. deleted during the disconnection gap). This keeps the
@@ -439,9 +454,14 @@ internal class MessageStore {
439454 if (reorderByTime()) changed = true
440455 if (evictOldMessages()) changed = true
441456 // Reset the pivot: everything just merged is now the baseline, so only parts/info
442− // that stream in *after* this seed can outrank the next reconnect's snapshot.
443− streamedSincePivot.clear()
444− messageInfoUpdatedSincePivot.clear()
457+ // that stream in *after* this seed can outrank the next reconnect's snapshot. The
458+ // on-disk cache seed passes clearPivot = false: it runs *before* the authoritative
459+ // initial REST seed, so clearing here would strip pivot protection from SSE parts/info
460+ // that streamed in since subscribe, letting that REST seed regress them.
461+ if (clearPivot) {
462+ streamedSincePivot.clear()
463+ messageInfoUpdatedSincePivot.clear()
464+ }
445465 return changed
446466 }
447467