fix(chat): preserve usage baseline across local aborts
This commit is contained in:
1 parent
d16a2ea35c
commit
22b948f671
4 files changed
+80
-13
No files matched your search
@@ -1676,17 +1676,6 @@ class ChatViewModel(
|
|||||||
userMessage: ChatMessage? = null,
|
userMessage: ChatMessage? = null,
|
||||||
) {
|
) {
|
||||||
val dispatchGeneration = sessionGeneration
|
val dispatchGeneration = sessionGeneration
|
||||||
// Capture before attachment work or prompt submission can advance the
|
|
||||||
// cumulative counters. The reducer preserves this baseline across
|
|
||||||
// every assistant segment in a tool loop.
|
|
||||||
if (!_streamingState.value.turnUsageBaselineCaptured) {
|
|
||||||
_streamingState.update {
|
|
||||||
it.copy(
|
|
||||||
turnUsageBaseline = _uiState.value.sessionUsage,
|
|
||||||
turnUsageBaselineCaptured = true,
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
AuthManager.setLastOpenedSessionId(storageSessionId)
|
AuthManager.setLastOpenedSessionId(storageSessionId)
|
||||||
val msgToPersist =
|
val msgToPersist =
|
||||||
userMessage ?: ChatMessage(
|
userMessage ?: ChatMessage(
|
||||||
@@ -1826,6 +1815,7 @@ class ChatViewModel(
|
|||||||
if (dispatchGeneration == sessionGeneration) {
|
if (dispatchGeneration == sessionGeneration) {
|
||||||
ActiveSessionHolder.set(agentSessionId, storageSessionId)
|
ActiveSessionHolder.set(agentSessionId, storageSessionId)
|
||||||
}
|
}
|
||||||
|
captureTurnUsageBaselineIfNeeded()
|
||||||
if (wasStreaming && attachments.isEmpty()) {
|
if (wasStreaming && attachments.isEmpty()) {
|
||||||
wsClient.sendRedirect(
|
wsClient.sendRedirect(
|
||||||
agentSessionId,
|
agentSessionId,
|
||||||
@@ -1859,6 +1849,20 @@ class ChatViewModel(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private fun captureTurnUsageBaselineIfNeeded() {
|
||||||
|
if (_streamingState.value.turnUsageBaselineCaptured) return
|
||||||
|
_streamingState.update {
|
||||||
|
if (it.turnUsageBaselineCaptured) {
|
||||||
|
it
|
||||||
|
} else {
|
||||||
|
it.copy(
|
||||||
|
turnUsageBaseline = _uiState.value.sessionUsage,
|
||||||
|
turnUsageBaselineCaptured = true,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/** Snapshot one URI to private cache while enforcing the outbound frame limit. */
|
/** Snapshot one URI to private cache while enforcing the outbound frame limit. */
|
||||||
private suspend fun prepareAttachment(attachment: Attachment): PrepareAttachmentResult {
|
private suspend fun prepareAttachment(attachment: Attachment): PrepareAttachmentResult {
|
||||||
var encodedFile: File? = null
|
var encodedFile: File? = null
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ private fun List<ChatMessage>.upsertById(message: ChatMessage): List<ChatMessage
|
|||||||
|
|
||||||
private fun Long.toIntOrNullSafely(): Int? = takeIf { it in Int.MIN_VALUE.toLong()..Int.MAX_VALUE.toLong() }?.toInt()
|
private fun Long.toIntOrNullSafely(): Int? = takeIf { it in Int.MIN_VALUE.toLong()..Int.MAX_VALUE.toLong() }?.toInt()
|
||||||
|
|
||||||
|
private fun Double?.validTpsOrNull(): Double? = takeIf { it != null && it.isFinite() && it > 0.0 }
|
||||||
|
|
||||||
internal fun applyUsageSnapshot(
|
internal fun applyUsageSnapshot(
|
||||||
state: ChatUiState,
|
state: ChatUiState,
|
||||||
snapshot: com.m57.hermescontrol.data.model.UsageSnapshotResponse,
|
snapshot: com.m57.hermescontrol.data.model.UsageSnapshotResponse,
|
||||||
@@ -432,7 +434,7 @@ object ChatWsEventReducer {
|
|||||||
streamingState = StreamingState(),
|
streamingState = StreamingState(),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
val tps = (turnUsage?.avgTps ?: usageState.latestTps)?.takeIf { it.isFinite() && it > 0.0 }
|
val tps = turnUsage?.avgTps.validTpsOrNull() ?: usageState.latestTps.validTpsOrNull()
|
||||||
val tokenCount =
|
val tokenCount =
|
||||||
if (turnUsage?.outputTokens != null) {
|
if (turnUsage?.outputTokens != null) {
|
||||||
turnUsage.outputTokens.toIntOrNullSafely()
|
turnUsage.outputTokens.toIntOrNullSafely()
|
||||||
|
|||||||
@@ -1883,7 +1883,7 @@ class ChatViewModelTest {
|
|||||||
@Test
|
@Test
|
||||||
fun sendMessage_oversizedDuringSnapshotLeavesNoPersistedGhostMessage() =
|
fun sendMessage_oversizedDuringSnapshotLeavesNoPersistedGhostMessage() =
|
||||||
runTest {
|
runTest {
|
||||||
val (viewModel, _) = createViewModelWithSession()
|
val (viewModel, sessionId) = createViewModelWithSession()
|
||||||
val persistedBeforeSend = fakeRepo.dao.count()
|
val persistedBeforeSend = fakeRepo.dao.count()
|
||||||
val uriString = "content://unknown/growing-file"
|
val uriString = "content://unknown/growing-file"
|
||||||
val mockUri = mockk<Uri>()
|
val mockUri = mockk<Uri>()
|
||||||
@@ -1920,6 +1920,7 @@ class ChatViewModelTest {
|
|||||||
viewModel.sendMessage("Inspect changing file")
|
viewModel.sendMessage("Inspect changing file")
|
||||||
advanceUntilIdle()
|
advanceUntilIdle()
|
||||||
|
|
||||||
|
assertFalse(viewModel.streamingState.value.turnUsageBaselineCaptured)
|
||||||
assertEquals(persistedBeforeSend, fakeRepo.dao.count())
|
assertEquals(persistedBeforeSend, fakeRepo.dao.count())
|
||||||
assertFalse(
|
assertFalse(
|
||||||
fakeRepo.dao
|
fakeRepo.dao
|
||||||
@@ -1931,6 +1932,24 @@ class ChatViewModelTest {
|
|||||||
viewModel.uiState.value.messages
|
viewModel.uiState.value.messages
|
||||||
.any { it.content == "Inspect changing file" },
|
.any { it.content == "Inspect changing file" },
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mockEventsFlow.emit(
|
||||||
|
WsEvent.SessionUsage(
|
||||||
|
data = mapOf("usage" to mapOf("output" to 1300L)),
|
||||||
|
sessionId = sessionId,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
advanceUntilIdle()
|
||||||
|
viewModel.removeAttachment(0)
|
||||||
|
viewModel.sendMessage("Retry without attachment")
|
||||||
|
advanceUntilIdle()
|
||||||
|
|
||||||
|
assertEquals(
|
||||||
|
1300L,
|
||||||
|
viewModel.streamingState.value.turnUsageBaseline
|
||||||
|
?.outputTokens,
|
||||||
|
)
|
||||||
|
assertTrue(viewModel.streamingState.value.turnUsageBaselineCaptured)
|
||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
@@ -4191,6 +4210,7 @@ class ChatViewModelTest {
|
|||||||
advanceUntilIdle()
|
advanceUntilIdle()
|
||||||
|
|
||||||
viewModel.sendMessage("prompt")
|
viewModel.sendMessage("prompt")
|
||||||
|
advanceUntilIdle()
|
||||||
|
|
||||||
assertEquals(
|
assertEquals(
|
||||||
1000L,
|
1000L,
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import com.m57.hermescontrol.data.model.UsageSnapshotResponse
|
|||||||
import com.m57.hermescontrol.data.ws.WsEvent
|
import com.m57.hermescontrol.data.ws.WsEvent
|
||||||
import org.junit.Assert.assertEquals
|
import org.junit.Assert.assertEquals
|
||||||
import org.junit.Assert.assertFalse
|
import org.junit.Assert.assertFalse
|
||||||
|
import org.junit.Assert.assertNull
|
||||||
import org.junit.Assert.assertTrue
|
import org.junit.Assert.assertTrue
|
||||||
import org.junit.Test
|
import org.junit.Test
|
||||||
|
|
||||||
@@ -1185,6 +1186,46 @@ class ChatWsEventReducerTest {
|
|||||||
assertTrue(!result.streamingState.turnUsageBaselineCaptured)
|
assertTrue(!result.streamingState.turnUsageBaselineCaptured)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun testMessageComplete_invalidFinalTpsFallsBackToLatestValidTps() {
|
||||||
|
fun complete(
|
||||||
|
finalTps: Double?,
|
||||||
|
latestTps: Double?,
|
||||||
|
): Double? {
|
||||||
|
val state =
|
||||||
|
ChatUiState(
|
||||||
|
currentSessionId = "session-1",
|
||||||
|
latestTps = latestTps,
|
||||||
|
sessionUsage = UsageSnapshotResponse(outputTokens = 1000),
|
||||||
|
)
|
||||||
|
val start =
|
||||||
|
ChatWsEventReducer.reduce(
|
||||||
|
state,
|
||||||
|
StreamingState(),
|
||||||
|
WsEvent.MessageStart("session-1"),
|
||||||
|
"session-1",
|
||||||
|
)
|
||||||
|
val usage = mutableMapOf<String, Any?>("output" to 1250L)
|
||||||
|
if (finalTps != null) usage["avg_tps"] = finalTps
|
||||||
|
return ChatWsEventReducer
|
||||||
|
.reduce(
|
||||||
|
state,
|
||||||
|
start.streamingState,
|
||||||
|
WsEvent.MessageComplete("reply", "session-1", rawPayload = mapOf("usage" to usage)),
|
||||||
|
"session-1",
|
||||||
|
).state.messages
|
||||||
|
.single()
|
||||||
|
.tps
|
||||||
|
}
|
||||||
|
|
||||||
|
assertEquals(42.5, complete(42.5, 30.0) ?: 0.0, 0.001)
|
||||||
|
assertEquals(37.5, complete(0.0, 37.5) ?: 0.0, 0.001)
|
||||||
|
assertEquals(37.5, complete(-2.0, 37.5) ?: 0.0, 0.001)
|
||||||
|
assertEquals(37.5, complete(null, 37.5) ?: 0.0, 0.001)
|
||||||
|
assertNull(complete(-2.0, null))
|
||||||
|
assertNull(complete(0.0, -1.0))
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun testMessageComplete_withoutBaselineFallsBackToVisibleTextEstimate() {
|
fun testMessageComplete_withoutBaselineFallsBackToVisibleTextEstimate() {
|
||||||
val result =
|
val result =
|
||||||
|
|||||||
Reference in new issue
Block a user