git.iko.soy/opencode-android
on 61cdc75 help sign in
opencode-android/app/src/main/java/soy/iko/opencode/data/repo/DraftStore.kt
260 lines 13 kB history raw
1package soy.iko.opencode.data.repo
3import android.content.Context
4import kotlinx.coroutines.CoroutineScope
5import kotlinx.coroutines.Dispatchers
6import kotlinx.coroutines.flow.MutableStateFlow
7import kotlinx.coroutines.flow.StateFlow
8import kotlinx.coroutines.flow.asStateFlow
9import kotlinx.coroutines.flow.update
10import android.util.Log
11import kotlinx.coroutines.launch
12import java.util.concurrent.Executors
14/**
15 * Persists the in-progress message draft for each session so it survives
16 * back-navigation and process death. Backed by SharedPreferences because the
17 * UI needs the initial value synchronously and drafts are low-stakes data.
18 *
19 * The prefs file is loaded on a background thread at construction time. Until
20 * the load completes, [drafts] reports an empty map and [get] returns "" rather
21 * than blocking the caller (previously a [CountDownLatch] could stall the main
22 * thread for up to 2 seconds on a slow disk).
23 */
24open class DraftStore private constructor(
25 private val appContext: Context?,
26 private val scope: CoroutineScope?,
27 @Suppress("unused") private val testMode: Boolean,
28) {
29 constructor(context: Context, scope: CoroutineScope) : this(context.applicationContext, scope, false)
30 protected constructor() : this(null, null, true)
32 private val prefs by lazy {
33 appContext?.getSharedPreferences("drafts", Context.MODE_PRIVATE)
34 ?: error("No context — override methods in test subclass")
35 }
37 /** Single-thread executor for flushDraft when prefs aren't initialized yet,
38 * so the main thread is never blocked on disk I/O. */
39 private val flushExecutor by lazy {
40 Executors.newSingleThreadExecutor { r ->
41 Thread(r, "DraftStore-flush").apply { isDaemon = true }
42 }
43 }
45 private val _drafts = MutableStateFlow<Map<String, String>>(emptyMap())
46 /** The full in-memory snapshot of all drafts, updated when prefs finish loading. */
47 open val drafts: StateFlow<Map<String, String>> = _drafts.asStateFlow()
49 /** Queued follow-up messages per session, persisted under a namespaced key so a
50 * follow-up queued mid-run survives process death (the in-memory-only StateFlow in
51 * ChatViewModel would otherwise lose it, along with the cleared draft). */
52 private val _followUps = MutableStateFlow<Map<String, String>>(emptyMap())
54 private val _ready = MutableStateFlow(false)
55 /** True once the background prefs load has completed. */
56 open val ready: StateFlow<Boolean> = _ready.asStateFlow()
58 /** Keys written in-memory (set / cleared / removed) *before* the async prefs load
59 * finished. For these the in-memory state is authoritative — the load must not
60 * resurrect a persisted value the user has since cleared, which a plain putAll(current)
61 * merge cannot express (it can add/override, but not delete). Consulted once, in the
62 * init merge. Separate sets because drafts and follow-ups share the session-id key space
63 * but live under different prefs keys. */
64 // ConcurrentHashMap-backed sets: weakly-consistent iteration so the init merge's forEach
65 // can't throw ConcurrentModificationException if a mutator adds a key concurrently.
66 private val touchedDraftKeys: MutableSet<String> = java.util.concurrent.ConcurrentHashMap.newKeySet()
67 private val touchedFollowUpKeys: MutableSet<String> = java.util.concurrent.ConcurrentHashMap.newKeySet()
69 private fun markDraftTouched(sessionId: String) {
70 if (!_ready.value) touchedDraftKeys.add(sessionId)
71 }
72 private fun markFollowUpTouched(sessionId: String) {
73 if (!_ready.value) touchedFollowUpKeys.add(sessionId)
74 }
76 init {
77 // Eagerly load the prefs file on a background thread so the first main-thread
78 // get() doesn't block on disk I/O. The loaded snapshot is published via [_drafts]
79 // so callers can observe readiness without blocking.
80 if (scope != null && appContext != null) {
81 scope.launch(Dispatchers.IO) {
82 val all = runCatching {
83 prefs.all.mapValues { it.value as? String ?: "" }
84 }.getOrDefault(emptyMap())
85 // Partition persisted keys: follow-up keys carry a prefix, everything else
86 // is a plain per-session draft.
87 val drafts = mutableMapOf<String, String>()
88 val followUps = mutableMapOf<String, String>()
89 all.forEach { (k, v) ->
90 if (k.startsWith(FOLLOWUP_PREFIX)) followUps[k.removePrefix(FOLLOWUP_PREFIX)] = v
91 else drafts[k] = v
92 }
93 // Merge instead of overwriting: any in-memory writes that arrived before
94 // the background load completed (e.g. a share-intent draft injected via
95 // setImmediate) take priority over the persisted snapshot so they aren't
96 // clobbered. For keys the user explicitly touched before the load finished,
97 // drop the persisted entry first so a value they *cleared* isn't resurrected
98 // (the cleared key is absent from `current`, so putAll alone can't remove it).
99 _drafts.update { current ->
100 drafts.apply {
101 touchedDraftKeys.forEach { remove(it) }
102 putAll(current)
103 }
104 }
105 _followUps.update { current ->
106 followUps.apply {
107 touchedFollowUpKeys.forEach { remove(it) }
108 putAll(current)
109 }
110 }
111 _ready.value = true
112 }
113 } else {
114 _ready.value = true
115 }
116 }
118 /** Returns the draft for [sessionId], or "" if not found or still loading. */
119 open fun get(sessionId: String): String = _drafts.value[sessionId].orEmpty()
121 /** Returns the queued follow-up for [sessionId], or "" if none / still loading. */
122 open fun getFollowUp(sessionId: String): String = _followUps.value[sessionId].orEmpty()
124 /** Persist a queued follow-up immediately (fire-and-forget, main-thread-safe).
125 * Mirrors [flushDraft]: async apply() when prefs are open, else a background thread,
126 * so a follow-up queued just before process death is flushed before the process exits.
127 * A blank [text] clears the stored follow-up. */
128 open fun flushFollowUp(sessionId: String, text: String) {
129 markFollowUpTouched(sessionId)
130 _followUps.update { current ->
131 val updated = current.toMutableMap()
132 if (text.isBlank()) updated.remove(sessionId) else updated[sessionId] = text
133 updated
134 }
135 writeAsync(FOLLOWUP_PREFIX + sessionId, text)
136 }
138 /** Persist [key]=[text] (or remove it when blank) without blocking the main thread on
139 * first-time disk I/O — see [flushDraft] for the prefs-not-yet-open rationale. */
140 private fun writeAsync(key: String, text: String) {
141 val write = {
142 runCatching {
143 prefs.edit().apply {
144 if (text.isBlank()) remove(key) else putString(key, text)
145 }.apply()
146 }.onFailure { Log.w("DraftStore", "Failed to persist $key", it) }
147 }
148 // FIX: always submit to the single-thread flushExecutor (never apply directly) so
149 // writes persist in submission order. Previously the post-init fast path applied
150 // synchronously on the caller's thread, which could land BEFORE an earlier still-queued
151 // write, persisting stale text. The task resolves the prefs lazy and calls apply().
152 runCatching { flushExecutor.execute { write() } }
153 }
155 /** Update the in-memory draft immediately without persisting to disk.
156 * Used when navigation must see the new value synchronously (e.g. share-intent
157 * injection before ChatScreen composes). Call [set] afterwards to persist. */
158 open fun setImmediate(sessionId: String, text: String) {
159 markDraftTouched(sessionId)
160 _drafts.update { current ->
161 val updated = current.toMutableMap()
162 if (text.isBlank()) updated.remove(sessionId) else updated[sessionId] = text
163 updated
164 }
165 }
167 open suspend fun set(sessionId: String, text: String) {
168 markDraftTouched(sessionId)
169 _drafts.update { current ->
170 val updated = current.toMutableMap()
171 if (text.isBlank()) updated.remove(sessionId) else updated[sessionId] = text
172 updated
173 }
174 // Route persistence through the single-thread flushExecutor (like the other writers)
175 // rather than the shared IO pool, so this write can't race ahead of an earlier queued
176 // flush and land stale. writeAsync persists sessionId=text (or removes it when blank),
177 // which is exactly this draft's persistence.
178 writeAsync(sessionId, text)
179 }
181 /** Remove the draft for a session (call on session deletion to avoid orphaned entries).
182 * Also clears any queued follow-up so a deleted session leaves nothing behind. */
183 open suspend fun remove(sessionId: String) {
184 markDraftTouched(sessionId)
185 markFollowUpTouched(sessionId)
186 _drafts.update { current ->
187 val updated = current.toMutableMap()
188 updated.remove(sessionId)
189 updated
190 }
191 _followUps.update { current ->
192 val updated = current.toMutableMap()
193 updated.remove(sessionId)
194 updated
195 }
196 // Submit the (combined draft + follow-up) removal to the same single-thread
197 // flushExecutor the other writers use, so it can't race ahead of an earlier queued
198 // write and resurrect stale text. Mirrors flushDraft's submission pattern.
199 runCatching {
200 flushExecutor.execute {
201 runCatching {
202 prefs.edit().remove(sessionId).remove(FOLLOWUP_PREFIX + sessionId).apply()
203 }.onFailure { Log.w("DraftStore", "Failed to remove draft for $sessionId", it) }
204 }
205 }
206 }
208 /**
209 * Synchronously persist a draft. Intended for [ViewModel.onCleared] where the
210 * viewModelScope is already cancelled and a suspending write can't be used.
211 *
212 * If the SharedPreferences file has already been opened (the normal case — the
213 * background init load touches it shortly after construction), the write uses
214 * [apply] (asynchronous) and returns immediately without blocking the main thread.
215 *
216 * If the prefs file has NOT been opened yet (a rare race where onCleared fires
217 * before the background init), the write is dispatched to a background thread
218 * instead of blocking the main thread on a synchronous disk read. The draft
219 * may be lost if the process exits before the write completes, but this is
220 * preferable to an ANR.
221 */
222 open fun flushDraft(sessionId: String, text: String) {
223 markDraftTouched(sessionId)
224 _drafts.update { current ->
225 val updated = current.toMutableMap()
226 if (text.isBlank()) updated.remove(sessionId) else updated[sessionId] = text
227 updated
228 }
229 // FIX: always route the persistence through the single-thread flushExecutor so writes
230 // land in submission order. Previously, once prefs were open this applied synchronously
231 // on the calling (main) thread, so a newer direct write could persist BEFORE an older
232 // still-queued write completed, clobbering the draft with stale text. The executor task
233 // resolves the prefs lazy and uses apply() (async) exactly as before; the write may be
234 // lost only if the process exits before the queued task runs (as previously).
235 runCatching {
236 flushExecutor.execute {
237 runCatching {
238 prefs.edit().apply {
239 if (text.isBlank()) remove(sessionId) else putString(sessionId, text)
240 }.apply()
241 }.onFailure { Log.w("DraftStore", "Failed to flush draft for $sessionId", it) }
242 }
243 }
244 }
246 private companion object {
247 /** Key prefix distinguishing persisted follow-ups from plain drafts in the shared
248 * prefs file. The leading/trailing spaces can't collide with a real opencode session
249 * id (which are non-blank, space-free tokens like "ses_..."). */
250 const val FOLLOWUP_PREFIX = " followup "
251 }
253 /** Shut down the background flush executor. Call from AppContainer.shutdown(). */
254 open fun shutdown() {
255 flushExecutor.shutdown()
256 // Await pending flushDraft writes so drafts aren't lost if the process exits
257 // immediately after shutdown (e.g. ANR-triggered process kill).
258 runCatching { flushExecutor.awaitTermination(2, java.util.concurrent.TimeUnit.SECONDS) }
259 }