correx/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt

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, sourceId: String = id) =
        ContextEntry(
            id = ContextEntryId(id),
            layer = layer,
            content = id,
            sourceType = sourceType,
            sourceId = sourceId,
            tokenEstimate = 10,
            role = role,
        )

    @Test
    fun `stage tool loop renders chronologically with the task before tool turns`() = kotlinx.coroutines.runBlocking {
        // 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", sourceId = "tool1"),
            entry("result1", ContextLayer.L2, EntryRole.TOOL, "toolResult", sourceId = "tool1"),
            entry("call2", ContextLayer.L2, EntryRole.ASSISTANT, "assistantToolCall", sourceId = "tool2"),
            entry("result2", ContextLayer.L2, EntryRole.TOOL, "toolResult", sourceId = "tool2"),
        )
        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 },
        )
        Unit
    }

    @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("system" to "memory", "user" to "question"), messages.map { it.role to it.content })
    }

    @Test
    fun `steering directives are re-anchored as a trailing user reminder`() {
        // A steering note folds into the leading system block AND is restated at the tail so a
        // weak local model still sees the active constraint next to the final query.
        val pack = ContextPack(
            id = ContextPackId("p"),
            sessionId = sessionId,
            stageId = stageId,
            layers = mapOf(
                ContextLayer.L0 to listOf(entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt")),
                ContextLayer.L2 to listOf(entry("directive", ContextLayer.L2, EntryRole.SYSTEM, "steeringNote")),
                ContextLayer.L1 to listOf(entry("question", ContextLayer.L1, EntryRole.USER, "chat")),
            ),
            budgetUsed = 30,
            budgetLimit = 4000,
        )
        val messages = PromptRenderer.render(pack)
        assertEquals(
            listOf(
                "system" to "sys\n\ndirective",
                "user" to "question",
                "user" to "Reminder — active steering directive(s):\ndirective",
            ),
            messages.map { it.role to it.content },
        )
    }

    @Test
    fun `non-L0 SYSTEM entries fold into the single leading system message`() {
        // Strict chat templates (Qwen) 400 on any system message after index 0. Recalled L3
        // memory and L2 summaries carry EntryRole.SYSTEM and must merge into the system block.
        val pack = ContextPack(
            id = ContextPackId("p"),
            sessionId = sessionId,
            stageId = stageId,
            layers = mapOf(
                ContextLayer.L0 to listOf(entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt")),
                ContextLayer.L1 to listOf(entry("question", ContextLayer.L1, EntryRole.USER, "chat")),
                ContextLayer.L3 to listOf(entry("recalled", ContextLayer.L3, EntryRole.SYSTEM, "recalledMemory")),
            ),
            budgetUsed = 30,
            budgetLimit = 4000,
        )
        val messages = PromptRenderer.render(pack)
        assertEquals(
            listOf("system" to "sys\n\nrecalled", "user" to "question"),
            messages.map { it.role to it.content },
        )
    }
}