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:
2026-05-28 14:06:58 +04:00
parent 92bea6c2c4
commit e05532e7b2
11 changed files with 316 additions and 133 deletions
@@ -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,