fix(context): preserve chronological turn order through pack assembly
Tool-calling loops were rendered as all-assistant-calls-then-all-tool-results with the original task last, causing the model to re-issue the same tool call indefinitely. Two stacked reorderings caused it: - DefaultContextPackBuilder grouped entries by sourceType for per-type compression, discarding a,t,a,t interleaving. - PromptRenderer forced L1 (the live user turn) to render last, pushing the task after the entire transcript. - SessionOrchestrator's tool loop re-fed currentContext.layers.values.flatten() (grouped by layer) each round, compounding the scramble. Add a chronological `ordinal` to ContextEntry, stamped by the builder from input order and restored after grouping/compression (the compressor preserves entry identity, so ordinals survive). PromptRenderer now orders non-system messages by ordinal, with the old L1-last layer priority kept only as a tiebreak so router chat (ordinal 0) is unchanged. The orchestrator keeps a running ordered accumulator instead of reading back the grouped pack. Adds builder + renderer ordering regression tests.
This commit is contained in:
+9
-2
@@ -23,7 +23,12 @@ class DefaultContextPackBuilder(
|
||||
entries: List<ContextEntry>,
|
||||
budget: TokenBudget
|
||||
): ContextPack {
|
||||
val (pinned, compressible) = entries.partition { it.sourceType in neverDropSourceTypes }
|
||||
// Stamp chronological order from the caller's input sequence. Grouping by
|
||||
// sourceType (for compression) and by layer reorders entries; the ordinal lets
|
||||
// us restore true turn order afterwards so a tool loop reads assistant→tool→…
|
||||
// instead of all-assistants-then-all-tools.
|
||||
val ordered = entries.mapIndexed { index, entry -> entry.copy(ordinal = index) }
|
||||
val (pinned, compressible) = ordered.partition { it.sourceType in neverDropSourceTypes }
|
||||
val pinnedTokens = pinned.sumOf { it.tokenEstimate }
|
||||
var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)
|
||||
|
||||
@@ -41,7 +46,9 @@ class DefaultContextPackBuilder(
|
||||
result
|
||||
}
|
||||
|
||||
val retained = pinned + compressed
|
||||
// Restore chronological order before layering — groupBy preserves encounter
|
||||
// order, so each layer's list comes out in true turn order.
|
||||
val retained = (pinned + compressed).sortedBy { it.ordinal }
|
||||
val layers = retained.groupBy { it.layer }
|
||||
val budgetUsed = retained.sumOf { it.tokenEstimate }
|
||||
val droppedCount = entries.size - retained.size
|
||||
|
||||
@@ -14,4 +14,10 @@ data class ContextEntry(
|
||||
val sourceId: String,
|
||||
val tokenEstimate: Int,
|
||||
val role: EntryRole = EntryRole.USER,
|
||||
// Monotonic chronological position within a single context assembly. Stamped by
|
||||
// DefaultContextPackBuilder from input order so that grouping/compression (which
|
||||
// reorder by sourceType/layer) and PromptRenderer can recover true turn order.
|
||||
// Entries built outside that builder (e.g. router chat) leave this at 0 and fall
|
||||
// back to layer-priority ordering in PromptRenderer.
|
||||
val ordinal: Int = 0,
|
||||
)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
package com.correx.core.inference
|
||||
|
||||
import com.correx.core.context.model.ContextEntry
|
||||
import com.correx.core.context.model.ContextLayer
|
||||
import com.correx.core.context.model.ContextPack
|
||||
import com.correx.core.context.model.EntryRole
|
||||
@@ -12,12 +13,13 @@ data class ChatMessage(
|
||||
)
|
||||
|
||||
object PromptRenderer {
|
||||
// L1 (live conversation, ends with the current user turn) must render last so
|
||||
// the model template sees the user query as the final message. Background/memory
|
||||
// layers (L2, L3, L4) are injected before it.
|
||||
private val nonL0RenderOrder: Comparator<ContextLayer> = Comparator { a, b ->
|
||||
val priority = { layer: ContextLayer -> if (layer == ContextLayer.L1) Int.MAX_VALUE else layer.ordinal }
|
||||
priority(a) - priority(b)
|
||||
// Tiebreak only: when entries carry no chronological ordinal (all 0 — e.g. router
|
||||
// chat, which assembles its pack directly), fall back to the old layer priority that
|
||||
// renders L1 (the live user turn) last so the template sees a user query at the end.
|
||||
// When ordinals are present (orchestrator stage path), they dominate and produce true
|
||||
// chronological order — a tool loop reads task → assistant → tool → assistant → …
|
||||
private val layerPriority: (ContextLayer) -> Int = { layer ->
|
||||
if (layer == ContextLayer.L1) Int.MAX_VALUE else layer.ordinal
|
||||
}
|
||||
|
||||
fun render(contextPack: ContextPack): List<ChatMessage> {
|
||||
@@ -29,24 +31,23 @@ object PromptRenderer {
|
||||
.takeIf { it.isNotBlank() }
|
||||
val conversationMessages = contextPack.layers.entries
|
||||
.filter { it.key != ContextLayer.L0 }
|
||||
.sortedWith(Comparator.comparing({ it.key }, nonL0RenderOrder))
|
||||
.flatMap { (_, entries) ->
|
||||
entries.map { entry ->
|
||||
ChatMessage(
|
||||
role = when (entry.role) {
|
||||
EntryRole.SYSTEM -> "system"
|
||||
EntryRole.ASSISTANT -> "assistant"
|
||||
EntryRole.TOOL -> "tool"
|
||||
EntryRole.USER -> "user"
|
||||
},
|
||||
content = entry.content,
|
||||
)
|
||||
}
|
||||
}
|
||||
.flatMap { (layer, entries) -> entries.map { layer to it } }
|
||||
.sortedWith(compareBy({ it.second.ordinal }, { layerPriority(it.first) }))
|
||||
.map { (_, entry) -> entry.toChatMessage() }
|
||||
val messages = buildList {
|
||||
systemContent?.let { add(ChatMessage("system", it)) }
|
||||
addAll(conversationMessages)
|
||||
}
|
||||
return messages.ifEmpty { listOf(ChatMessage("user", "")) }
|
||||
}
|
||||
|
||||
private fun ContextEntry.toChatMessage(): ChatMessage = ChatMessage(
|
||||
role = when (role) {
|
||||
EntryRole.SYSTEM -> "system"
|
||||
EntryRole.ASSISTANT -> "assistant"
|
||||
EntryRole.TOOL -> "tool"
|
||||
EntryRole.USER -> "user"
|
||||
},
|
||||
content = content,
|
||||
)
|
||||
}
|
||||
|
||||
+8
-3
@@ -281,11 +281,16 @@ abstract class SessionOrchestrator(
|
||||
val schemaEntries = buildSchemaEntries(responseFormat, stageId)
|
||||
val steeringEntries = buildSteeringNoteEntries(sessionId)
|
||||
|
||||
// Maintain the running transcript in true chronological order. Reading it back
|
||||
// from currentContext.layers (a Map grouped by layer) would scramble turn order
|
||||
// across rounds; instead we grow our own ordered list and let the builder restamp
|
||||
// ordinals from it each round.
|
||||
var accumulatedEntries = systemPrompt + schemaEntries + promptEntries + steeringEntries
|
||||
val contextPack = contextPackBuilder.build(
|
||||
id = ContextPackId(UUID.randomUUID().toString()),
|
||||
sessionId = sessionId,
|
||||
stageId = stageId,
|
||||
entries = systemPrompt + schemaEntries + promptEntries + steeringEntries,
|
||||
entries = accumulatedEntries,
|
||||
budget = TokenBudget(limit = stageConfig.tokenBudget),
|
||||
)
|
||||
emitContextTruncationIfNeeded(sessionId, stageId, contextPack)
|
||||
@@ -324,12 +329,12 @@ abstract class SessionOrchestrator(
|
||||
retryable = true,
|
||||
)
|
||||
}
|
||||
val allEntries = currentContext.layers.values.flatten() + toolEntries
|
||||
accumulatedEntries = accumulatedEntries + toolEntries
|
||||
currentContext = contextPackBuilder.build(
|
||||
id = ContextPackId(UUID.randomUUID().toString()),
|
||||
sessionId = sessionId,
|
||||
stageId = stageId,
|
||||
entries = allEntries,
|
||||
entries = accumulatedEntries,
|
||||
budget = TokenBudget(limit = stageConfig.tokenBudget),
|
||||
)
|
||||
emitContextTruncationIfNeeded(sessionId, stageId, currentContext)
|
||||
|
||||
Reference in New Issue
Block a user