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