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 val fresh = runCatchingCancellable { api.listMessages(sessionId) }

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

105 .getOrDefault(emptyList())

106 var seedChanged = false

107 lock.withLock {

108 if (generation == seedGeneration) {

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

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

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

112 // Stale messages from deletions during the disconnect

113 // gap are eventually removed via MessageRemoved events

114 // or on the next app restart.

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

116 }

117 }

118 if (seedChanged) publish()

119 }

120 hasConnectedBefore = true

121 }

122 }

123 }

124

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

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

127 .getOrDefault(emptyList())

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

129 publish()

130

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

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

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

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

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

136 dirty.consumeAsFlow().collect {

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

138 send(snapshot)

139 }

140 }.conflate()

141

142 companion object {

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

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

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

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

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

148

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

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

151 }

152}

153

154/**

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

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

157 */

158internal class MessageStore {

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

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

161

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

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

164

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

166 internal val maxMessages = NetworkConfig.maxInMemoryMessages

167

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

169

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

171 var changed = false

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

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

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

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

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

177 if (prune) {

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

179 val before = messages.size

180 messages.keys.retainAll(incomingIds)

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

182 }

183

184 for (m in initial) {

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

186 if (existing == null) {

187 messages[m.info.id] = m

188 changed = true

189 } else {

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

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

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

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

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

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

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

197 val ordered = m.parts.map { p -> streamedById[p.id] ?: p }.toMutableList()

198 for (p in existing.parts) {

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

200 }

201 val info = if (existing.info is UnknownMessage) m.info else existing.info

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

203 if (merged != existing) {

204 messages[m.info.id] = merged

205 changed = true

206 }

207 }

208 }

209 if (evictOldMessages()) changed = true

210 return changed

211 }

212

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

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

215 return when (event) {

216 is MessageUpdated -> {

217 val info = event.properties.info

218 if (info.sessionID != sessionId) {

219 false

220 } else {

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

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

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

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

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

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

227 } else {

228 info.id

229 }

230 val existing = messages[key]

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

232 if (existing == null) evictOldMessages()

233 true

234 }

235 }

236

237 is MessagePartUpdated -> {

238 val part = event.properties.part

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

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

241 if (messageId == null || (partSession != null && partSession != sessionId)) {

242 false

243 } else {

244 upsertPart(messageId, part)

245 }

246 }

247

248 is MessagePartRemoved -> handlePartRemoved(sessionId, event)

249

250 is MessageRemoved -> handleMessageRemoved(sessionId, event)

251

252 else -> false

253 }

254 }

255

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

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

258 eventSession == null || eventSession == sessionId

259

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

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

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

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

264 return removePart(messageId, partId)

265 }

266

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

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

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

270 return messages.remove(id) != null

271 }

272

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

274 val current = messages[messageId]

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

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

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

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

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

280 } else {

281 parts + part

282 }

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

284 ?: MessageWithParts(

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

286 parts = newParts,

287 )

288 if (current == null) evictOldMessages()

289 return true

290 }

291

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

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

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

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

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

297 return true

298 }

299

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

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

302 private fun evictOldMessages(): Boolean {

303 var evicted = false

304 while (messages.size > maxMessages) {

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

306 messages.remove(oldestKey)

307 evicted = true

308 }

309 return evicted

310 }

311}