feat: inject recalled L3 memory into router context with budget

Record L3 retrieval as an event carrying hit text (invariant #9), then rebuild router state and inject recalled memories as a SYSTEM L3 layer in the context pack. Apply token budget: protected frames (L0 immutable + current user turn) are never dropped; honest budgetUsed is reported and 'BudgetExceeded' is flagged in appliedStrategies when they overflow. L1/L2 fit newest-first so oldest entries drop first. Emit ContextTruncatedEvent when entries are dropped. L3MemoryRetrievedEvent is emitted on every CHAT turn (empty hits reset recalled memory).
2026-05-30 18:40:08 +04:00
parent e6239515bd
commit 8e6a3e1470
6 changed files with 852 additions and 77 deletions
@@ -37,3 +37,13 @@ data class L3RetrievedHit(
    val text: String,
    val score: Float,
 )
+
+@Serializable
+@SerialName("ContextTruncated")
+data class ContextTruncatedEvent(
+    val sessionId: SessionId,
+    val turnId: String,
+    val entriesDropped: Int,
+    val truncatedLayers: List<String>,
+    val timestampMs: Long,
+) : EventPayload
@@ -9,6 +9,7 @@ import com.correx.core.events.events.ArtifactValidatedEvent
 import com.correx.core.events.events.ArtifactValidatingEvent
 import com.correx.core.events.events.ChatSessionStartedEvent
 import com.correx.core.events.events.ChatTurnEvent
+import com.correx.core.events.events.ContextTruncatedEvent
 import com.correx.core.events.events.EventPayload
 import com.correx.core.events.events.L3MemoryRetrievedEvent
 import com.correx.core.events.events.InferenceCompletedEvent
@@ -75,6 +76,7 @@ val eventModule = SerializersModule {
        subclass(ChatSessionStartedEvent::class)
        subclass(ChatTurnEvent::class)
        subclass(L3MemoryRetrievedEvent::class)
+        subclass(ContextTruncatedEvent::class)
    }
 }

@@ -17,6 +17,7 @@ import com.correx.core.router.model.RouterState
 import com.correx.core.router.model.RouterTurn
 import com.correx.core.router.model.TurnRole
 import java.util.*
+import kotlinx.coroutines.CancellationException

 interface RouterContextBuilder {
    suspend fun build(state: RouterState, budget: TokenBudget): ContextPack
@@ -30,56 +31,128 @@ class DefaultRouterContextBuilder(
    companion object {
        private const val SYSTEM_PROMPT =
            "You are a routing assistant. Provide guidance based on workflow state and conversation context."
+        private const val RECALLED_MEMORY_PREFIX = "[recalled memory]"
    }

    override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
-        var remainingBudget = budget.limit
-        val allEntries = mutableListOf<ContextEntry>()
-        var droppedCount = 0
-
+        // Protected frame: system prompt and workflow status are ALWAYS included regardless of budget.
        val systemPrompt = buildContextEntry(
            sourceType = "systemPrompt",
            sourceId = "router-system",
            content = SYSTEM_PROMPT,
+            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
-        remainingBudget -= systemPrompt.tokenEstimate
-        if (remainingBudget < 0) remainingBudget = 0
-        allEntries += systemPrompt
-
        val workflowStatusEntry = buildContextEntry(
            sourceType = "workflowStatus",
            sourceId = state.currentStageId?.value ?: "none",
            content = buildWorkflowStatusContent(state),
+            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
-        remainingBudget -= workflowStatusEntry.tokenEstimate
-        if (remainingBudget < 0) remainingBudget = 0
+
+        // Compute remaining budget after protected frame.
+        val protectedTokens = systemPrompt.tokenEstimate + workflowStatusEntry.tokenEstimate
+        var remainingBudget = (budget.limit - protectedTokens).coerceAtLeast(0)
+
+        val allEntries = mutableListOf<ContextEntry>()
+        allEntries += systemPrompt
        allEntries += workflowStatusEntry

-        val recentTurns = state.conversationHistory.takeLast(config.conversationKeepLast)
-        for (turn in recentTurns) {
+        var droppedCount = 0
+        val truncatedLayers = mutableSetOf<ContextLayer>()
+
+        // --- Conversation turns (L1) ---
+        // The current (protected) user turn is the last element in conversationHistory,
+        // regardless of conversationKeepLast — it must always be present.
+        val protectedUserTurn: RouterTurn? = state.conversationHistory.lastOrNull()
+
+        // Apply conversationKeepLast cap. The protected turn may or may not fall within this window.
+        val cappedTurns = state.conversationHistory.takeLast(config.conversationKeepLast)
+
+        // Build entries for the capped window, excluding the protected turn (to avoid double-counting).
+        val nonProtectedCapped = if (protectedUserTurn != null && cappedTurns.lastOrNull() == protectedUserTurn) {
+            cappedTurns.dropLast(1)
+        } else {
+            cappedTurns
+        }
+
+        // Build the protected user turn entry.
+        val protectedTurnEntry: ContextEntry? = protectedUserTurn?.let { turn ->
            val role = when (turn.role) {
                TurnRole.USER -> EntryRole.USER
                TurnRole.ROUTER -> EntryRole.ASSISTANT
            }
-            val entry = buildContextEntry(
+            buildContextEntry(
                sourceType = "conversation",
                sourceId = "${turn.role.name}-${turn.hashCode()}",
                content = turn.content,
                layer = ContextLayer.L1,
                role = role,
            )
+        }
+
+        // Reserve budget for the protected user turn first.
+        if (protectedTurnEntry != null) {
+            remainingBudget = (remainingBudget - protectedTurnEntry.tokenEstimate).coerceAtLeast(0)
+        }
+
+        // Fit remaining capped turns newest-to-oldest (excluding protected), so oldest get dropped when tight.
+        val nonProtectedTurnEntries = nonProtectedCapped.map { turn ->
+            val role = when (turn.role) {
+                TurnRole.USER -> EntryRole.USER
+                TurnRole.ROUTER -> EntryRole.ASSISTANT
+            }
+            buildContextEntry(
+                sourceType = "conversation",
+                sourceId = "${turn.role.name}-${turn.hashCode()}",
+                content = turn.content,
+                layer = ContextLayer.L1,
+                role = role,
+            )
+        }
+
+        val fittedTurns = mutableListOf<ContextEntry>()
+        for (entry in nonProtectedTurnEntries.asReversed()) {
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
-                if (remainingBudget < 0) remainingBudget = 0
-                allEntries += entry
+                fittedTurns.add(entry)
            } else {
                droppedCount++
+                truncatedLayers.add(ContextLayer.L1)
            }
        }

-        for (l2Entry in state.l2Memory) {
+        // Re-order fitted turns oldest-to-newest (reverse back) and append protected turn at end.
+        fittedTurns.reverse()
+        allEntries.addAll(fittedTurns)
+        protectedTurnEntry?.let { allEntries.add(it) }
+
+        // --- L3 memory (recalled cross-session) ---
+        // Sort by score descending (highest relevance first), add what fits.
+        val sortedHits = state.lastRetrievedMemory.sortedByDescending { it.score }
+        for (hit in sortedHits) {
+            val content = "$RECALLED_MEMORY_PREFIX ${hit.text}"
+            val entry = buildContextEntry(
+                sourceType = "recalledMemory",
+                sourceId = hit.entryId,
+                content = content,
+                layer = ContextLayer.L3,
+                role = EntryRole.SYSTEM,
+            )
+            if (remainingBudget >= entry.tokenEstimate) {
+                remainingBudget -= entry.tokenEstimate
+                allEntries += entry
+            } else {
+                droppedCount++
+                truncatedLayers.add(ContextLayer.L3)
+            }
+        }
+
+        // --- L2 memory (compressed session summaries) ---
+        // Iterate newest-to-oldest so that when budget is tight the OLDEST entries are dropped.
+        val fittedL2 = mutableListOf<ContextEntry>()
+        for (l2Entry in state.l2Memory.asReversed()) {
            val content = buildL2Content(l2Entry)
            val entry = buildContextEntry(
                sourceType = "stageSummary",
@@ -90,16 +163,31 @@ class DefaultRouterContextBuilder(
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
-                if (remainingBudget < 0) remainingBudget = 0
-                allEntries += entry
+                fittedL2.add(entry)
            } else {
                droppedCount++
+                truncatedLayers.add(ContextLayer.L2)
            }
        }
+        // Restore chronological (oldest-to-newest) order.
+        fittedL2.reverse()
+        allEntries.addAll(fittedL2)
+
+        // Build layers map sorted by ContextLayer ordinal for deterministic ordering.
+        val layers = allEntries
+            .groupBy { it.layer }
+            .toSortedMap(compareBy { it.ordinal })

-        val layers = allEntries.groupBy { it.layer }
        val budgetUsed = allEntries.sumOf { it.tokenEstimate }

+        // Build appliedStrategies list — deterministic, stable order.
+        val appliedStrategies = buildList {
+            add("L0Immutable")
+            add("Conversation")
+            if (state.lastRetrievedMemory.isNotEmpty()) add("L3Recall")
+            if (budgetUsed > budget.limit) add("BudgetExceeded")
+        }
+
        return ContextPack(
            id = ContextPackId("${state.sessionId?.value ?: "unknown"}-router-pack"),
            sessionId = state.sessionId ?: SessionId("unknown"),
@@ -108,8 +196,8 @@ class DefaultRouterContextBuilder(
            budgetUsed = budgetUsed,
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
-                appliedStrategies = listOf("L0Immutable", "Conversation"),
-                truncatedLayers = emptyList(),
+                appliedStrategies = appliedStrategies,
+                truncatedLayers = truncatedLayers.toList(),
                entriesDropped = droppedCount,
            ),
        )
@@ -150,7 +238,10 @@ class DefaultRouterContextBuilder(
    private suspend fun estimateTokens(content: String): Int {
        val t = tokenizer
        if (t != null) {
-            return runCatching { t.countTokens(content) }.getOrElse { fallbackEstimate(content) }
+            return runCatching { t.countTokens(content) }.getOrElse { e ->
+                if (e is CancellationException) throw e
+                fallbackEstimate(content)
+            }
        }
        return fallbackEstimate(content)
    }
@@ -158,4 +249,4 @@ class DefaultRouterContextBuilder(
    private fun fallbackEstimate(content: String): Int {
        return (content.length / 4).coerceAtLeast(1)
    }
-}
+}
@@ -2,6 +2,7 @@ package com.correx.core.router

 import com.correx.core.events.events.ChatTurnEvent
 import com.correx.core.events.events.ChatTurnRole
+import com.correx.core.events.events.ContextTruncatedEvent
 import com.correx.core.events.events.EventMetadata
 import com.correx.core.events.events.L3MemoryRetrievedEvent
 import com.correx.core.events.events.L3RetrievedHit
@@ -57,10 +58,12 @@ class DefaultRouterFacade(
        val userTurnId = emitChatTurn(sessionId, input, ChatTurnRole.USER)

        // Rebuild state with user turn appended
-        val stateWithUserTurn = routerRepository.getRouterState(sessionId)
-        val effectiveStageId = stateWithUserTurn.currentStageId ?: StageId.NONE
+        var state = routerRepository.getRouterState(sessionId)
+        val effectiveStageId = state.currentStageId ?: StageId.NONE

-        // L3 retrieval — non-fatal; does NOT alter what contextBuilder receives in this slice
+        // L3 retrieval — non-fatal; results fed back into state via event.
+        // Always emit L3MemoryRetrievedEvent for every CHAT turn where retrieval was attempted,
+        // including with empty hits — so lastRetrievedMemory is always current-turn-scoped.
        val retrieved = runCatching {
            val queryVector = embedder.embed(input)
            l3MemoryStore.query(L3Query(vector = queryVector, k = config.retrievalK))
@@ -71,40 +74,67 @@ class DefaultRouterFacade(
            }
        }.getOrElse { emptyList() }

-        val inSessionTurnIds = stateWithUserTurn.conversationHistory.map { it.turnId }.toSet()
+        val inSessionTurnIds = state.conversationHistory.map { it.turnId }.toSet()
        val deduped = retrieved.filter { it.entry.turnId !in inSessionTurnIds }

-        if (deduped.isNotEmpty()) {
-            val nowMs = Clock.System.now().toEpochMilliseconds()
+        // Always emit L3MemoryRetrievedEvent (even with empty hits) so that
+        // lastRetrievedMemory always reflects this turn and a prior turn's hits
+        // are never injected into the next turn's context.
+        val retrievalNow = Clock.System.now()
+        eventStore.append(
+            NewEvent(
+                metadata = EventMetadata(
+                    eventId = EventId(UUID.randomUUID().toString()),
+                    sessionId = sessionId,
+                    timestamp = retrievalNow,
+                    schemaVersion = 1,
+                    causationId = null,
+                    correlationId = null,
+                ),
+                payload = L3MemoryRetrievedEvent(
+                    sessionId = sessionId,
+                    queryTurnId = userTurnId,
+                    hits = deduped.map { hit ->
+                        L3RetrievedHit(
+                            entryId = hit.entry.id,
+                            sourceSessionId = hit.entry.sessionId,
+                            sourceTurnId = hit.entry.turnId,
+                            text = hit.entry.text,
+                            score = hit.score,
+                        )
+                    },
+                    timestampMs = retrievalNow.toEpochMilliseconds(),
+                ),
+            ),
+        )
+        // Rebuild state so lastRetrievedMemory reflects this turn
+        state = routerRepository.getRouterState(sessionId)
+
+        val contextPack = routerContextBuilder.build(state, config.tokenBudget)
+
+        if (contextPack.compressionMetadata.entriesDropped > 0) {
+            val truncNow = Clock.System.now()
            eventStore.append(
                NewEvent(
                    metadata = EventMetadata(
                        eventId = EventId(UUID.randomUUID().toString()),
                        sessionId = sessionId,
-                        timestamp = Clock.System.now(),
+                        timestamp = truncNow,
                        schemaVersion = 1,
                        causationId = null,
                        correlationId = null,
                    ),
-                    payload = L3MemoryRetrievedEvent(
+                    payload = ContextTruncatedEvent(
                        sessionId = sessionId,
-                        queryTurnId = userTurnId,
-                        hits = deduped.map { hit ->
-                            L3RetrievedHit(
-                                entryId = hit.entry.id,
-                                sourceSessionId = hit.entry.sessionId,
-                                sourceTurnId = hit.entry.turnId,
-                                text = hit.entry.text,
-                                score = hit.score,
-                            )
-                        },
-                        timestampMs = nowMs,
+                        turnId = userTurnId,
+                        entriesDropped = contextPack.compressionMetadata.entriesDropped,
+                        truncatedLayers = contextPack.compressionMetadata.truncatedLayers.map { it.name },
+                        timestampMs = truncNow.toEpochMilliseconds(),
                    ),
                ),
            )
        }

-        val contextPack = routerContextBuilder.build(stateWithUserTurn, config.tokenBudget)
        val provider = inferenceRouter.route(effectiveStageId, setOf(ModelCapability.General))
        val inferenceRequest = InferenceRequest(
            requestId = InferenceRequestId(UUID.randomUUID().toString()),
@@ -132,13 +162,13 @@ class DefaultRouterFacade(

    private suspend fun emitChatTurn(sessionId: SessionId, content: String, role: ChatTurnRole): String {
        val turnId = UUID.randomUUID().toString()
-        val nowMs = Clock.System.now().toEpochMilliseconds()
+        val now = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
-                    timestamp = Clock.System.now(),
+                    timestamp = now,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
@@ -148,7 +178,7 @@ class DefaultRouterFacade(
                    turnId = turnId,
                    role = role,
                    content = content,
-                    timestampMs = nowMs,
+                    timestampMs = now.toEpochMilliseconds(),
                ),
            ),
        )
@@ -162,7 +192,7 @@ class DefaultRouterFacade(
                    turnId = turnId,
                    text = content,
                    vector = vector,
-                    timestampMs = nowMs,
+                    timestampMs = now.toEpochMilliseconds(),
                )
            )
        }.also { result ->
@@ -176,12 +206,13 @@ class DefaultRouterFacade(
    }

    private suspend fun emitSteeringNote(sessionId: SessionId, content: String, effectiveStageId: StageId) {
+        val now = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
-                    timestamp = Clock.System.now(),
+                    timestamp = now,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,