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 `retry repair mandate renders as the final user turn after the tool transcript`() = kotlinx.coroutines.runBlocking {
        // #293: even though the repair mandate is stamped mid-transcript by input order, the renderer
        // must lift it to the very end (after assistant/tool evidence) and never fold it into system.
        val builder = DefaultContextPackBuilder(DefaultContextCompressor())
        val entries = listOf(
            entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt"),
            entry("task", ContextLayer.L1, EntryRole.USER, "agentPrompt"),
            entry("repair", ContextLayer.L1, EntryRole.USER, "retryFeedback"),
            entry("call1", ContextLayer.L2, EntryRole.ASSISTANT, "assistantToolCall", sourceId = "tool1"),
            entry("result1", ContextLayer.L2, EntryRole.TOOL, "toolResult", sourceId = "tool1"),
        )
        val pack = builder.build(ContextPackId("p"), sessionId, stageId, entries, TokenBudget(limit = 4000))

        val messages = PromptRenderer.render(pack)
        assertEquals("user" to "repair", messages.last().role to messages.last().content)
        assertEquals(1, messages.count { it.content == "repair" }, "mandate must appear exactly once")
        org.junit.jupiter.api.Assertions.assertFalse(
            messages.first().content.contains("repair"),
            "leading system must not carry the repair mandate",
        )
        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 `USER-role context entries render as separate turns and never fold into system`() {
        // #290: intent, repo map, docs catalog and decision journal carry EntryRole.USER so the
        // leading system block stays pure policy/schema. They must appear as user turns, not merge in.
        val pack = ContextPack(
            id = ContextPackId("p"),
            sessionId = sessionId,
            stageId = stageId,
            layers = mapOf(
                ContextLayer.L0 to listOf(entry("policy", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt")),
                ContextLayer.L1 to listOf(entry("intent", ContextLayer.L1, EntryRole.USER, "initialIntent")),
                ContextLayer.L3 to listOf(entry("journal", ContextLayer.L3, EntryRole.USER, "decisionJournal")),
            ),
            budgetUsed = 30,
            budgetLimit = 4000,
        )
        val messages = PromptRenderer.render(pack)
        assertEquals("policy", messages.first().content, "leading system must be policy only")
        org.junit.jupiter.api.Assertions.assertEquals(1, messages.count { it.role == "system" })
        org.junit.jupiter.api.Assertions.assertTrue(messages.any { it.role == "user" && it.content == "intent" })
        org.junit.jupiter.api.Assertions.assertTrue(messages.any { it.role == "user" && it.content == "journal" })
    }

    @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 },
        )
    }
}