feat(#919): adopt session.usage WS push event for real-time context & compression updates (#927)

Co-authored-by: M57 <hy4ri@users.noreply.github.com>
This commit is contained in:
M57andM57 authored and GitHub committed 2026-08-17 23:31:07 +03:00
1 parent d6892cde7f
commit acd82c0085
6 files changed
+141 -9

No files matched your search

@@ -634,4 +634,38 @@ class EventParserTest {
assertEquals("sub-456", subagentEvent.payload?.get("subagent_id"))
assertEquals("analyzing docs...", subagentEvent.payload?.get("text"))
}
@Test
fun testParseSessionUsage_returnsSessionUsageEvent() {
val response =
createJsonRpcResponse(
jsonrpc = "2.0",
id = null,
result = null,
error = null,
method = "event",
params =
mapOf(
"type" to "session.usage",
"session_id" to "sess-123",
"payload" to
mapOf(
"usage" to
mapOf(
"compressions" to 3,
"context_used" to 12500,
"context_max" to 128000,
"total" to 45000,
),
),
),
)
val event = EventParser.parse(response)
assertTrue(event is WsEvent.SessionUsage)
val sessionUsage = event as WsEvent.SessionUsage
assertEquals("sess-123", sessionUsage.sessionId)
val usage = sessionUsage.data?.get("usage") as? Map<*, *>
assertEquals(3, (usage?.get("compressions") as? Number)?.toInt())
assertEquals(12500, (usage?.get("context_used") as? Number)?.toInt())
}
}
@@ -869,4 +869,51 @@ class ChatWsEventReducerTest {
assertEquals(listOf(line1, line2, summary), allAssistant.map { it.content })
assertEquals(null, result.streamingState.streamingMessage)
}
@Test
fun testSessionUsage_updatesCompressionCountAndContextTokens() {
val state = ChatUiState(currentSessionId = "session-1")
val usagePayload =
mapOf(
"usage" to
mapOf(
"compressions" to 2,
"context_used" to 8500L,
"context_max" to 128000L,
"total" to 25000L,
),
)
val result =
ChatWsEventReducer.reduce(
state = state,
streamingState = StreamingState(),
event = WsEvent.SessionUsage(data = usagePayload, sessionId = "session-1"),
currentSessionId = "session-1",
)
assertEquals(2, result.state.compressionCount)
assertEquals(8500L, result.state.usedContextTokens)
assertEquals(128000L, result.state.fullContextTokens)
}
@Test
fun testSessionUsage_differentSession_isIgnored() {
val state = ChatUiState(currentSessionId = "session-1", compressionCount = 1)
val usagePayload =
mapOf(
"usage" to
mapOf(
"compressions" to 5,
"context_used" to 99999L,
),
)
val result =
ChatWsEventReducer.reduce(
state = state,
streamingState = StreamingState(),
event = WsEvent.SessionUsage(data = usagePayload, sessionId = "session-other"),
currentSessionId = "session-1",
)
assertEquals(1, result.state.compressionCount)
assertEquals(null, result.state.usedContextTokens)
}
}