feat: implement CreateGrant handler and grant-aware approval engine
- 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.
This commit is contained in:
@@ -1,5 +1,8 @@
|
||||
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
|
||||
@@ -23,6 +26,10 @@ 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
|
||||
@@ -41,7 +48,7 @@ class RouterContextBuilderTest {
|
||||
RouterTurn(TurnRole.USER, "hi", clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
assertEquals(10000, pack.budgetLimit)
|
||||
assertTrue(pack.budgetUsed <= pack.budgetLimit)
|
||||
}
|
||||
@@ -63,7 +70,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 100))
|
||||
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)
|
||||
}
|
||||
@@ -83,7 +90,7 @@ class RouterContextBuilderTest {
|
||||
)
|
||||
// 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 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()
|
||||
@@ -127,7 +134,7 @@ class RouterContextBuilderTest {
|
||||
workflowStatus = WorkflowStatus.IDLE,
|
||||
currentStageId = null,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 0))
|
||||
val pack = buildPack(state, TokenBudget(limit = 0))
|
||||
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
|
||||
assertTrue(l0Entries.any { it.sourceType == "systemPrompt" })
|
||||
}
|
||||
@@ -139,7 +146,7 @@ class RouterContextBuilderTest {
|
||||
workflowStatus = WorkflowStatus.FAILED,
|
||||
currentStageId = StageId("failed-stage"),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 0))
|
||||
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" }
|
||||
@@ -156,7 +163,7 @@ class RouterContextBuilderTest {
|
||||
conversationHistory = emptyList(),
|
||||
l2Memory = emptyList(),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 0))
|
||||
val pack = buildPack(state, TokenBudget(limit = 0))
|
||||
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
|
||||
assertEquals(2, l0Entries.size) // systemPrompt + workflowStatus
|
||||
}
|
||||
@@ -174,7 +181,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s1"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 50))
|
||||
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()
|
||||
@@ -225,7 +232,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s2"), "summary", StageOutcomeKind.FAILURE, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
|
||||
assertEquals(2, l2Entries.size)
|
||||
assertEquals(0, pack.compressionMetadata.entriesDropped)
|
||||
@@ -241,7 +248,7 @@ class RouterContextBuilderTest {
|
||||
sessionId = sessionId,
|
||||
workflowStatus = WorkflowStatus.IDLE,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
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 })
|
||||
@@ -254,7 +261,7 @@ class RouterContextBuilderTest {
|
||||
workflowStatus = WorkflowStatus.RUNNING,
|
||||
currentStageId = stageId,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l0Entries = pack.layers[ContextLayer.L0]
|
||||
assertTrue(l0Entries!!.any { it.sourceType == "workflowStatus" && it.layer == ContextLayer.L0 })
|
||||
}
|
||||
@@ -270,7 +277,7 @@ class RouterContextBuilderTest {
|
||||
RouterTurn(TurnRole.ROUTER, "hi", clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l1Entries = pack.layers[ContextLayer.L1]
|
||||
assertNotNull(l1Entries)
|
||||
assertEquals(2, l1Entries!!.size)
|
||||
@@ -288,7 +295,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l2Entries = pack.layers[ContextLayer.L2]
|
||||
assertNotNull(l2Entries)
|
||||
assertEquals(1, l2Entries!!.size)
|
||||
@@ -309,7 +316,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
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))
|
||||
@@ -336,7 +343,7 @@ class RouterContextBuilderTest {
|
||||
RouterTurn(TurnRole.USER, "five", Instant.parse("2026-01-05T00:00:00Z")),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
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)
|
||||
@@ -357,7 +364,7 @@ class RouterContextBuilderTest {
|
||||
RouterTurn(TurnRole.ROUTER, "fourth", Instant.parse("2026-01-04T00:00:00Z")),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
|
||||
assertEquals(3, l1Entries.size)
|
||||
// The last 3: second, third, fourth
|
||||
@@ -373,7 +380,7 @@ class RouterContextBuilderTest {
|
||||
@Test
|
||||
fun `build with empty state produces L0 only`() {
|
||||
val state = RouterState()
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l0Entries = pack.layers[ContextLayer.L0]
|
||||
assertNotNull(l0Entries)
|
||||
assertEquals(2, l0Entries!!.size)
|
||||
@@ -384,7 +391,7 @@ class RouterContextBuilderTest {
|
||||
@Test
|
||||
fun `pack contains correct metadata on empty state`() {
|
||||
val state = RouterState()
|
||||
val pack = builder.build(state, TokenBudget(limit = 5000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 5000))
|
||||
assertEquals(5000, pack.budgetLimit)
|
||||
assertEquals(0, pack.compressionMetadata.entriesDropped)
|
||||
assertEquals(listOf("L0Immutable", "Conversation"), pack.compressionMetadata.appliedStrategies)
|
||||
@@ -398,7 +405,7 @@ class RouterContextBuilderTest {
|
||||
workflowStatus = WorkflowStatus.RUNNING,
|
||||
currentStageId = stageId,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 1000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 1000))
|
||||
assertEquals("test-session-router-pack", pack.id.value)
|
||||
assertEquals(sessionId, pack.sessionId)
|
||||
assertEquals(stageId, pack.stageId)
|
||||
@@ -411,7 +418,7 @@ class RouterContextBuilderTest {
|
||||
@Test
|
||||
fun `system prompt has expected content`() {
|
||||
val state = RouterState()
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val systemEntry = pack.layers[ContextLayer.L0]?.find { it.sourceType == "systemPrompt" }
|
||||
assertNotNull(systemEntry)
|
||||
assertEquals(
|
||||
@@ -427,7 +434,7 @@ class RouterContextBuilderTest {
|
||||
workflowStatus = WorkflowStatus.COMPLETED,
|
||||
currentStageId = null,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
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"))
|
||||
@@ -444,10 +451,11 @@ class RouterContextBuilderTest {
|
||||
RouterTurn(TurnRole.USER, "user message", clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
|
||||
assertEquals(1, l1Entries.size)
|
||||
assertEquals("[USER] user message", l1Entries[0].content)
|
||||
assertEquals("user message", l1Entries[0].content)
|
||||
assertEquals(EntryRole.USER, l1Entries[0].role)
|
||||
}
|
||||
|
||||
@Test
|
||||
@@ -460,7 +468,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("stage-x"), "completed with 3 items", StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
|
||||
assertEquals(1, l2Entries.size)
|
||||
val entry = l2Entries[0]
|
||||
@@ -487,7 +495,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
val allEntries = pack.layers.values.flatten()
|
||||
val computedSum = allEntries.sumOf { it.tokenEstimate }
|
||||
assertEquals(computedSum, pack.budgetUsed)
|
||||
@@ -509,7 +517,7 @@ class RouterContextBuilderTest {
|
||||
RouterL2Entry(StageId("s2"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 65))
|
||||
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)
|
||||
@@ -548,7 +556,7 @@ class RouterContextBuilderTest {
|
||||
),
|
||||
),
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 10000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 10000))
|
||||
|
||||
// L0: system prompt + workflow status
|
||||
val l0 = pack.layers[ContextLayer.L0]!!
|
||||
@@ -577,7 +585,7 @@ class RouterContextBuilderTest {
|
||||
val state = RouterState(
|
||||
workflowStatus = WorkflowStatus.IDLE,
|
||||
)
|
||||
val pack = builder.build(state, TokenBudget(limit = 1000))
|
||||
val pack = buildPack(state, TokenBudget(limit = 1000))
|
||||
assertTrue(pack.id.value.contains("unknown"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -109,7 +109,7 @@ class RouterFacadeTest {
|
||||
)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter("response"),
|
||||
eventStore = mockStore,
|
||||
@@ -129,7 +129,7 @@ class RouterFacadeTest {
|
||||
RouterState(sessionId = sessionId, workflowStatus = WorkflowStatus.IDLE, currentStageId = null)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter("response"),
|
||||
eventStore = mockStore,
|
||||
@@ -157,7 +157,7 @@ class RouterFacadeTest {
|
||||
)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
capturedStates.add(state)
|
||||
return emptyContextPack()
|
||||
}
|
||||
@@ -190,7 +190,7 @@ class RouterFacadeTest {
|
||||
RouterState(sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
capturedStates.add(state)
|
||||
return emptyContextPack()
|
||||
}
|
||||
@@ -218,7 +218,7 @@ class RouterFacadeTest {
|
||||
fun `state is passed through to context builder`(): Unit = runBlocking {
|
||||
val capturedState = mutableListOf<RouterState>()
|
||||
val mockContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
capturedState.add(state)
|
||||
return emptyContextPack()
|
||||
}
|
||||
@@ -288,7 +288,7 @@ class RouterFacadeTest {
|
||||
)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter,
|
||||
eventStore = mockEventStore(),
|
||||
@@ -319,7 +319,7 @@ class RouterFacadeTest {
|
||||
)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter,
|
||||
eventStore = mockEventStore(),
|
||||
@@ -345,7 +345,7 @@ class RouterFacadeTest {
|
||||
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter,
|
||||
eventStore = mockEventStore(),
|
||||
@@ -373,7 +373,7 @@ class RouterFacadeTest {
|
||||
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter,
|
||||
eventStore = mockEventStore(),
|
||||
@@ -394,7 +394,7 @@ class RouterFacadeTest {
|
||||
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
|
||||
val pack = ContextPack(
|
||||
id = ContextPackId("test-pack"),
|
||||
sessionId = state.sessionId ?: SessionId("unknown"),
|
||||
@@ -449,7 +449,7 @@ class RouterFacadeTest {
|
||||
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = object : InferenceRouter {
|
||||
override suspend fun route(
|
||||
@@ -497,7 +497,7 @@ class RouterFacadeTest {
|
||||
)
|
||||
},
|
||||
routerContextBuilder = object : RouterContextBuilder {
|
||||
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||
},
|
||||
inferenceRouter = mockInferenceRouter("inference response"),
|
||||
eventStore = eventStore,
|
||||
|
||||
Reference in New Issue
Block a user