app/src/main/java/soy/iko/opencode/data/repo/SessionRepository.kt · opencode-android

1package soy.iko.opencode.data.repo

2

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.channels.Channel

21import kotlinx.coroutines.flow.Flow

22import kotlinx.coroutines.flow.channelFlow

23import kotlinx.coroutines.flow.consumeAsFlow

24import kotlinx.coroutines.flow.conflate

25import kotlinx.coroutines.launch

26import kotlinx.coroutines.sync.Mutex

27import kotlinx.coroutines.sync.withLock

28

29/**

30 * Source of truth for a session's conversation. It seeds from `GET /session/:id/message`

31 * and then reduces [BusEvent]s from the SSE stream into an ordered list of messages,

32 * each keyed by message id with parts keyed by part id (idempotent upserts).

33 */

34open class SessionRepository(

35 private val api: OpencodeApiClient,

36 private val eventStream: EventStreamClient,

37) {

38 open suspend fun listSessions() = api.listSessions()

39 open suspend fun createSession(title: String? = null) = api.createSession(title)

40 open suspend fun deleteSession(id: String) = api.deleteSession(id)

41 open suspend fun abort(sessionId: String) = api.abort(sessionId)

42

43 open suspend fun sendPrompt(sessionId: String, text: String, model: ModelRef?, agent: String? = null) =

44 api.sendPrompt(sessionId, text, model, agent)

45

46 open suspend fun runCommand(

47 sessionId: String,

48 command: String,

49 arguments: String = "",

50 agent: String? = null,

51 ) = api.runCommand(sessionId, command, arguments = arguments, agent = agent)

52

53 /**

54 * A live, ordered view of [sessionId]'s messages. Begins collecting the event

55 * stream before the initial REST load so no streamed part is missed, then reduces

56 * events into the in-memory state. The flow is conflated so a fast burst of

57 * streaming tokens (each publishing a fresh snapshot) never back-pressures the

58 * event reducer — a slow collector simply sees the most recent snapshot.

59 */

60 open fun observeMessages(sessionId: String): Flow<List<MessageWithParts>> = channelFlow {

61 val store = MessageStore()

62 val lock = Mutex()

63 // Conflated dirty signal: a burst of reduce()/seed() calls coalesces into a

64 // single snapshot+send per demand cycle. Previously every changed event called

65 // snapshot() — allocating an O(N) copy of the whole message list — and ~97% of

66 // those copies were immediately discarded by the outer .conflate() during fast

67 // streaming (hundreds of tokens/sec). Marking dirty and snapshotting in one

68 // drain loop means N events in a burst allocate ~1 snapshot, not N.

69 val dirty = Channel<Unit>(Channel.CONFLATED)

70

71 // Mark the store dirty; the drain loop below snapshots and publishes. CONFLATED

72 // collapses repeated marks into a single pending signal, so a tight burst of

73 // events doesn't even allocate N marks.

74 fun publish() { dirty.trySend(Unit) }

75

76 // Subscribe to events first so we don't miss early deltas during the initial load.

77 launch {

78 eventStream.events.collect { event ->

79 if (lock.withLock { store.reduce(sessionId, event) }) publish()

80 }

81 }

82

83 // Re-seed from REST when the SSE stream reconnects after a drop. The REST fetch

84 // is performed *outside* the lock so events arriving via SSE during the fetch are

85 // not blocked. The seed itself is under the lock; a seed generation counter ensures

86 // a stale fetch (from an earlier reconnect) doesn't clobber a newer one.

87 launch {

88 // Re-seed from REST when the SSE stream reconnects after a drop. The

89 // hasConnectedBefore flag is NOT reset on disconnect — it stays true after

90 // the first successful connection so a subsequent reconnect triggers the

91 // re-seed. Resetting it on disconnect would make the condition always false

92 // after the first cycle, turning the re-seed into dead code.

93 //

94 // A seed generation counter guards against a stale re-seed clobbering a

95 // newer one: if two reconnects fire in quick succession, the older fetch

96 // (which returns later) is discarded because its generation no longer matches.

97 var hasConnectedBefore = false

98 var seedGeneration = 0

99 eventStream.state.collect { state ->

100 if (state == EventStreamClient.ConnectionState.Connected) {

101 if (hasConnectedBefore) {

102 val generation = ++seedGeneration

103 // Pivot before the fetch: parts already in memory predate this

104 // snapshot and must yield to it, while parts that stream in during

105 // the fetch (added to streamedSincePivot via reduce) outrank it.

106 lock.withLock { store.beginReseed() }

107 val fresh = runCatchingCancellable { api.listMessages(sessionId) }

108 .onFailure { Log.w("SessionRepository", "Re-seed message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }

109 .getOrDefault(emptyList())

110 var seedChanged = false

111 lock.withLock {

112 if (generation == seedGeneration) {

113 // Merge without pruning: SSE events arriving during the

114 // REST fetch may have added messages not yet indexed by

115 // the REST endpoint. Pruning would silently delete them.

116 // Stale messages from deletions during the disconnect

117 // gap are eventually removed via MessageRemoved events

118 // or on the next app restart.

119 seedChanged = store.seed(fresh, prune = false)

120 }

121 }

122 if (seedChanged) publish()

123 }

124 hasConnectedBefore = true

125 }

126 }

127 }

128

129 val initial = runCatchingCancellable { api.listMessages(sessionId) }

130 .onFailure { Log.w("SessionRepository", "Initial message load failed for $sessionId; relying on SSE: ${safeExceptionSummary(it)}") }

131 .getOrDefault(emptyList())

132 lock.withLock { store.seed(initial) }

133 publish()

134

135 // Single drain: one snapshot+send per conflated dirty signal. The snapshot is

136 // taken under the lock so a reader never sees a half-reduced state; send() runs

137 // outside the lock so a slow downstream collector can't stall the drain. This

138 // loop also keeps the flow alive until the collector cancels; the launched jobs

139 // above are children of this scope and are torn down with it.

140 dirty.consumeAsFlow().collect {

141 val snapshot = lock.withLock { store.snapshot() }

142 send(snapshot)

143 }

144 }.conflate()

145

146 companion object {

147 /** Convenience: is this event a run-completion signal for [sessionId]?

148 * A null sessionID is NOT treated as a wildcard — an idle event with no

149 * session id must not reset the running state of every open chat. */

150 fun isIdle(event: BusEvent, sessionId: String): Boolean =

151 event is SessionIdle && event.properties.sessionID == sessionId

152

153 fun isError(event: BusEvent, sessionId: String): Boolean =

154 event is SessionError && event.properties.sessionID == sessionId

155 }

156}

157

158/**

159 * In-memory reduction state for one observed session. Not thread-safe; guard with a Mutex.

160 * Exposed as `internal` so the reducer logic can be unit-tested directly.

161 */

162internal class MessageStore {

163 // messageId -> (info + parts), insertion-ordered.

164 private val messages = LinkedHashMap<String, MessageWithParts>()

165

166 /** Part ids upserted via [reduce] since the last reseed pivot ([beginReseed]) or seed.

167 * On a reconnect re-seed these are the parts that streamed in *during* the REST fetch,

168 * so they're newer than the snapshot and must win over it; every other part is taken

169 * from the authoritative REST snapshot. Without this distinction the merge kept every

170 * overlapping in-memory part, so a part that changed server-side during the disconnect

171 * (e.g. a tool going running -> completed) stayed stale until app restart. */

172 private val streamedSincePivot = mutableSetOf<String>()

173

174 /** Message ids whose `info` was updated via a live SSE [reduce] since the last reseed

175 * pivot ([beginReseed]) or seed. Parallels [streamedSincePivot] but for message info:

176 * on a reconnect re-seed these are the messages whose info changed *during* the REST

177 * fetch, so the in-memory info is newer than the snapshot and must win; every other

178 * message's info is taken from the authoritative REST snapshot. Without this the merge

179 * kept the in-memory info, so a run that finished server-side during the disconnect

180 * (e.g. cost/token totals, completion time) stayed stale until app restart. */

181 private val messageInfoUpdatedSincePivot = mutableSetOf<String>()

182

183 /** Mark the point a reconnect re-seed's REST fetch begins: parts/info streamed from here

184 * on are newer than the fetched snapshot. Call under the same lock as [seed]/[reduce]. */

185 fun beginReseed() {

186 streamedSincePivot.clear()

187 messageInfoUpdatedSincePivot.clear()

188 }

189

190 /** Monotonic counter for synthetic UnknownMessage keys, avoiding nanoTime collisions. */

191 private val unknownCounter = java.util.concurrent.atomic.AtomicLong(0)

192

193 /** Maximum number of messages to keep in memory; oldest are evicted when exceeded. */

194 internal val maxMessages = NetworkConfig.maxInMemoryMessages

195

196 fun snapshot(): List<MessageWithParts> = messages.values.toList()

197

198 fun seed(initial: List<MessageWithParts>, prune: Boolean = false): Boolean {

199 var changed = false

200 // On re-seed (after SSE reconnect), remove messages that are no longer in the

201 // server snapshot (e.g. deleted during the disconnection gap). This keeps the

202 // in-memory state in sync with the authoritative REST snapshot rather than

203 // accumulating stale messages forever. Skipped on the initial seed to avoid

204 // racing with just-arrived SSE events whose messages may not yet be in REST.

205 if (prune) {

206 val incomingIds = initial.mapTo(mutableSetOf()) { it.info.id }

207 val before = messages.size

208 messages.keys.retainAll(incomingIds)

209 if (messages.size != before) changed = true

210 }

211

212 for (m in initial) {

213 val existing = messages[m.info.id]

214 if (existing == null) {

215 messages[m.info.id] = m

216 changed = true

217 } else {

218 // A part streamed in between subscribe and this initial load may already

219 // have populated this message (see observeMessages). Merge instead of

220 // overwriting so that newer streamed part isn't discarded: take the

221 // snapshot's part order, swap in the streamed version where ids overlap,

222 // append any streamed-only parts, and adopt the authoritative REST info.

223 val streamedById = existing.parts.associateBy { it.id }

224 val snapshotIds = m.parts.mapTo(mutableSetOf()) { it.id }

225 // Prefer the in-memory version of an overlapping part ONLY if it streamed

226 // in since the reseed pivot (i.e. during this fetch), so it's newer than

227 // the snapshot. Otherwise the REST snapshot is authoritative and wins,

228 // discarding an in-memory part that went stale during a disconnect.

229 val ordered = m.parts.map { p ->

230 if (p.id in streamedSincePivot) streamedById[p.id] ?: p else p

231 }.toMutableList()

232 for (p in existing.parts) {

233 if (p.id !in snapshotIds) ordered.add(p)

234 }

235 // Adopt the authoritative REST info unless the in-memory info was itself

236 // updated live since the pivot (i.e. during this fetch), in which case it's

237 // newer than the snapshot. Otherwise the REST snapshot wins, discarding

238 // in-memory info that went stale during a disconnect.

239 val info = if (existing.info is UnknownMessage || m.info.id !in messageInfoUpdatedSincePivot) m.info else existing.info

240 val merged = MessageWithParts(info = info, parts = ordered)

241 if (merged != existing) {

242 messages[m.info.id] = merged

243 changed = true

244 }

245 }

246 }

247 if (evictOldMessages()) changed = true

248 // Reset the pivot: everything just merged is now the baseline, so only parts/info

249 // that stream in *after* this seed can outrank the next reconnect's snapshot.

250 streamedSincePivot.clear()

251 messageInfoUpdatedSincePivot.clear()

252 return changed

253 }

254

255 /** Returns true if the state changed (and a new snapshot should be published). */

256 fun reduce(sessionId: String, event: BusEvent): Boolean {

257 return when (event) {

258 is MessageUpdated -> {

259 val info = event.properties.info

260 if (info.sessionID != sessionId) {

261 false

262 } else {

263 // An UnknownMessage with an empty id (e.g. from an unrecognized

264 // server role with no id field) would collide with other such

265 // messages in the map. Generate a unique synthetic key so each

266 // unknown message gets its own entry instead of overwriting others.

267 val key = if (info.id.isEmpty() && info is UnknownMessage) {

268 "unknown-${unknownCounter.incrementAndGet()}"

269 } else {

270 info.id

271 }

272 val existing = messages[key]

273 messages[key] = existing?.copy(info = info) ?: MessageWithParts(info)

274 // Record that this message's info streamed in live, so a re-seed after a

275 // reconnect knows it's newer than the REST snapshot and keeps it (see

276 // seed()/beginReseed()).

277 messageInfoUpdatedSincePivot.add(key)

278 if (existing == null) evictOldMessages()

279 true

280 }

281 }

282

283 is MessagePartUpdated -> {

284 val part = event.properties.part

285 val messageId = part.messageID ?: event.properties.messageID

286 val partSession = part.sessionID ?: event.properties.sessionID

287 when {

288 messageId == null -> false

289 partSession != null -> if (partSession == sessionId) upsertPart(messageId, part) else false

290 // No session id anywhere on the event: a null session is NOT a wildcard

291 // (mirroring isIdle/isError). Every open session's store sees the shared

292 // event stream, so accepting it unconditionally would leak the part into

293 // unrelated conversations (e.g. both panes in two-pane mode). Attribute it

294 // only if this store already holds the parent message — MessageUpdated,

295 // which always carries a session id, will have created it here first.

296 else -> if (messages.containsKey(messageId)) upsertPart(messageId, part) else false

297 }

298 }

299

300 is MessagePartRemoved -> handlePartRemoved(sessionId, event)

301

302 is MessageRemoved -> handleMessageRemoved(sessionId, event)

303

304 else -> false

305 }

306 }

307

308 /** Returns false if [eventSession] is non-null and doesn't match [sessionId]. */

309 private fun matchesSession(eventSession: String?, sessionId: String): Boolean =

310 eventSession == null || eventSession == sessionId

311

312 private fun handlePartRemoved(sessionId: String, event: MessagePartRemoved): Boolean {

313 val messageId = event.properties.messageID ?: return false

314 val partId = event.properties.partID ?: return false

315 if (!matchesSession(event.properties.sessionID, sessionId)) return false

316 return removePart(messageId, partId)

317 }

318

319 private fun handleMessageRemoved(sessionId: String, event: MessageRemoved): Boolean {

320 val id = event.properties.messageID ?: return false

321 if (!matchesSession(event.properties.sessionID, sessionId)) return false

322 return messages.remove(id) != null

323 }

324

325 private fun upsertPart(messageId: String, part: Part): Boolean {

326 val current = messages[messageId]

327 val parts = current?.parts.orEmpty()

328 val idx = parts.indexOfFirst { it.id == part.id }

329 if (idx >= 0 && parts[idx] == part) return false

330 val newParts = if (idx >= 0) {

331 parts.toMutableList().also { it[idx] = part }

332 } else {

333 parts + part

334 }

335 messages[messageId] = current?.copy(parts = newParts)

336 ?: MessageWithParts(

337 info = UnknownMessage(id = messageId, sessionID = part.sessionID ?: ""),

338 parts = newParts,

339 )

340 // Record that this part streamed in, so a re-seed after a reconnect knows it's

341 // newer than the REST snapshot and keeps it (see seed()/beginReseed()).

342 streamedSincePivot.add(part.id)

343 if (current == null) evictOldMessages()

344 return true

345 }

346

347 private fun removePart(messageId: String, partId: String): Boolean {

348 val current = messages[messageId] ?: return false

349 val newParts = current.parts.filterNot { it.id == partId }

350 if (newParts.size == current.parts.size) return false

351 messages[messageId] = current.copy(parts = newParts)

352 return true

353 }

354

355 /** Evict the oldest messages when the store exceeds [maxMessages], keeping memory bounded.

356 * Returns true if any messages were evicted. */

357 private fun evictOldMessages(): Boolean {

358 var evicted = false

359 while (messages.size > maxMessages) {

360 val oldestKey = messages.keys.iterator().next()

361 messages.remove(oldestKey)

362 evicted = true

363 }

364 return evicted

365 }

366}