kami
e05532e7b2
- GlobalStreamHandler.handleCreateGrant validates scope (SESSION/STAGE) and sends ProtocolError to client on validation failure instead of silent log.warn + drop. - Derive event stageId directly from GrantScope type, removing redundant stageIdForEvent tuple. - SessionOrchestrator integrates ApprovalEngine to check active grants before requesting user approval for T2+ tool calls. - ContextEntry gains EntryRole (SYSTEM/USER/ASSISTANT/TOOL) field for proper chat message role mapping. - RouterContextBuilder.build is now suspend; uses Tokenizer for accurate token estimation with fallback to character-based estimate. - LlamaCppInferenceProvider maps EntryRole to ChatMessage role instead of heuristic layer-based inference.
592 lines
25 KiB
Kotlin
592 lines
25 KiB
Kotlin
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 kotlinx.coroutines.runBlocking
|
|
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 fun buildPack(state: RouterState, budget: TokenBudget): ContextPack = runBlocking {
|
|
buildPack(state, budget)
|
|
}
|
|
|
|
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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(state, TokenBudget(limit = 10000))
|
|
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
|
|
assertEquals(1, l1Entries.size)
|
|
assertEquals("user message", l1Entries[0].content)
|
|
assertEquals(EntryRole.USER, l1Entries[0].role)
|
|
}
|
|
|
|
@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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(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 = buildPack(state, TokenBudget(limit = 1000))
|
|
assertTrue(pack.id.value.contains("unknown"))
|
|
}
|
|
}
|