git.iko.soy/opencode-android
on 325e359 help sign in
opencode-android/app/src/main/java/soy/iko/opencode/data/repo/SessionRepository.kt
284 lines 14 kB history raw
1package soy.iko.opencode.data.repo
3import soy.iko.opencode.data.model.BusEvent
4import soy.iko.opencode.data.model.MessagePartRemoved
5import soy.iko.opencode.data.model.MessagePartUpdated
6import soy.iko.opencode.data.model.MessageRemoved
7import soy.iko.opencode.data.model.MessageUpdated
8import soy.iko.opencode.data.model.MessageWithParts
9import soy.iko.opencode.data.model.ModelRef
10import soy.iko.opencode.data.model.Part
11import soy.iko.opencode.data.model.SessionError
12import soy.iko.opencode.data.model.SessionIdle
13import soy.iko.opencode.data.model.UnknownMessage
14import soy.iko.opencode.data.network.EventStreamClient
15import soy.iko.opencode.data.network.NetworkConfig
16import soy.iko.opencode.data.network.OpencodeApiClient
17import soy.iko.opencode.util.runCatchingCancellable
18import soy.iko.opencode.util.safeExceptionSummary
19import android.util.Log
20import kotlinx.coroutines.flow.Flow
21import kotlinx.coroutines.flow.channelFlow
22import kotlinx.coroutines.flow.conflate
23import kotlinx.coroutines.launch
24import kotlinx.coroutines.sync.Mutex
25import kotlinx.coroutines.sync.withLock
27/**
28 * Source of truth for a session's conversation. It seeds from `GET /session/:id/message`
29 * and then reduces [BusEvent]s from the SSE stream into an ordered list of messages,
30 * each keyed by message id with parts keyed by part id (idempotent upserts).
31 */
32open class SessionRepository(
33 private val api: OpencodeApiClient,
34 private val eventStream: EventStreamClient,
35) {
36 open suspend fun listSessions() = api.listSessions()
37 open suspend fun createSession(title: String? = null) = api.createSession(title)
38 open suspend fun deleteSession(id: String) = api.deleteSession(id)
39 open suspend fun abort(sessionId: String) = api.abort(sessionId)
41 open suspend fun sendPrompt(sessionId: String, text: String, model: ModelRef?, agent: String? = null) =
42 api.sendPrompt(sessionId, text, model, agent)
44 open suspend fun runCommand(
45 sessionId: String,
46 command: String,
47 arguments: String = "",
48 agent: String? = null,
49 ) = api.runCommand(sessionId, command, arguments = arguments, agent = agent)
51 /**
52 * A live, ordered view of [sessionId]'s messages. Begins collecting the event
53 * stream before the initial REST load so no streamed part is missed, then reduces
54 * events into the in-memory state. The flow is conflated so a fast burst of
55 * streaming tokens (each publishing a fresh snapshot) never back-pressures the
56 * event reducer — a slow collector simply sees the most recent snapshot.
57 */
58 open fun observeMessages(sessionId: String): Flow<List<MessageWithParts>> = channelFlow {
59 val store = MessageStore()
60 val lock = Mutex()
62 // Collect the snapshot under the lock, then publish outside it so a slow
63 // downstream collector can't stall event processing.
64 suspend fun publish() {
65 val snapshot = lock.withLock { store.snapshot() }
66 send(snapshot)
67 }
69 // Subscribe to events first so we don't miss early deltas during the initial load.
70 val job = launch {
71 eventStream.events.collect { event ->
72 val changed = lock.withLock { store.reduce(sessionId, event) }
73 if (changed) publish()
74 }
75 }
77 // Re-seed from REST when the SSE stream reconnects after a drop. The REST fetch
78 // is performed *outside* the lock so events arriving via SSE during the fetch are
79 // not blocked. The seed itself is under the lock; a seed generation counter ensures
80 // a stale fetch (from an earlier reconnect) doesn't clobber a newer one.
81 launch {
82 // Re-seed from REST when the SSE stream reconnects after a drop. The
83 // hasConnectedBefore flag is NOT reset on disconnect — it stays true after
84 // the first successful connection so a subsequent reconnect triggers the
85 // re-seed. Resetting it on disconnect would make the condition always false
86 // after the first cycle, turning the re-seed into dead code.
87 //
88 // A seed generation counter guards against a stale re-seed clobbering a
89 // newer one: if two reconnects fire in quick succession, the older fetch
90 // (which returns later) is discarded because its generation no longer matches.
91 var hasConnectedBefore = false
92 var seedGeneration = 0
93 eventStream.state.collect { state ->
94 if (state == EventStreamClient.ConnectionState.Connected) {
95 if (hasConnectedBefore) {
96 val generation = ++seedGeneration
97 val fresh = runCatchingCancellable { api.listMessages(sessionId) }
98 .onFailure { Log.w("SessionRepository", "Re-seed message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }
99 .getOrDefault(emptyList())
100 lock.withLock {
101 if (generation == seedGeneration) {
102 // Merge without pruning: SSE events arriving during the
103 // REST fetch may have added messages not yet indexed by
104 // the REST endpoint. Pruning would silently delete them.
105 // Stale messages from deletions during the disconnect
106 // gap are eventually removed via MessageRemoved events
107 // or on the next app restart.
108 store.seed(fresh, prune = false)
109 }
110 }
111 publish()
112 }
113 hasConnectedBefore = true
114 }
115 }
116 }
118 val initial = runCatchingCancellable { api.listMessages(sessionId) }
119 .onFailure { Log.w("SessionRepository", "Initial message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }
120 .getOrDefault(emptyList())
121 lock.withLock { store.seed(initial) }
122 publish()
124 // Keep the flow alive until the collector cancels; the launched job is torn down with it.
125 job.join()
126 }.conflate()
128 companion object {
129 /** Convenience: is this event a run-completion signal for [sessionId]?
130 * A null sessionID is NOT treated as a wildcard — an idle event with no
131 * session id must not reset the running state of every open chat. */
132 fun isIdle(event: BusEvent, sessionId: String): Boolean =
133 event is SessionIdle && event.properties.sessionID == sessionId
135 fun isError(event: BusEvent, sessionId: String): Boolean =
136 event is SessionError && event.properties.sessionID == sessionId
137 }
140/**
141 * In-memory reduction state for one observed session. Not thread-safe; guard with a Mutex.
142 * Exposed as `internal` so the reducer logic can be unit-tested directly.
143 */
144internal class MessageStore {
145 // messageId -> (info + parts), insertion-ordered.
146 private val messages = LinkedHashMap<String, MessageWithParts>()
148 /** Monotonic counter for synthetic UnknownMessage keys, avoiding nanoTime collisions. */
149 private val unknownCounter = java.util.concurrent.atomic.AtomicLong(0)
151 /** Maximum number of messages to keep in memory; oldest are evicted when exceeded. */
152 internal val maxMessages = NetworkConfig.maxInMemoryMessages
154 fun snapshot(): List<MessageWithParts> = messages.values.toList()
156 fun seed(initial: List<MessageWithParts>, prune: Boolean = false) {
157 // On re-seed (after SSE reconnect), remove messages that are no longer in the
158 // server snapshot (e.g. deleted during the disconnection gap). This keeps the
159 // in-memory state in sync with the authoritative REST snapshot rather than
160 // accumulating stale messages forever. Skipped on the initial seed to avoid
161 // racing with just-arrived SSE events whose messages may not yet be in REST.
162 if (prune) {
163 val incomingIds = initial.mapTo(mutableSetOf()) { it.info.id }
164 messages.keys.retainAll(incomingIds)
165 }
167 for (m in initial) {
168 val existing = messages[m.info.id]
169 if (existing == null) {
170 messages[m.info.id] = m
171 } else {
172 // A part streamed in between subscribe and this initial load may already
173 // have populated this message (see observeMessages). Merge instead of
174 // overwriting so that newer streamed part isn't discarded: take the
175 // snapshot's part order, swap in the streamed version where ids overlap,
176 // append any streamed-only parts, and adopt the authoritative REST info.
177 val streamedById = existing.parts.associateBy { it.id }
178 val snapshotIds = m.parts.mapTo(mutableSetOf()) { it.id }
179 val ordered = m.parts.map { p -> streamedById[p.id] ?: p }.toMutableList()
180 for (p in existing.parts) {
181 if (p.id !in snapshotIds) ordered.add(p)
182 }
183 val info = if (existing.info is UnknownMessage) m.info else existing.info
184 messages[m.info.id] = MessageWithParts(info = info, parts = ordered)
185 }
186 }
187 evictOldMessages()
188 }
190 /** Returns true if the state changed (and a new snapshot should be published). */
191 fun reduce(sessionId: String, event: BusEvent): Boolean {
192 return when (event) {
193 is MessageUpdated -> {
194 val info = event.properties.info
195 if (info.sessionID != sessionId) {
196 false
197 } else {
198 // An UnknownMessage with an empty id (e.g. from an unrecognized
199 // server role with no id field) would collide with other such
200 // messages in the map. Generate a unique synthetic key so each
201 // unknown message gets its own entry instead of overwriting others.
202 val key = if (info.id.isEmpty() && info is UnknownMessage) {
203 "unknown-${unknownCounter.incrementAndGet()}"
204 } else {
205 info.id
206 }
207 val existing = messages[key]
208 messages[key] = existing?.copy(info = info) ?: MessageWithParts(info)
209 if (existing == null) evictOldMessages()
210 true
211 }
212 }
214 is MessagePartUpdated -> {
215 val part = event.properties.part
216 val messageId = part.messageID ?: event.properties.messageID
217 val partSession = part.sessionID ?: event.properties.sessionID
218 if (messageId == null || (partSession != null && partSession != sessionId)) {
219 false
220 } else {
221 upsertPart(messageId, part)
222 }
223 }
225 is MessagePartRemoved -> handlePartRemoved(sessionId, event)
227 is MessageRemoved -> handleMessageRemoved(sessionId, event)
229 else -> false
230 }
231 }
233 /** Returns false if [eventSession] is non-null and doesn't match [sessionId]. */
234 private fun matchesSession(eventSession: String?, sessionId: String): Boolean =
235 eventSession == null || eventSession == sessionId
237 private fun handlePartRemoved(sessionId: String, event: MessagePartRemoved): Boolean {
238 val messageId = event.properties.messageID ?: return false
239 val partId = event.properties.partID ?: return false
240 if (!matchesSession(event.properties.sessionID, sessionId)) return false
241 return removePart(messageId, partId)
242 }
244 private fun handleMessageRemoved(sessionId: String, event: MessageRemoved): Boolean {
245 val id = event.properties.messageID ?: return false
246 if (!matchesSession(event.properties.sessionID, sessionId)) return false
247 return messages.remove(id) != null
248 }
250 private fun upsertPart(messageId: String, part: Part): Boolean {
251 val current = messages[messageId]
252 val parts = current?.parts.orEmpty()
253 val idx = parts.indexOfFirst { it.id == part.id }
254 if (idx >= 0 && parts[idx] == part) return false
255 val newParts = if (idx >= 0) {
256 parts.toMutableList().also { it[idx] = part }
257 } else {
258 parts + part
259 }
260 messages[messageId] = current?.copy(parts = newParts)
261 ?: MessageWithParts(
262 info = UnknownMessage(id = messageId, sessionID = part.sessionID ?: ""),
263 parts = newParts,
264 )
265 if (current == null) evictOldMessages()
266 return true
267 }
269 private fun removePart(messageId: String, partId: String): Boolean {
270 val current = messages[messageId] ?: return false
271 val newParts = current.parts.filterNot { it.id == partId }
272 if (newParts.size == current.parts.size) return false
273 messages[messageId] = current.copy(parts = newParts)
274 return true
275 }
277 /** Evict the oldest messages when the store exceeds [maxMessages], keeping memory bounded. */
278 private fun evictOldMessages() {
279 while (messages.size > maxMessages) {
280 val oldestKey = messages.keys.iterator().next()
281 messages.remove(oldestKey)
282 }
283 }