feat(ws): adopt event replay via session.events.since on reconnect (#1016) (#1027)

Co-authored-by: M57 <hy4ri@users.noreply.github.com>
This commit is contained in:
M57andM57 authored and GitHub committed 2026-08-31 19:32:17 +03:00
1 parent 53c8eca714
commit 3fd9ac767f
5 files changed
+408

No files matched your search

@@ -31,6 +31,17 @@ object EventParser {
// ── Notification / event (no id, has method) ─────────────────────
@Suppress("UNCHECKED_CAST")
val params = response.params?.toAny() as? Map<String, Any?> ?: return WsEvent.Unknown(rawJson)
return parseParams(params, rawJson)
}
/**
* Parses an event object (bare params map containing `type`, `session_id`, `seq`, `payload`).
* Used both for live notifications and for event lists returned by `session.events.since`.
*/
fun parseParams(
params: Map<String, Any?>,
rawJson: String = "",
): WsEvent {
val eventType = params["type"] as? String ?: return WsEvent.Unknown(rawJson)
@Suppress("UNCHECKED_CAST")
@@ -27,12 +27,14 @@ import kotlinx.coroutines.flow.filterIsInstance
import kotlinx.coroutines.launch
import kotlinx.serialization.decodeFromString
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.JsonElement
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okio.utf8Size
import java.util.Collections
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.atomic.AtomicBoolean
@@ -128,6 +130,19 @@ object HermesWsClient {
@Volatile
private var outboundDrainJob: Job? = null
// ── Sequence tracking & replay on reconnect (issue #1016) ─────────────
private val lastSeenSeq = ConcurrentHashMap<String, Int>()
@Volatile
private var replayEpoch: String? = null
@Volatile
private var replayInFlight: Boolean = false
private val replayHold = ConcurrentHashMap<String, MutableList<Pair<Int?, WsEvent>>>()
@Volatile
private var replayJob: Job? = null
// ── Health and Ping/Pong tracking ────────────────────────────────────
@Volatile
@@ -277,6 +292,20 @@ object HermesWsClient {
.filterIsInstance<WsEvent.ChangeEvent>()
.collect { ChangeEventHub.emit(it) }
}
// Track gateway replay_epoch across gateway.ready events (issue #1016)
wsScope.launch {
events.collect { event ->
if (event is WsEvent.GatewayReady) {
val epoch = event.data?.get("replay_epoch") as? String
if (!epoch.isNullOrEmpty()) {
if (replayEpoch != null && replayEpoch != epoch) {
lastSeenSeq.clear()
}
replayEpoch = epoch
}
}
}
}
}
// ── Connection helpers ────────────────────────────────────────────────
@@ -576,6 +605,12 @@ object HermesWsClient {
queuedMessagesById.clear()
pendingPromptSubmits.clear()
pendingReply = false
lastSeenSeq.clear()
replayHold.clear()
replayEpoch = null
replayJob?.cancel()
replayJob = null
replayInFlight = false
}
connected.set(false)
_connectionStatus.value = ConnectionStatus.DISCONNECTED
@@ -819,6 +854,135 @@ object HermesWsClient {
onSent = onSent,
)
// ── Sequence tracking & replay on reconnect (issue #1016) ─────────────
private fun triggerReplay() {
if (lastSeenSeq.isEmpty() || replayInFlight) return
replayJob?.cancel()
replayJob =
wsScope.launch {
fetchReplay()
}
}
private suspend fun fetchReplay() {
if (replayInFlight || lastSeenSeq.isEmpty()) return
replayInFlight = true
val sessionsToReplay = lastSeenSeq.keys().toList()
for (sid in sessionsToReplay) {
replayHold[sid] = Collections.synchronizedList(mutableListOf())
}
try {
for (sid in sessionsToReplay) {
val lastSeen = lastSeenSeq[sid] ?: continue
try {
val deferred =
request(
method = WsMethods.SESSION_EVENTS_SINCE,
params = mapOf("session_id" to sid, "last_seen" to lastSeen, "since_seq" to lastSeen),
timeoutMs = 10_000L,
)
val result = deferred.await()
@Suppress("UNCHECKED_CAST")
val resultMap =
when (result) {
is JsonElement -> result.toAny() as? Map<String, Any?>
is Map<*, *> -> result as? Map<String, Any?>
else -> null
}
if (resultMap != null) {
val epoch = resultMap["epoch"] as? String
if (!epoch.isNullOrEmpty()) {
if (replayEpoch != null && replayEpoch != epoch) {
replayEpoch = epoch
lastSeenSeq.clear()
continue
}
replayEpoch = epoch
}
val eventsList = (resultMap["events"] as? List<*>)?.filterIsInstance<Map<String, Any?>>()
if (eventsList != null) {
for (eventMap in eventsList) {
val eventSeq = (eventMap["seq"] as? Number)?.toInt()
val eventSid = (eventMap["session_id"] as? String) ?: sid
if (eventSeq != null && eventSid.isNotEmpty()) {
val prev = lastSeenSeq[eventSid] ?: 0
if (eventSeq <= prev) continue
lastSeenSeq[eventSid] = eventSeq
}
val parsedEvent = EventParser.parseParams(eventMap)
val finalEvent =
if (parsedEvent is WsEvent.MessageComplete) {
parsedEvent.copy(
storedSessionId =
ActiveSessionHolder.resolveStoredSessionId(
parsedEvent.sessionId,
),
)
} else {
parsedEvent
}
parsedEvents.tryEmit(finalEvent)
}
}
}
} catch (e: Exception) {
Log.w(TAG, "Replay failed for session $sid: ${e.message}")
}
}
} finally {
for (sid in sessionsToReplay) {
val held = replayHold.remove(sid)
if (held != null) {
synchronized(held) {
for ((heldSeq, heldEvent) in held) {
if (heldSeq != null) {
val prev = lastSeenSeq[sid] ?: 0
if (heldSeq <= prev) continue
lastSeenSeq[sid] = heldSeq
}
parsedEvents.tryEmit(heldEvent)
}
}
}
}
replayInFlight = false
}
}
@VisibleForTesting
internal fun getSeqWatermarks(): Map<String, Int> = HashMap(lastSeenSeq)
@VisibleForTesting
internal fun setSeqWatermark(
sessionId: String,
seq: Int,
) {
lastSeenSeq[sessionId] = seq
}
@VisibleForTesting
internal fun clearSeqWatermarks() {
lastSeenSeq.clear()
replayHold.clear()
replayEpoch = null
replayInFlight = false
}
@VisibleForTesting
internal fun isReplayInFlight(): Boolean = replayInFlight
@VisibleForTesting
internal fun setReplayEpochForTest(epoch: String?) {
replayEpoch = epoch
}
@VisibleForTesting
internal suspend fun fetchReplayForTest() {
fetchReplay()
}
// ── Internal ─────────────────────────────────────────────────────────
private fun openSocket() {
@@ -937,6 +1101,7 @@ object HermesWsClient {
messageQueue.poll()
markQueuedMessageSent(msg)
}
triggerReplay()
}
disconnectIfIdleInBackground()
}
@@ -963,6 +1128,39 @@ object HermesWsClient {
removeQueuedMessage(rpcId)
resolvePending(rpcId, rpc.result, null)
}
if (rpc.id == null) {
@Suppress("UNCHECKED_CAST")
val params = rpc.params?.toAny() as? Map<String, Any?>
val sid =
(params?.get("session_id") as? String)
?: ((params?.get("payload") as? Map<*, *>)?.get("session_id") as? String)
val seq = (params?.get("seq") as? Number)?.toInt()
if (!sid.isNullOrBlank() && seq != null) {
val prev = lastSeenSeq[sid] ?: 0
if (seq <= prev) {
return
}
val parsed = EventParser.parse(rpc, text)
val finalEvent =
if (parsed is WsEvent.MessageComplete) {
parsed.copy(
storedSessionId = ActiveSessionHolder.resolveStoredSessionId(parsed.sessionId),
)
} else {
parsed
}
if (replayInFlight && replayHold.containsKey(sid)) {
replayHold[sid]?.add(seq to finalEvent)
return
}
lastSeenSeq[sid] = seq
parsedEvents.tryEmit(finalEvent)
return
}
}
val parsed = EventParser.parse(rpc, text)
if (parsed is WsEvent.MessageComplete) {
parsed.copy(
@@ -15,6 +15,12 @@ object WsMethods {
const val SESSION_TITLE = "session.title"
const val SESSION_BRANCH = "session.branch"
/** Replay recorded events newer than client's last-seen seq for a session. */
const val SESSION_EVENTS_SINCE = "session.events.since"
/** Replay buffer telemetry for ops/debug. */
const val SESSION_EVENTS_STATS = "session.events.stats"
/** Live context-window occupancy for the chat meter (desktop-mirror). */
const val SESSION_CONTEXT_BREAKDOWN = "session.context_breakdown"
@@ -668,4 +668,36 @@ class EventParserTest {
assertEquals(3, (usage?.get("compressions") as? Number)?.toInt())
assertEquals(12500, (usage?.get("context_used") as? Number)?.toInt())
}
@Test
fun testParseParams_parsesBareEventMap() {
val params =
mapOf(
"type" to "message.token",
"session_id" to "sess-replay-1",
"seq" to 42,
"payload" to mapOf("text" to "hello replay"),
)
val event = EventParser.parseParams(params)
assertTrue(event is WsEvent.MessageToken)
val tokenEvent = event as WsEvent.MessageToken
assertEquals("hello replay", tokenEvent.token)
assertEquals("sess-replay-1", tokenEvent.sessionId)
}
@Test
fun testParseParams_parsesToolEvent() {
val params =
mapOf(
"type" to "tool.start",
"session_id" to "sess-replay-2",
"seq" to 43,
"payload" to mapOf("name" to "terminal", "tool_id" to "t1"),
)
val event = EventParser.parseParams(params)
assertTrue(event is WsEvent.ToolStart)
val toolEvent = event as WsEvent.ToolStart
assertEquals("terminal", toolEvent.name)
assertEquals("sess-replay-2", toolEvent.sessionId)
}
}
@@ -1527,4 +1527,165 @@ class HermesWsClientTest {
HermesWsClient.connectionStatus.value != ConnectionStatus.CONNECTED,
)
}
@Test
fun testSequenceNumberDeduplication() {
var serverWebSocket: WebSocket? = null
val serverLatch = CountDownLatch(1)
mockWebServer.enqueue(
MockResponse().withWebSocketUpgrade(
object : WebSocketListener() {
override fun onOpen(
webSocket: WebSocket,
response: okhttp3.Response,
) {
serverWebSocket = webSocket
serverLatch.countDown()
}
},
),
)
HermesWsClient.connect()
runBlocking {
withTimeout(5000) {
HermesWsClient.connectionStatus.first { it == ConnectionStatus.CONNECTED }
}
}
assertTrue(serverLatch.await(5, TimeUnit.SECONDS))
val ws = serverWebSocket
assertNotNull(ws)
val receivedTokens = mutableListOf<String>()
val collectorJob =
kotlinx.coroutines.CoroutineScope(kotlinx.coroutines.Dispatchers.IO).launch {
HermesWsClient.events.collect { event ->
if (event is WsEvent.MessageToken) {
receivedTokens.add(event.token)
}
}
}
// Send seq 1
ws!!.send(
"""{"method":"event","params":{"type":"message.token","session_id":"s1","seq":1,"payload":{"text":"A"}}}""",
)
Thread.sleep(100)
assertEquals(1, HermesWsClient.getSeqWatermarks()["s1"])
assertEquals(listOf("A"), receivedTokens)
// Send duplicate seq 1 -> should be dropped
ws.send(
"""{"method":"event","params":{"type":"message.token","session_id":"s1","seq":1,"payload":{"text":"A-dup"}}}""",
)
Thread.sleep(100)
assertEquals(1, HermesWsClient.getSeqWatermarks()["s1"])
assertEquals(listOf("A"), receivedTokens)
// Send seq 2 -> should be accepted
ws.send(
"""{"method":"event","params":{"type":"message.token","session_id":"s1","seq":2,"payload":{"text":"B"}}}""",
)
Thread.sleep(100)
assertEquals(2, HermesWsClient.getSeqWatermarks()["s1"])
assertEquals(listOf("A", "B"), receivedTokens)
collectorJob.cancel()
}
@Test
fun testEpochChangeClearsWatermarks() {
var serverWebSocket: WebSocket? = null
val serverLatch = CountDownLatch(1)
mockWebServer.enqueue(
MockResponse().withWebSocketUpgrade(
object : WebSocketListener() {
override fun onOpen(
webSocket: WebSocket,
response: okhttp3.Response,
) {
serverWebSocket = webSocket
serverLatch.countDown()
}
},
),
)
HermesWsClient.connect()
runBlocking {
withTimeout(5000) {
HermesWsClient.connectionStatus.first { it == ConnectionStatus.CONNECTED }
}
}
assertTrue(serverLatch.await(5, TimeUnit.SECONDS))
val ws = serverWebSocket
assertNotNull(ws)
// Initial gateway.ready with epoch-1
ws!!.send(
"""{"method":"event","params":{"type":"gateway.ready","payload":{"replay_epoch":"epoch-1"}}}""",
)
Thread.sleep(100)
HermesWsClient.setSeqWatermark("s1", 10)
assertEquals(10, HermesWsClient.getSeqWatermarks()["s1"])
// Second gateway.ready with new epoch-2 -> backend restarted, watermarks cleared
ws.send(
"""{"method":"event","params":{"type":"gateway.ready","payload":{"replay_epoch":"epoch-2"}}}""",
)
Thread.sleep(100)
assertTrue(HermesWsClient.getSeqWatermarks().isEmpty())
}
@Test
fun testReplayOnReconnectFlow() {
var serverWebSocket: WebSocket? = null
val serverLatch = CountDownLatch(1)
val requestLatch = CountDownLatch(1)
var receivedMethod: String? = null
mockWebServer.enqueue(
MockResponse().withWebSocketUpgrade(
object : WebSocketListener() {
override fun onOpen(
webSocket: WebSocket,
response: okhttp3.Response,
) {
serverWebSocket = webSocket
serverLatch.countDown()
}
override fun onMessage(
webSocket: WebSocket,
text: String,
) {
if (text.contains(WsMethods.SESSION_EVENTS_SINCE)) {
receivedMethod = WsMethods.SESSION_EVENTS_SINCE
// Extract ID from JSON-RPC request
val id = Regex(""""id":"([^"]+)"""").find(text)?.groupValues?.get(1) ?: "1"
webSocket.send(
"""{"jsonrpc":"2.0","id":"$id","result":{"epoch":"ep1","events":[{"type":"message.token","session_id":"s1","seq":6,"payload":{"text":"replayed"}}]}}""",
)
requestLatch.countDown()
}
}
},
),
)
HermesWsClient.setSeqWatermark("s1", 5)
HermesWsClient.connect()
runBlocking {
withTimeout(5000) {
HermesWsClient.connectionStatus.first { it == ConnectionStatus.CONNECTED }
}
}
assertTrue(serverLatch.await(5, TimeUnit.SECONDS))
assertTrue(requestLatch.await(5, TimeUnit.SECONDS))
assertEquals(WsMethods.SESSION_EVENTS_SINCE, receivedMethod)
// Wait for replay processing
Thread.sleep(200)
assertEquals(6, HermesWsClient.getSeqWatermarks()["s1"])
}
}