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