diff --git a/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt b/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt index ad5091ad..14516d7b 100644 --- a/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt +++ b/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt @@ -23,7 +23,12 @@ class DefaultContextPackBuilder( entries: List, 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 diff --git a/core/context/src/main/kotlin/com/correx/core/context/model/ContextEntry.kt b/core/context/src/main/kotlin/com/correx/core/context/model/ContextEntry.kt index ed2f53bc..016aa326 100644 --- a/core/context/src/main/kotlin/com/correx/core/context/model/ContextEntry.kt +++ b/core/context/src/main/kotlin/com/correx/core/context/model/ContextEntry.kt @@ -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, ) diff --git a/core/inference/src/main/kotlin/com/correx/core/inference/PromptRenderer.kt b/core/inference/src/main/kotlin/com/correx/core/inference/PromptRenderer.kt index cbd7b767..a52d821a 100644 --- a/core/inference/src/main/kotlin/com/correx/core/inference/PromptRenderer.kt +++ b/core/inference/src/main/kotlin/com/correx/core/inference/PromptRenderer.kt @@ -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 = 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 { @@ -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, + ) } diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index 9b1e6bd5..1a495270 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt @@ -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) diff --git a/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt b/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt index 8a5a6c5f..4f867479 100644 --- a/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt +++ b/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt @@ -94,4 +94,23 @@ class DefaultContextPackBuilderTest { assertTrue(allRetained.any { it.id == ContextEntryId("steering") }, "steeringNote must be retained") assertTrue(allRetained.any { it.id == ContextEntryId("history") }, "eventHistory must be retained") } + + @Test + fun `tool dialogue keeps chronological order despite sourceType grouping`() { + // A multi-round tool loop: assistant call then tool result, interleaved. Compression + // groups by sourceType internally; the builder must restore input order so the L2 + // layer reads call, result, call, result — not all calls then all results. + val entries = listOf( + entry("task", ContextLayer.L1, 10).copy(sourceType = "agentPrompt"), + entry("call1", ContextLayer.L2, 10).copy(sourceType = "assistantToolCall"), + entry("result1", ContextLayer.L2, 10).copy(sourceType = "toolResult"), + entry("call2", ContextLayer.L2, 10).copy(sourceType = "assistantToolCall"), + entry("result2", ContextLayer.L2, 10).copy(sourceType = "toolResult"), + ) + val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 4000)) + val l2Ids = pack.layers[ContextLayer.L2]!!.map { it.id.value } + assertEquals(listOf("call1", "result1", "call2", "result2"), l2Ids) + // ordinals are stamped from input order + assertEquals(listOf(1, 2, 3, 4), pack.layers[ContextLayer.L2]!!.map { it.ordinal }) + } } diff --git a/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt b/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt new file mode 100644 index 00000000..7e6c6553 --- /dev/null +++ b/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt @@ -0,0 +1,72 @@ +import com.correx.core.context.builder.DefaultContextPackBuilder +import com.correx.core.context.compression.DefaultContextCompressor +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 +import com.correx.core.context.model.TokenBudget +import com.correx.core.events.types.ContextEntryId +import com.correx.core.events.types.ContextPackId +import com.correx.core.events.types.SessionId +import com.correx.core.events.types.StageId +import com.correx.core.inference.PromptRenderer +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Test + +class PromptRendererOrderingTest { + + private val sessionId = SessionId("s1") + private val stageId = StageId("stage-1") + + private fun entry(id: String, layer: ContextLayer, role: EntryRole, sourceType: String) = ContextEntry( + id = ContextEntryId(id), + layer = layer, + content = id, + sourceType = sourceType, + sourceId = id, + tokenEstimate = 10, + role = role, + ) + + @Test + fun `stage tool loop renders chronologically with the task before tool turns`() { + // The builder stamps ordinals from input order; the renderer must honour them so the + // model sees: system, task, assistant call, tool result, assistant call, tool result. + val builder = DefaultContextPackBuilder(DefaultContextCompressor()) + val entries = listOf( + entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt"), + entry("task", ContextLayer.L1, EntryRole.USER, "agentPrompt"), + entry("call1", ContextLayer.L2, EntryRole.ASSISTANT, "assistantToolCall"), + entry("result1", ContextLayer.L2, EntryRole.TOOL, "toolResult"), + entry("call2", ContextLayer.L2, EntryRole.ASSISTANT, "assistantToolCall"), + entry("result2", ContextLayer.L2, EntryRole.TOOL, "toolResult"), + ) + val pack = builder.build(ContextPackId("p"), sessionId, stageId, entries, TokenBudget(limit = 4000)) + + val messages = PromptRenderer.render(pack) + assertEquals( + listOf("system" to "sys", "user" to "task", "assistant" to "call1", + "tool" to "result1", "assistant" to "call2", "tool" to "result2"), + messages.map { it.role to it.content }, + ) + } + + @Test + fun `without ordinals the live user turn still renders last (router chat fallback)`() { + // Packs built outside DefaultContextPackBuilder (router chat) leave ordinal at 0; + // the renderer falls back to layer priority so L1 (the user query) lands last. + val pack = ContextPack( + id = ContextPackId("p"), + sessionId = sessionId, + stageId = stageId, + layers = mapOf( + ContextLayer.L1 to listOf(entry("question", ContextLayer.L1, EntryRole.USER, "chat")), + ContextLayer.L2 to listOf(entry("memory", ContextLayer.L2, EntryRole.SYSTEM, "l2")), + ), + budgetUsed = 20, + budgetLimit = 4000, + ) + val messages = PromptRenderer.render(pack) + assertEquals(listOf("memory", "question"), messages.map { it.content }) + } +}