feat(chat): display token usage and tps in message footers (#1094)

This commit is contained in:
M57 authored and GitHub committed 2026-09-12 16:04:13 +03:00
1 parent 0bb087c41e
commit 065c7fd75f
24 files changed
+448 -3

No files matched your search

@@ -42,4 +42,8 @@ data class ChatMessageEntity(
val isStreaming: Boolean = false,
@ColumnInfo(name = "display_kind")
val displayKind: String? = null,
@ColumnInfo(name = "token_count")
val tokenCount: Int? = null,
@ColumnInfo(name = "tps")
val tps: Double? = null,
)
@@ -35,6 +35,8 @@ fun ChatMessageEntity.toUiModel(): ChatMessage =
else -> null
},
displayKind = displayKind,
tokenCount = tokenCount,
tps = tps,
)
fun ChatMessage.toEntity(sessionId: String): ChatMessageEntity =
@@ -50,4 +52,6 @@ fun ChatMessage.toEntity(sessionId: String): ChatMessageEntity =
toolStatus = toolStatus?.name,
isStreaming = isStreaming,
displayKind = displayKind,
tokenCount = tokenCount,
tps = tps,
)
@@ -12,7 +12,7 @@ import java.io.File
@Database(
entities = [ChatMessageEntity::class],
version = 6,
version = 7,
exportSchema = true,
)
abstract class HermesDatabase : RoomDatabase() {
@@ -67,6 +67,18 @@ abstract class HermesDatabase : RoomDatabase() {
}
}
val MIGRATION_6_7: Migration =
object : Migration(6, 7) {
override suspend fun migrate(connection: SQLiteConnection) {
connection.execSQL(
"ALTER TABLE `chat_messages` ADD COLUMN `token_count` INTEGER",
)
connection.execSQL(
"ALTER TABLE `chat_messages` ADD COLUMN `tps` REAL",
)
}
}
fun get(context: Context): HermesDatabase =
instance ?: synchronized(this) {
// SQLCipher can't open plaintext SQLite databases — if an old
@@ -97,6 +109,7 @@ abstract class HermesDatabase : RoomDatabase() {
MIGRATION_3_4,
MIGRATION_4_5,
MIGRATION_5_6,
MIGRATION_6_7,
).fallbackToDestructiveMigration(false)
.build()
.also { instance = it }
@@ -4,6 +4,7 @@ import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.doubleOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.longOrNull
@Serializable
@@ -50,10 +51,18 @@ data class SessionMessage(
val display_kind: String? = null,
/** Display-only metadata for the marker (e.g. delegation result counts). */
val display_metadata: JsonElement? = null,
/** Token count recorded by the backend. */
val token_count: JsonElement? = null,
) {
val timestampText: String?
get() = (timestamp as? JsonPrimitive)?.content
val tokenCount: Int?
get() {
val prim = token_count as? JsonPrimitive ?: return null
return prim.intOrNull ?: prim.content.toIntOrNull()
}
val timestampEpochMs: Long?
get() {
val prim = timestamp as? JsonPrimitive ?: return null
@@ -20,6 +20,8 @@ data class UsageSnapshotResponse(
val contextMax: Long? = null,
/** Cumulative session total tokens. */
val totalTokens: Long? = null,
/** Rolling tokens-per-second generation speed. */
val avgTps: Double? = null,
)
/**
@@ -38,5 +40,6 @@ fun parseUsageSnapshot(result: Any?): UsageSnapshotResponse? {
contextUsed = (usageMap["context_used"] as? Number)?.toLong(),
contextMax = (usageMap["context_max"] as? Number)?.toLong(),
totalTokens = (usageMap["total"] as? Number)?.toLong(),
avgTps = (usageMap["avg_tps"] as? Number)?.toDouble(),
)
}
@@ -87,7 +87,7 @@ object EventParser {
"message.complete" -> {
val text = payload?.get("text") as? String ?: ""
val reasoning = payload?.get("reasoning") as? String
WsEvent.MessageComplete(text, sessionId, reasoning)
WsEvent.MessageComplete(text, sessionId, reasoning, rawPayload = payload)
}
"message.done" -> {
@@ -61,6 +61,8 @@ sealed class WsEvent {
val reasoning: String? = null,
/** Stored session id captured before background disconnect clears the active mapping. */
val storedSessionId: String? = null,
/** Full raw payload map (including usage/avg_tps) emitted with message.complete. */
val rawPayload: Map<String, Any?>? = null,
) : WsEvent()
data class MessageDone(
@@ -34,6 +34,8 @@ data class GroupChatMessage(
val isSystem: Boolean = false,
val isPass: Boolean = false,
val toolCalls: List<GroupChatToolCall> = emptyList(),
val tokenCount: Int? = null,
val tps: Double? = null,
)
const val STOPPED_SYSTEM_TEXT = "Discussion stopped"
@@ -49,6 +49,7 @@ import com.m57.hermescontrol.ui.bots.group.CAPPED_SYSTEM_TEXT
import com.m57.hermescontrol.ui.bots.group.GroupChatMessage
import com.m57.hermescontrol.ui.bots.group.STOPPED_SYSTEM_TEXT
import com.m57.hermescontrol.ui.chat.MarkdownText
import com.m57.hermescontrol.ui.chat.TokenEstimator
import com.m57.hermescontrol.ui.chat.formatTimestamp
import com.m57.hermescontrol.ui.common.BotAvatar
import com.m57.hermescontrol.util.BidiUtils
@@ -195,6 +196,26 @@ fun GroupMessageCard(
color = MaterialTheme.colorScheme.onPrimary.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
)
val userTokens =
message.tokenCount ?: TokenEstimator.estimate(message.text).takeIf { it > 0 }
if (userTokens != null && userTokens > 0) {
Spacer(modifier = Modifier.width(4.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onPrimary.copy(alpha = 0.5f),
)
Spacer(modifier = Modifier.width(4.dp))
Text(
text =
stringResource(
R.string.chat_msg_tokens,
TokenEstimator.formatTokenCount(userTokens),
),
color = MaterialTheme.colorScheme.onPrimary.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
)
}
}
}
}
@@ -304,6 +325,44 @@ fun GroupMessageCard(
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
)
val botTokens =
message.tokenCount ?: TokenEstimator.estimate(message.text).takeIf { it > 0 }
if (botTokens != null && botTokens > 0) {
Spacer(modifier = Modifier.width(4.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.4f),
)
Spacer(modifier = Modifier.width(4.dp))
Text(
text =
stringResource(
R.string.chat_msg_tokens,
TokenEstimator.formatTokenCount(botTokens),
),
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
)
}
if (message.tps != null && message.tps > 0.0) {
Spacer(modifier = Modifier.width(4.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.4f),
)
Spacer(modifier = Modifier.width(4.dp))
Text(
text =
stringResource(
R.string.chat_msg_tps,
TokenEstimator.formatTps(message.tps),
),
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
)
}
}
}
}
@@ -257,6 +257,44 @@ fun ChatBubble(
color = userBubbleTextColor.copy(alpha = 0.6f),
style = MaterialTheme.typography.labelSmall,
)
if (message.tokenCount != null && message.tokenCount > 0) {
Spacer(modifier = Modifier.width(4.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = userBubbleTextColor.copy(alpha = 0.4f),
)
Spacer(modifier = Modifier.width(4.dp))
Text(
text =
stringResource(
R.string.chat_msg_tokens,
TokenEstimator.formatTokenCount(message.tokenCount),
),
color = userBubbleTextColor.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
modifier = Modifier.testTag("bubble_token_count"),
)
}
if (message.tps != null && message.tps > 0.0) {
Spacer(modifier = Modifier.width(4.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = userBubbleTextColor.copy(alpha = 0.4f),
)
Spacer(modifier = Modifier.width(4.dp))
Text(
text =
stringResource(
R.string.chat_msg_tps,
TokenEstimator.formatTps(message.tps),
),
color = userBubbleTextColor.copy(alpha = 0.7f),
style = MaterialTheme.typography.labelSmall,
modifier = Modifier.testTag("bubble_tps"),
)
}
}
}
}
@@ -79,6 +79,10 @@ data class ChatMessage(
* When present, rendered at the bottom beside the copy button.
*/
val finishTimestamp: Long? = null,
/** Token count used by this message (persisted to SQLite). */
val tokenCount: Int? = null,
/** Tokens per second generation speed (persisted to SQLite). */
val tps: Double? = null,
)
/**
@@ -171,6 +171,7 @@ internal fun mapServerMessages(
}
}
val tokenCount = msg.tokenCount ?: TokenEstimator.estimate(finalContent).takeIf { it > 0 }
mapped.add(
ChatMessage(
id = restId,
@@ -182,6 +183,7 @@ internal fun mapServerMessages(
timestamp = timestamp,
isStreaming = false,
displayKind = msg.display_kind,
tokenCount = tokenCount,
),
)
}
@@ -146,6 +146,8 @@ data class ChatUiState(
// until the first successful session.usage fetch) — drives the
// "compressed ×N" badge on the context chip.
val compressionCount: Int? = null,
/** Rolling output tokens/sec over the last ~10 calls. */
val latestTps: Double? = null,
// Attachment state
val pendingAttachments: List<Attachment> = emptyList(),
/** One-shot composer recovery after an attachment is rejected before send. */
@@ -1298,6 +1300,7 @@ class ChatViewModel(
role = MessageRole.USER,
content = text,
attachments = if (attachments.isNotEmpty()) attachments else null,
tokenCount = TokenEstimator.estimate(text).takeIf { it > 0 },
)
// Update UI immediately
@@ -1342,6 +1345,7 @@ class ChatViewModel(
role = MessageRole.USER,
content = text,
attachments = if (attachments.isNotEmpty()) attachments else null,
tokenCount = TokenEstimator.estimate(text).takeIf { it > 0 },
)
// Upload attachments then submit prompt
@@ -1624,7 +1628,12 @@ class ChatViewModel(
val displayContent =
if (result is SlashResult.QueuePrompt) result.displayContent else command
val userMsg = ChatMessage(role = MessageRole.USER, content = displayContent)
val userMsg =
ChatMessage(
role = MessageRole.USER,
content = displayContent,
tokenCount = TokenEstimator.estimate(displayContent).takeIf { it > 0 },
)
val sessionId = _uiState.value.currentSessionId
_uiState.update { it.copy(messages = it.messages + userMsg) }
@@ -221,6 +221,8 @@ object ChatWsEventReducer {
streamingState.streamingMessage.copy(
isStreaming = false,
finishTimestamp = System.currentTimeMillis(),
tokenCount = TokenEstimator.estimate(streamingState.streamingMessage.content).takeIf { it > 0 },
tps = state.latestTps?.takeIf { it > 0.0 },
)
orphan = finalized
state.copy(
@@ -379,17 +381,25 @@ object ChatWsEventReducer {
streamingState = streamingState.copy(streamingMessage = null),
)
}
val rawUsage = event.rawPayload?.get("usage") as? Map<*, *>
val tpsFromPayload = (rawUsage?.get("avg_tps") as? Number)?.toDouble()
val tps = (tpsFromPayload ?: state.latestTps)?.takeIf { it > 0.0 }
val tokenCount = TokenEstimator.estimate(text).takeIf { it > 0 }
val msg =
streaming?.copy(
content = text,
isStreaming = false,
reasoningText = reasoning,
finishTimestamp = System.currentTimeMillis(),
tokenCount = tokenCount,
tps = tps,
) ?: ChatMessage(
role = MessageRole.ASSISTANT,
content = text,
reasoningText = reasoning,
finishTimestamp = System.currentTimeMillis(),
tokenCount = tokenCount,
tps = tps,
)
val effects = mutableListOf<ReducerEffect>()
val sid = state.currentSessionId
@@ -435,11 +445,14 @@ object ChatWsEventReducer {
} else {
streaming.reasoningText
}
val tokenCount = TokenEstimator.estimate(streaming.content).takeIf { it > 0 }
val msg =
streaming.copy(
isStreaming = false,
reasoningText = reasoning,
finishTimestamp = System.currentTimeMillis(),
tokenCount = streaming.tokenCount ?: tokenCount,
tps = streaming.tps ?: state.latestTps?.takeIf { it > 0.0 },
)
val effects = mutableListOf<ReducerEffect>()
val sid = state.currentSessionId
@@ -504,6 +517,8 @@ object ChatWsEventReducer {
isStreaming = false,
reasoningText = reasoning,
finishTimestamp = System.currentTimeMillis(),
tokenCount = TokenEstimator.estimate(streamingState.streamingMessage.content).takeIf { it > 0 },
tps = state.latestTps?.takeIf { it > 0.0 },
)
orphanToPersist = finalized
state.copy(
@@ -962,6 +977,9 @@ object ChatWsEventReducer {
if (snapshot.contextMax != null && snapshot.contextMax > 0L) {
updatedState = updatedState.copy(fullContextTokens = snapshot.contextMax)
}
if (snapshot.avgTps != null && snapshot.avgTps > 0.0) {
updatedState = updatedState.copy(latestTps = snapshot.avgTps)
}
return ReducerResult(state = updatedState, streamingState = streamingState)
}
}
@@ -0,0 +1,55 @@
package com.m57.hermescontrol.ui.chat
import java.util.Locale
/**
* Token count estimation and formatting utilities.
*
* The estimation algorithm matches the backend's `estimate_tokens_rough`
* in `agent/model_metadata.py`:
* - CJK ideographs, Hangul, Kana count as ~1 token each.
* - Non-CJK characters count as ceil(UTF-8 bytes / 4).
*/
object TokenEstimator {
fun estimate(text: String?): Int {
if (text.isNullOrEmpty()) return 0
var dense = 0
var nonDenseByteCount = 0
for (i in 0 until text.length) {
val ch = text[i]
if (isCjkDense(ch)) {
dense++
} else {
val code = ch.code
nonDenseByteCount +=
when {
code <= 0x7F -> 1
code <= 0x7FF -> 2
code <= 0xFFFF -> 3
else -> 4
}
}
}
return dense + ((nonDenseByteCount + 3) / 4)
}
private fun isCjkDense(ch: Char): Boolean {
val code = ch.code
return (code in 0x4E00..0x9FFF) ||
(code in 0x3400..0x4DBF) ||
(code in 0xAC00..0xD7AF) ||
(code in 0x1100..0x11FF) ||
(code in 0x3040..0x309F) ||
(code in 0x30A0..0x30FF)
}
fun formatTokenCount(tokens: Int): String =
when {
tokens >= 1_000_000 -> String.format(Locale.US, "%.1fM", tokens / 1_000_000.0)
tokens >= 10_000 -> String.format(Locale.US, "%.1fk", tokens / 1_000.0)
tokens >= 1_000 -> String.format(Locale.US, "%,d", tokens)
else -> tokens.toString()
}
fun formatTps(tps: Double): String = Math.round(tps).toString()
}
@@ -40,6 +40,7 @@ import com.m57.hermescontrol.ui.chat.ChatMessage
import com.m57.hermescontrol.ui.chat.ImageViewerModel
import com.m57.hermescontrol.ui.chat.InlineAttachment
import com.m57.hermescontrol.ui.chat.MarkdownText
import com.m57.hermescontrol.ui.chat.TokenEstimator
import com.m57.hermescontrol.ui.chat.components.ReasoningCard
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
@@ -171,6 +172,44 @@ internal fun FullBleedAgentMessage(
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
modifier = Modifier.testTag("fullbleed_finish_time"),
)
if (message.tokenCount != null && message.tokenCount > 0) {
Spacer(modifier = Modifier.width(6.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.5f),
)
Spacer(modifier = Modifier.width(6.dp))
Text(
text =
stringResource(
R.string.chat_msg_tokens,
TokenEstimator.formatTokenCount(message.tokenCount),
),
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
modifier = Modifier.testTag("fullbleed_token_count"),
)
}
if (message.tps != null && message.tps > 0.0) {
Spacer(modifier = Modifier.width(6.dp))
Text(
text = "•",
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.5f),
)
Spacer(modifier = Modifier.width(6.dp))
Text(
text =
stringResource(
R.string.chat_msg_tps,
TokenEstimator.formatTps(message.tps),
),
style = MaterialTheme.typography.labelSmall,
color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
modifier = Modifier.testTag("fullbleed_tps"),
)
}
}
}
}