fix(chat): bind tool milestone dividers to message id and exclude synthetic rows (#1095)

This commit is contained in:
M57 authored and GitHub committed 2026-09-12 16:29:40 +03:00
1 parent 065c7fd75f
commit 47b6955c9a
3 files changed
+108 -32

No files matched your search

@@ -25,33 +25,66 @@ import androidx.compose.ui.unit.sp
*/
const val TOOL_CALL_DIVIDER_INTERVAL = 5
private const val MAX_ITERATIONS_SYSTEM_MARKER =
"You've reached the maximum number of tool-calling iterations allowed."
/**
* Returns the message indices that complete a tool-call milestone — a [MessageRole.TOOL]
* A genuine user message that establishes a new turn boundary.
* Timeline markers (`displayKind != null`) and synthetic system rows
* (such as max-iterations nudges) ride with `role == MessageRole.USER`
* in storage but belong to the surrounding agent turn.
*/
internal fun ChatMessage.isUserTurnBoundary(): Boolean =
role == MessageRole.USER &&
displayKind == null &&
!content.startsWith(MAX_ITERATIONS_SYSTEM_MARKER)
/**
* Returns the message IDs that complete a tool-call milestone — a [MessageRole.TOOL]
* message whose per-turn tool-call count is a multiple of [TOOL_CALL_DIVIDER_INTERVAL] —
* mapped to the per-turn count at that point.
*
* The counter resets at each [MessageRole.USER] message (a new turn), so
* The counter resets at each genuine user turn boundary ([isUserTurnBoundary]), so
* milestones read `5, 10, 15...` within a turn and restart at `5` on the next.
*/
fun toolCallMilestones(messages: List<ChatMessage>): Map<Int, Int> {
val milestones = mutableMapOf<Int, Int>()
fun toolCallMilestones(messages: List<ChatMessage>): Map<String, Int> {
val milestones = mutableMapOf<String, Int>()
var toolCount = 0
messages.forEachIndexed { index, message ->
when (message.role) {
MessageRole.USER -> {
messages.forEach { message ->
when {
message.isUserTurnBoundary() -> {
toolCount = 0
}
MessageRole.TOOL -> {
message.role == MessageRole.TOOL -> {
toolCount += 1
if (toolCount % TOOL_CALL_DIVIDER_INTERVAL == 0) {
milestones[message.id] = toolCount
}
}
}
}
return milestones
}
/**
* Index-based milestone map for position-based callers and backwards compatibility.
*/
fun toolCallMilestoneIndices(messages: List<ChatMessage>): Map<Int, Int> {
val milestones = mutableMapOf<Int, Int>()
var toolCount = 0
messages.forEachIndexed { index, message ->
when {
message.isUserTurnBoundary() -> {
toolCount = 0
}
message.role == MessageRole.TOOL -> {
toolCount += 1
if (toolCount % TOOL_CALL_DIVIDER_INTERVAL == 0) {
milestones[index] = toolCount
}
}
else -> {
Unit
}
}
}
return milestones
@@ -85,7 +85,7 @@ fun FullBleedChatList(
)
}
} else {
val toolMilestones = toolCallMilestones(messages)
val toolMilestones = remember(messages) { toolCallMilestones(messages) }
val turns =
remember(messages, streamingMessage) {
groupIntoTurnsWithStreaming(messages, streamingMessage)
@@ -148,14 +148,12 @@ fun FullBleedChatList(
}
}
var entryIndex = 0
turns.forEach { turn ->
when (turn) {
is ChatTurn.User -> {
// Eager captures: item lambda reads these at
// composition time (lazy), so capture now.
val userMessage = turn.message
val milestone = toolMilestones[entryIndex]
item(key = "user-${userMessage.id}") {
Column(modifier = Modifier.padding(bottom = 12.dp)) {
renderChatBubble(
@@ -170,12 +168,8 @@ fun FullBleedChatList(
openingAttachmentPath = openingAttachmentPath,
onImageClick = onImageClick,
)
milestone?.let { count ->
ToolCallDivider(count = count, maxPerTurn = maxToolCallsPerTurn)
}
}
}
entryIndex++
}
is ChatTurn.Agent -> {
@@ -210,7 +204,6 @@ fun FullBleedChatList(
val hoistedReasoning =
turnReasoning != null &&
proseMessage.id == turnReasoning.message.id
val milestone = toolMilestones[entryIndex]
item(key = "prose-${proseMessage.id}") {
Column(modifier = Modifier.padding(bottom = 12.dp)) {
if (proseMessage.isStreaming && typingEffectEnabled) {
@@ -248,18 +241,14 @@ fun FullBleedChatList(
onImageClick = onImageClick,
)
}
milestone?.let { count ->
ToolCallDivider(count = count, maxPerTurn = maxToolCallsPerTurn)
}
}
}
firstProseSeen = true
entryIndex++
}
is AgentEntry.ToolRow -> {
val toolMessage = entry.message
val milestone = toolMilestones[entryIndex]
val milestone = toolMilestones[toolMessage.id]
item(key = "tool-${toolMessage.id}") {
Column(modifier = Modifier.padding(bottom = 6.dp)) {
FullBleedToolRow(toolMessage)
@@ -268,7 +257,6 @@ fun FullBleedChatList(
}
}
}
entryIndex++
}
is AgentEntry.SystemEvent -> {
@@ -288,7 +276,6 @@ fun FullBleedChatList(
}
}
}
entryIndex++
}
}
}
@@ -31,13 +31,18 @@ class ChatListEntriesTest {
@Test
fun fifthToolCallIsAMilestone() {
val messages = List(5) { toolMessage() }
assertEquals(mapOf(4 to 5), toolCallMilestones(messages))
assertEquals(mapOf(messages[4].id to 5), toolCallMilestones(messages))
assertEquals(mapOf(4 to 5), toolCallMilestoneIndices(messages))
}
@Test
fun milestoneAtEveryFifthToolCallWithinTurn() {
val messages = List(12) { toolMessage() }
assertEquals(mapOf(4 to 5, 9 to 10), toolCallMilestones(messages))
assertEquals(
mapOf(messages[4].id to 5, messages[9].id to 10),
toolCallMilestones(messages),
)
assertEquals(mapOf(4 to 5, 9 to 10), toolCallMilestoneIndices(messages))
}
@Test
@@ -57,7 +62,56 @@ class ChatListEntriesTest {
toolMessage(), // 10
toolMessage(), // 11 ← 5th tool call of turn 2 (not 10)
)
assertEquals(mapOf(5 to 5, 11 to 5), toolCallMilestones(messages))
assertEquals(
mapOf(messages[5].id to 5, messages[11].id to 5),
toolCallMilestones(messages),
)
assertEquals(mapOf(5 to 5, 11 to 5), toolCallMilestoneIndices(messages))
}
@Test
fun timelineMarkerDoesNotResetToolCounter() {
val marker =
ChatMessage(
role = MessageRole.USER,
content = "[System: The active model has changed to gpt-5]",
displayKind = "model_switch",
)
val messages =
listOf(
userMessage(),
toolMessage(),
toolMessage(),
marker,
toolMessage(),
toolMessage(),
toolMessage(), // 5th tool call despite intervening marker
)
assertEquals(mapOf(messages[6].id to 5), toolCallMilestones(messages))
assertEquals(mapOf(6 to 5), toolCallMilestoneIndices(messages))
}
@Test
fun syntheticMaxIterationsNudgeDoesNotResetToolCounter() {
val nudge =
ChatMessage(
role = MessageRole.USER,
content =
"You've reached the maximum number of tool-calling iterations allowed. " +
"Please provide a summary.",
)
val messages =
listOf(
userMessage(),
toolMessage(),
toolMessage(),
toolMessage(),
nudge,
toolMessage(),
toolMessage(), // 5th tool call despite intervening nudge
)
assertEquals(mapOf(messages[6].id to 5), toolCallMilestones(messages))
assertEquals(mapOf(6 to 5), toolCallMilestoneIndices(messages))
}
@Test
@@ -73,13 +127,15 @@ class ChatListEntriesTest {
textMessage(), // 6
toolMessage(), // 7 ← 5th tool call
)
assertEquals(mapOf(7 to 5), toolCallMilestones(messages))
assertEquals(mapOf(messages[7].id to 5), toolCallMilestones(messages))
assertEquals(mapOf(7 to 5), toolCallMilestoneIndices(messages))
}
@Test
fun toolMessagesWithoutToolNameStillCount() {
val messages = List(5) { ChatMessage(role = MessageRole.TOOL, content = "", toolName = null) }
assertEquals(mapOf(4 to 5), toolCallMilestones(messages))
assertEquals(mapOf(messages[4].id to 5), toolCallMilestones(messages))
assertEquals(mapOf(4 to 5), toolCallMilestoneIndices(messages))
}
@Test