fix(#99): gate tool and clarify events by session id to prevent leaks (#168)

This commit is contained in:
M57 authored and GitHub committed 2026-06-20 17:27:34 +03:00
1 parent 6cb4e17882
commit 74714fc1df
4 files changed
+56 -3

No files matched your search

@@ -71,12 +71,12 @@ object EventParser {
"tool.start" -> {
val name = payload?.get("name") as? String
WsEvent.ToolStart(name, payload)
WsEvent.ToolStart(name, payload, sessionId)
}
"tool.complete" -> {
val name = payload?.get("name") as? String
WsEvent.ToolComplete(name, payload)
WsEvent.ToolComplete(name, payload, sessionId)
}
"clarify.request" -> {
@@ -86,7 +86,7 @@ object EventParser {
@Suppress("UNCHECKED_CAST")
val options = (rawOptions as? List<*>)?.filterIsInstance<String>()
WsEvent.ClarifyRequest(text, options, clarifyId)
WsEvent.ClarifyRequest(text, options, clarifyId, sessionId)
}
"status.update" -> {
@@ -44,11 +44,13 @@ sealed class WsEvent {
data class ToolStart(
val name: String?,
val data: Map<String, Any?>?,
val sessionId: String? = null,
) : WsEvent()
data class ToolComplete(
val name: String?,
val data: Map<String, Any?>?,
val sessionId: String? = null,
) : WsEvent()
// ── Interactive ──────────────────────────────────────────────────────
@@ -57,6 +59,7 @@ sealed class WsEvent {
val text: String?,
val options: List<String>?,
val clarifyId: String? = null,
val sessionId: String? = null,
) : WsEvent()
// ── Status ───────────────────────────────────────────────────────────
@@ -160,6 +160,7 @@ class ChatViewModel(
is WsEvent.ToolComplete -> handleToolComplete(event)
is WsEvent.ClarifyRequest -> {
if (!isCurrentSession(event.sessionId)) return
_uiState.update {
it.copy(
clarifyRequest =
@@ -369,6 +370,7 @@ class ChatViewModel(
// ── Tool events ──────────────────────────────────────────────────────
private fun handleToolStart(event: WsEvent.ToolStart) {
if (!isCurrentSession(event.sessionId)) return
var orphanToPersist: ChatMessage? = null
val sessionId = _uiState.value.currentSessionId
@@ -409,6 +411,7 @@ class ChatViewModel(
}
private fun handleToolComplete(event: WsEvent.ToolComplete) {
if (!isCurrentSession(event.sessionId)) return
var completedTool: ChatMessage? = null
_uiState.update { state ->
@@ -610,4 +610,51 @@ class ChatViewModelTest {
assertEquals("{\"result\":\"4\",\"exit_code\":0}", state.messages[1].content)
assertEquals(ToolStatus.COMPLETED, state.messages[1].toolStatus)
}
@Test
fun testSessionMismatchEventsAreIgnored() =
runTest {
val viewModel = ChatViewModel(app, startCleanup = false)
advanceUntilIdle()
// Setup active session "session-active"
var createReqId = ""
every { HermesWsClient.send(WsMethods.SESSION_CREATE, any(), any()) } answers {
createReqId = "create-req-mismatch-test"
val onSent = arg<((String) -> Unit)?>(2)
onSent?.invoke(createReqId)
createReqId
}
mockEventsFlow.emit(WsEvent.GatewayReady(null))
advanceUntilIdle()
mockEventsFlow.emit(WsEvent.RpcResult(createReqId, mapOf("session_id" to "session-active")))
advanceUntilIdle()
// Emit events for another session "session-other" — should be ignored
mockEventsFlow.emit(
WsEvent.ToolStart(
name = "calculator",
data = mapOf("input" to "2+2"),
sessionId = "session-other",
),
)
mockEventsFlow.emit(
WsEvent.ClarifyRequest(
text = "Choose:",
options = listOf("Yes"),
clarifyId = "clarify-1",
sessionId = "session-other",
),
)
mockEventsFlow.emit(WsEvent.MessageStart("session-other"))
mockEventsFlow.emit(WsEvent.MessageToken("Hello", "session-other"))
advanceUntilIdle()
val state = viewModel.uiState.value
// Only the "Session created" message should be present (mismatch events ignored)
assertEquals(1, state.messages.size)
assertEquals("Session created", state.messages[0].content)
assertNull(state.streamingMessage)
assertNull(state.clarifyRequest)
}
}