import com.correx.core.context.model.ContextLayer import com.correx.core.context.model.TokenBudget import com.correx.core.events.types.SessionId import com.correx.core.events.types.StageId import com.correx.core.router.DefaultRouterContextBuilder import com.correx.core.router.model.RouterConfig import com.correx.core.router.model.RouterL2Entry import com.correx.core.router.model.RouterState import com.correx.core.router.model.RouterTurn import com.correx.core.router.model.StageOutcomeKind import com.correx.core.router.model.TurnRole import com.correx.core.router.model.WorkflowStatus import kotlinx.datetime.Clock import kotlinx.datetime.Instant import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertNotNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test class RouterContextBuilderTest { private val config = RouterConfig(conversationKeepLast = 3, tokenBudget = TokenBudget(limit = 200)) private val builder = DefaultRouterContextBuilder(config) private val sessionId = SessionId("test-session") private val stageId = StageId("stage-1") private val clock = Clock.System // -------------------------------------------------------------------------- // Budget enforcement // -------------------------------------------------------------------------- @Test fun `build fits within budget when all entries are small`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "hi", clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) assertEquals(10000, pack.budgetLimit) assertTrue(pack.budgetUsed <= pack.budgetLimit) } @Test fun `build drops entries when budget is exceeded`() { // Build a state where conversation + L2 entries together exceed the tight budget val longContent = "x".repeat(800) // ~400 tokens each val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, longContent, clock.now()), RouterTurn(TurnRole.ROUTER, longContent, clock.now()), RouterTurn(TurnRole.USER, longContent, clock.now()), ), l2Memory = listOf( RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 100)) // L0 entries (system prompt + workflow status) always fit; L1/L2 should be dropped assertTrue(pack.compressionMetadata.entriesDropped > 0) } @Test fun `build drops entries oldest-first for L2 memory`() { val short = "ok" val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, l2Memory = listOf( RouterL2Entry(StageId("s1"), short, StageOutcomeKind.SUCCESS, Instant.parse("2026-01-01T00:00:00Z")), RouterL2Entry(StageId("s2"), short, StageOutcomeKind.SUCCESS, Instant.parse("2026-01-02T00:00:00Z")), RouterL2Entry(StageId("s3"), short, StageOutcomeKind.SUCCESS, Instant.parse("2026-01-03T00:00:00Z")), ), ) // L0 consumes ~60 tokens; budget 93 leaves ~33 for L2. // Each L2 entry is ~12 tokens; 2 fit (s1, s2), s3 is dropped val pack = builder.build(state, TokenBudget(limit = 93)) val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList() // Oldest-first eviction: s1 and s2 fit, s3 is the one dropped val remainingStageIds = l2Entries.map { it.sourceId }.toSet() assertEquals(2, l2Entries.size) assertTrue(remainingStageIds.contains("s1")) assertTrue(remainingStageIds.contains("s2")) assertEquals(0, l2Entries.count { it.sourceId == "s3" }) } @Test fun `build drops conversation turns oldest-first when conversationKeepLast exceeded`() { val configWide = RouterConfig(conversationKeepLast = 2, tokenBudget = TokenBudget(limit = 10000)) val builderWide = DefaultRouterContextBuilder(configWide) val longContent = "x".repeat(400) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, longContent, Instant.parse("2026-01-01T00:00:00Z")), RouterTurn(TurnRole.ROUTER, longContent, Instant.parse("2026-01-02T00:00:00Z")), RouterTurn(TurnRole.USER, longContent, Instant.parse("2026-01-03T00:00:00Z")), RouterTurn(TurnRole.ROUTER, longContent, Instant.parse("2026-01-04T00:00:00Z")), ), ) // Budget allows L0 + 2 long conversation turns; oldest 2 turns are dropped val pack = builderWide.build(state, TokenBudget(limit = 500)) val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList() // Only the last 2 turns (oldest-first drop) should remain assertEquals(2, l1Entries.size) } // -------------------------------------------------------------------------- // L0 immutability // -------------------------------------------------------------------------- @Test fun `L0 system prompt is always present`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.IDLE, currentStageId = null, ) val pack = builder.build(state, TokenBudget(limit = 0)) val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList() assertTrue(l0Entries.any { it.sourceType == "systemPrompt" }) } @Test fun `L0 workflow status is always present`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.FAILED, currentStageId = StageId("failed-stage"), ) val pack = builder.build(state, TokenBudget(limit = 0)) val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList() assertTrue(l0Entries.any { it.sourceType == "workflowStatus" }) val workflowEntry = l0Entries.find { it.sourceType == "workflowStatus" } assertNotNull(workflowEntry) assertTrue(workflowEntry!!.content.contains("failed")) } @Test fun `L0 entries survive zero budget`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.COMPLETED, currentStageId = null, conversationHistory = emptyList(), l2Memory = emptyList(), ) val pack = builder.build(state, TokenBudget(limit = 0)) val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList() assertEquals(2, l0Entries.size) // systemPrompt + workflowStatus } @Test fun `L0 entries survive budget that cannot cover L1 or L2`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "x".repeat(500), clock.now()), ), l2Memory = listOf( RouterL2Entry(StageId("s1"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 50)) val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList() val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList() val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList() assertEquals(2, l0Entries.size) assertEquals(0, l1Entries.size) assertEquals(0, l2Entries.size) assertTrue(pack.compressionMetadata.entriesDropped > 0) } // -------------------------------------------------------------------------- // L2 oldest-first eviction // -------------------------------------------------------------------------- @Test fun `L2 entries are evicted in insertion order (oldest first)`() { val configTight = RouterConfig(conversationKeepLast = 0, tokenBudget = TokenBudget(limit = 10)) val builderTight = DefaultRouterContextBuilder(configTight) val baseTime = Instant.parse("2026-01-01T00:00:00Z") val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, l2Memory = listOf( RouterL2Entry(StageId("s1"), "old", StageOutcomeKind.SUCCESS, baseTime), RouterL2Entry(StageId("s2"), "mid", StageOutcomeKind.SUCCESS, Instant.parse("2026-01-01T01:00:00Z")), RouterL2Entry(StageId("s3"), "new", StageOutcomeKind.SUCCESS, Instant.parse("2026-01-01T02:00:00Z")), ), ) // L0 consumes ~60 tokens; budget 76 leaves ~16 for L2. // Each L2 entry is ~12 tokens; only 1 fits — oldest (s1) survives val pack = builderTight.build(state, TokenBudget(limit = 76)) val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList() // Oldest-first eviction: s1 fits, s2 and s3 are dropped assertEquals(1, l2Entries.size) assertEquals("s1", l2Entries[0].sourceId) } @Test fun `L2 entries fit within budget are all retained`() { val config = RouterConfig(conversationKeepLast = 0, tokenBudget = TokenBudget(limit = 10000)) val builder = DefaultRouterContextBuilder(config) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, l2Memory = listOf( RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()), RouterL2Entry(StageId("s2"), "summary", StageOutcomeKind.FAILURE, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList() assertEquals(2, l2Entries.size) assertEquals(0, pack.compressionMetadata.entriesDropped) } // -------------------------------------------------------------------------- // Layer classification // -------------------------------------------------------------------------- @Test fun `system prompt is classified as L0`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.IDLE, ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l0Entries = pack.layers[ContextLayer.L0] assertNotNull(l0Entries) assertTrue(l0Entries!!.any { it.sourceType == "systemPrompt" && it.layer == ContextLayer.L0 }) } @Test fun `workflow status is classified as L0`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l0Entries = pack.layers[ContextLayer.L0] assertTrue(l0Entries!!.any { it.sourceType == "workflowStatus" && it.layer == ContextLayer.L0 }) } @Test fun `conversation turns are classified as L1`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "hello", clock.now()), RouterTurn(TurnRole.ROUTER, "hi", clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l1Entries = pack.layers[ContextLayer.L1] assertNotNull(l1Entries) assertEquals(2, l1Entries!!.size) assertTrue(l1Entries.all { it.layer == ContextLayer.L1 }) assertTrue(l1Entries.all { it.sourceType == "conversation" }) } @Test fun `l2 memory entries are classified as L2`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, l2Memory = listOf( RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l2Entries = pack.layers[ContextLayer.L2] assertNotNull(l2Entries) assertEquals(1, l2Entries!!.size) assertTrue(l2Entries.all { it.layer == ContextLayer.L2 }) assertTrue(l2Entries.all { it.sourceType == "stageSummary" }) } @Test fun `no entries other than L0 L1 L2 appear in pack`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "hello", clock.now()), ), l2Memory = listOf( RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val presentLayers = pack.layers.keys assertTrue(presentLayers.containsAll(listOf(ContextLayer.L0, ContextLayer.L1, ContextLayer.L2))) assertFalse(presentLayers.contains(ContextLayer.L3)) assertFalse(presentLayers.contains(ContextLayer.L4)) } // -------------------------------------------------------------------------- // Conversation history capping (conversationKeepLast) // -------------------------------------------------------------------------- @Test fun `conversationHistory respects conversationKeepLast from config`() { val config = RouterConfig(conversationKeepLast = 2, tokenBudget = TokenBudget(limit = 10000)) val builder = DefaultRouterContextBuilder(config) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "one", Instant.parse("2026-01-01T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "two", Instant.parse("2026-01-02T00:00:00Z")), RouterTurn(TurnRole.USER, "three", Instant.parse("2026-01-03T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "four", Instant.parse("2026-01-04T00:00:00Z")), RouterTurn(TurnRole.USER, "five", Instant.parse("2026-01-05T00:00:00Z")), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList() // conversationKeepLast=2, so only the last 2 turns should appear assertEquals(2, l1Entries.size) } @Test fun `conversationHistory takeLast preserves ordering`() { val config = RouterConfig(conversationKeepLast = 3, tokenBudget = TokenBudget(limit = 10000)) val builder = DefaultRouterContextBuilder(config) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "first", Instant.parse("2026-01-01T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "second", Instant.parse("2026-01-02T00:00:00Z")), RouterTurn(TurnRole.USER, "third", Instant.parse("2026-01-03T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "fourth", Instant.parse("2026-01-04T00:00:00Z")), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList() assertEquals(3, l1Entries.size) // The last 3: second, third, fourth assertTrue(l1Entries.any { it.content.contains("second") }) assertTrue(l1Entries.any { it.content.contains("third") }) assertTrue(l1Entries.any { it.content.contains("fourth") }) } // -------------------------------------------------------------------------- // Empty state // -------------------------------------------------------------------------- @Test fun `build with empty state produces L0 only`() { val state = RouterState() val pack = builder.build(state, TokenBudget(limit = 10000)) val l0Entries = pack.layers[ContextLayer.L0] assertNotNull(l0Entries) assertEquals(2, l0Entries!!.size) assertTrue(pack.layers[ContextLayer.L1].isNullOrEmpty()) assertTrue(pack.layers[ContextLayer.L2].isNullOrEmpty()) } @Test fun `pack contains correct metadata on empty state`() { val state = RouterState() val pack = builder.build(state, TokenBudget(limit = 5000)) assertEquals(5000, pack.budgetLimit) assertEquals(0, pack.compressionMetadata.entriesDropped) assertEquals(listOf("L0Immutable", "Conversation"), pack.compressionMetadata.appliedStrategies) assertTrue(pack.compressionMetadata.truncatedLayers.isEmpty()) } @Test fun `pack has correct context pack id and session info`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, ) val pack = builder.build(state, TokenBudget(limit = 1000)) assertEquals("test-session-router-pack", pack.id.value) assertEquals(sessionId, pack.sessionId) assertEquals(stageId, pack.stageId) } // -------------------------------------------------------------------------- // Content formatting // -------------------------------------------------------------------------- @Test fun `system prompt has expected content`() { val state = RouterState() val pack = builder.build(state, TokenBudget(limit = 10000)) val systemEntry = pack.layers[ContextLayer.L0]?.find { it.sourceType == "systemPrompt" } assertNotNull(systemEntry) assertEquals( "You are a routing assistant. Provide guidance based on workflow state and conversation context.", systemEntry!!.content, ) } @Test fun `workflow status entry contains status and stage`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.COMPLETED, currentStageId = null, ) val pack = builder.build(state, TokenBudget(limit = 10000)) val workflowEntry = pack.layers[ContextLayer.L0]?.find { it.sourceType == "workflowStatus" } assertNotNull(workflowEntry) assertTrue(workflowEntry!!.content.contains("Status: Completed")) assertTrue(workflowEntry.content.contains("Stage: none")) } @Test fun `conversation turn entry is formatted with role prefix`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "user message", clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList() assertEquals(1, l1Entries.size) assertEquals("[USER] user message", l1Entries[0].content) } @Test fun `l2 entry is formatted with stage outcome and summary`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, l2Memory = listOf( RouterL2Entry(StageId("stage-x"), "completed with 3 items", StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList() assertEquals(1, l2Entries.size) val entry = l2Entries[0] assertEquals("stage-x", entry.sourceId) assertTrue(entry.content.contains("Stage stage-x")) assertTrue(entry.content.contains("SUCCESS")) assertTrue(entry.content.contains("completed with 3 items")) } // -------------------------------------------------------------------------- // Token budget accounting // -------------------------------------------------------------------------- @Test fun `budgetUsed equals sum of tokenEstimate of all retained entries`() { val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "hello world", clock.now()), ), l2Memory = listOf( RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) val allEntries = pack.layers.values.flatten() val computedSum = allEntries.sumOf { it.tokenEstimate } assertEquals(computedSum, pack.budgetUsed) } @Test fun `compressionMetadata entriesDropped reflects dropped count`() { val config = RouterConfig(conversationKeepLast = 0, tokenBudget = TokenBudget(limit = 10)) val builder = DefaultRouterContextBuilder(config) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "x".repeat(200), clock.now()), ), l2Memory = listOf( RouterL2Entry(StageId("s1"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()), RouterL2Entry(StageId("s2"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()), ), ) val pack = builder.build(state, TokenBudget(limit = 65)) // L0 consumes ~60 tokens, leaving 5 — both L2 entries (each ~259 tokens) dropped // conversationKeepLast=0 means conversation entry is not included assertEquals(2, pack.compressionMetadata.entriesDropped) } // -------------------------------------------------------------------------- // Integration: full lifecycle // -------------------------------------------------------------------------- @Test fun `build with full state including all layers`() { val config = RouterConfig(conversationKeepLast = 2, tokenBudget = TokenBudget(limit = 10000)) val builder = DefaultRouterContextBuilder(config) val state = RouterState( sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING, currentStageId = stageId, conversationHistory = listOf( RouterTurn(TurnRole.USER, "first turn", Instant.parse("2026-01-01T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "first reply", Instant.parse("2026-01-02T00:00:00Z")), RouterTurn(TurnRole.USER, "second turn", Instant.parse("2026-01-03T00:00:00Z")), RouterTurn(TurnRole.ROUTER, "second reply", Instant.parse("2026-01-04T00:00:00Z")), ), l2Memory = listOf( RouterL2Entry( StageId("s1"), "stage one done", StageOutcomeKind.SUCCESS, Instant.parse("2026-01-01T10:00:00Z"), ), RouterL2Entry( StageId("s2"), "stage two failed", StageOutcomeKind.FAILURE, Instant.parse("2026-01-02T10:00:00Z"), ), ), ) val pack = builder.build(state, TokenBudget(limit = 10000)) // L0: system prompt + workflow status val l0 = pack.layers[ContextLayer.L0]!! assertEquals(2, l0.size) assertEquals("systemPrompt", l0[0].sourceType) assertEquals("workflowStatus", l0[1].sourceType) // L1: last 2 conversation turns val l1 = pack.layers[ContextLayer.L1]!! assertEquals(2, l1.size) assertEquals("conversation", l1[0].sourceType) // L2: both stage summaries fit val l2 = pack.layers[ContextLayer.L2]!! assertEquals(2, l2.size) assertEquals("stageSummary", l2[0].sourceType) // Budget and metadata assertTrue(pack.budgetUsed > 0) assertTrue(pack.budgetUsed <= pack.budgetLimit) assertEquals(0, pack.compressionMetadata.entriesDropped) } @Test fun `build with null sessionId produces unknown pack id`() { val state = RouterState( workflowStatus = WorkflowStatus.IDLE, ) val pack = builder.build(state, TokenBudget(limit = 1000)) assertTrue(pack.id.value.contains("unknown")) } }