feat: L3 memory retrieval on the router read path (record-only)

When a user sends input, embed it, query L3 across all sessions, dedup
against in-session turns, and record the retrieval as an event so replay
is deterministic (invariant #9). Hits land in RouterState; context
injection follows in a later slice.

- Add L3MemoryRetrievedEvent + L3RetrievedHit (registered in eventModule)
- RouterState.lastRetrievedMemory + reducer case; RouterTurn carries turnId
- RouterConfig.retrievalK (default 5)
- Harden L3 write path: runCatching + visible logging, cancellation
  re-thrown; event append stays ahead of the L3 write
- Warn prominently when the non-durable in_memory L3 backend is selected
This commit is contained in:
2026-05-30 14:36:30 +04:00
parent e8372e0643
commit e6239515bd
10 changed files with 340 additions and 37 deletions
@@ -19,3 +19,21 @@ enum class ChatTurnRole {
USER,
ROUTER,
}
@Serializable
@SerialName("L3MemoryRetrieved")
data class L3MemoryRetrievedEvent(
val sessionId: SessionId,
val queryTurnId: String,
val hits: List<L3RetrievedHit>,
val timestampMs: Long,
) : EventPayload
@Serializable
data class L3RetrievedHit(
val entryId: String,
val sourceSessionId: SessionId,
val sourceTurnId: String,
val text: String,
val score: Float,
)
@@ -10,6 +10,7 @@ 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.EventPayload
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceFailedEvent
import com.correx.core.events.events.InferenceStartedEvent
@@ -73,6 +74,7 @@ val eventModule = SerializersModule {
subclass(RiskAssessedEvent::class)
subclass(ChatSessionStartedEvent::class)
subclass(ChatTurnEvent::class)
subclass(L3MemoryRetrievedEvent::class)
}
}
+1
View File
@@ -10,6 +10,7 @@ dependencies {
implementation project(':core:context')
implementation project(':core:inference')
implementation project(':core:sessions')
implementation "org.slf4j:slf4j-api:2.0.16"
}
tasks.named("koverVerify").configure {
@@ -3,6 +3,8 @@ 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.EventMetadata
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.L3RetrievedHit
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.SteeringNoteAddedEvent
import com.correx.core.events.stores.EventStore
@@ -17,10 +19,15 @@ import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ResponseFormat
import com.correx.core.router.l3.L3MemoryEntry
import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query
import com.correx.core.router.model.RouterConfig
import com.correx.core.router.model.RouterResponse
import kotlinx.coroutines.CancellationException
import kotlinx.datetime.Clock
import java.util.*
import org.slf4j.LoggerFactory
import java.util.UUID
private val log = LoggerFactory.getLogger(DefaultRouterFacade::class.java)
interface RouterFacade {
suspend fun onUserInput(
@@ -46,13 +53,57 @@ class DefaultRouterFacade(
input: String,
mode: ChatMode,
): RouterResponse {
// Emit USER turn event
emitChatTurn(sessionId, input, ChatTurnRole.USER)
// Emit USER turn event and get the turnId for retrieval dedup
val userTurnId = emitChatTurn(sessionId, input, ChatTurnRole.USER)
// Rebuild state with user turn appended
val stateWithUserTurn = routerRepository.getRouterState(sessionId)
val effectiveStageId = stateWithUserTurn.currentStageId ?: StageId.NONE
// L3 retrieval — non-fatal; does NOT alter what contextBuilder receives in this slice
val retrieved = runCatching {
val queryVector = embedder.embed(input)
l3MemoryStore.query(L3Query(vector = queryVector, k = config.retrievalK))
}.also { result ->
result.exceptionOrNull()?.let { e ->
if (e is CancellationException) throw e
log.warn("L3 retrieval failed for session {}: {}", sessionId, e.message)
}
}.getOrElse { emptyList() }
val inSessionTurnIds = stateWithUserTurn.conversationHistory.map { it.turnId }.toSet()
val deduped = retrieved.filter { it.entry.turnId !in inSessionTurnIds }
if (deduped.isNotEmpty()) {
val nowMs = Clock.System.now().toEpochMilliseconds()
eventStore.append(
NewEvent(
metadata = EventMetadata(
eventId = EventId(UUID.randomUUID().toString()),
sessionId = sessionId,
timestamp = Clock.System.now(),
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 = nowMs,
),
),
)
}
val contextPack = routerContextBuilder.build(stateWithUserTurn, config.tokenBudget)
val provider = inferenceRouter.route(effectiveStageId, setOf(ModelCapability.General))
val inferenceRequest = InferenceRequest(
@@ -79,7 +130,7 @@ class DefaultRouterFacade(
return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING))
}
private suspend fun emitChatTurn(sessionId: SessionId, content: String, role: ChatTurnRole) {
private suspend fun emitChatTurn(sessionId: SessionId, content: String, role: ChatTurnRole): String {
val turnId = UUID.randomUUID().toString()
val nowMs = Clock.System.now().toEpochMilliseconds()
eventStore.append(
@@ -102,17 +153,26 @@ class DefaultRouterFacade(
),
)
val vector = embedder.embed(content)
l3MemoryStore.store(
L3MemoryEntry(
id = turnId,
sessionId = sessionId,
turnId = turnId,
text = content,
vector = vector,
timestampMs = nowMs,
runCatching {
val vector = embedder.embed(content)
l3MemoryStore.store(
L3MemoryEntry(
id = turnId,
sessionId = sessionId,
turnId = turnId,
text = content,
vector = vector,
timestampMs = nowMs,
)
)
)
}.also { result ->
result.exceptionOrNull()?.let { e ->
if (e is CancellationException) throw e
log.warn("L3 write failed for turn {}: {}", turnId, e.message)
}
}
return turnId
}
private suspend fun emitSteeringNote(sessionId: SessionId, content: String, effectiveStageId: StageId) {
@@ -1,6 +1,7 @@
package com.correx.core.router
import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.StageCompletedEvent
@@ -42,6 +43,7 @@ class DefaultRouterReducer : RouterReducer {
is StageCompletedEvent -> handleStageCompleted(state, event)
is StageFailedEvent -> handleStageFailed(state, event)
is ChatTurnEvent -> handleChatTurn(state, event)
is L3MemoryRetrievedEvent -> state.copy(lastRetrievedMemory = payload.hits)
else -> state
}
}
@@ -118,6 +120,7 @@ class DefaultRouterReducer : RouterReducer {
role = turnRole,
content = payload.content,
timestamp = Instant.fromEpochMilliseconds(payload.timestampMs),
turnId = payload.turnId,
)
return state.copy(
conversationHistory = state.conversationHistory + turn
@@ -7,6 +7,7 @@ import kotlinx.serialization.Serializable
@Serializable
data class RouterConfig(
val conversationKeepLast: Int = 6,
val retrievalK: Int = 5,
val tokenBudget: TokenBudget = TokenBudget(limit = 4096),
val generationConfig: GenerationConfig = GenerationConfig(
temperature = 0.7,
@@ -1,5 +1,6 @@
package com.correx.core.router.model
import com.correx.core.events.events.L3RetrievedHit
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import kotlinx.datetime.Instant
@@ -40,6 +41,7 @@ data class RouterTurn(
val role: TurnRole,
val content: String,
val timestamp: Instant,
val turnId: String = "",
)
@Serializable
@@ -49,4 +51,5 @@ data class RouterState(
val currentStageId: StageId? = null,
val l2Memory: List<RouterL2Entry> = emptyList(),
val conversationHistory: List<RouterTurn> = emptyList(),
val lastRetrievedMemory: List<L3RetrievedHit> = emptyList(),
)
+1
View File
@@ -27,6 +27,7 @@ dependencies {
implementation "io.ktor:ktor-client-cio:$ktor_version"
implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
implementation "org.slf4j:slf4j-api:2.0.16"
}
tasks.named("koverVerify").configure { enabled = false }
@@ -51,12 +51,15 @@ import com.correx.infrastructure.workflow.PromptLoader
import com.correx.infrastructure.workflow.TomlWorkflowLoader
import com.correx.infrastructure.workflow.WorkflowLoader
import kotlinx.coroutines.runBlocking
import org.slf4j.LoggerFactory
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths
import java.nio.file.StandardCopyOption
import java.sql.DriverManager
private val log = LoggerFactory.getLogger(InfrastructureModule::class.java)
@Suppress("TooManyFunctions")
object InfrastructureModule {
private val defaultDbPath: String =
@@ -199,7 +202,14 @@ object InfrastructureModule {
fun createL3MemoryStoreFromConfig(config: L3Config): L3MemoryStore {
return when (config.backend) {
"in_memory" -> createInMemoryL3MemoryStore()
"in_memory" -> {
log.warn(
"L3 backend is 'in_memory': non-durable, intended for tests/dev only — " +
"all router memory is lost on restart. " +
"Set [l3] backend = \"turbovec\" for persistence."
)
createInMemoryL3MemoryStore()
}
"turbovec" -> createTurboVecL3MemoryStore(config)
else -> throw IllegalArgumentException(
"Unknown L3 backend: '${config.backend}'. Supported: in_memory, turbovec"
@@ -1,5 +1,10 @@
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.ChatTurnRole
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.L3RetrievedHit
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
@@ -29,6 +34,7 @@ import com.correx.core.router.RouterFacade
import com.correx.core.router.RouterProjector
import com.correx.core.router.RouterRepository
import com.correx.core.router.l3.InMemoryL3MemoryStore
import com.correx.core.router.l3.L3MemoryEntry
import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query
import com.correx.core.router.model.RouterConfig
@@ -37,6 +43,7 @@ import com.correx.core.router.model.RouterState
import com.correx.core.router.model.TurnRole
import com.correx.core.router.model.WorkflowStatus
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.testing.fixtures.EventFixtures
import com.correx.testing.fixtures.inference.MockTokenizer
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking
@@ -73,14 +80,15 @@ class RouterFacadeTest {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.CHAT)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertEquals(2, mockStore.appendedEvents.size)
val userEvent = mockStore.appendedEvents[0].payload
val routerEvent = mockStore.appendedEvents[1].payload
assertTrue(userEvent is com.correx.core.events.events.ChatTurnEvent)
assertTrue(routerEvent is com.correx.core.events.events.ChatTurnEvent)
assertEquals("Hello!", (userEvent as com.correx.core.events.events.ChatTurnEvent).content)
val chatEvents = mockStore.appendedEvents
.map { it.payload }
.filterIsInstance<com.correx.core.events.events.ChatTurnEvent>()
assertEquals(2, chatEvents.size)
val userEvent = chatEvents[0]
val routerEvent = chatEvents[1]
assertEquals("Hello!", userEvent.content)
assertEquals(com.correx.core.events.events.ChatTurnRole.USER, userEvent.role)
assertEquals("inference response", (routerEvent as com.correx.core.events.events.ChatTurnEvent).content)
assertEquals("inference response", routerEvent.content)
assertEquals(com.correx.core.events.events.ChatTurnRole.ROUTER, routerEvent.role)
}
@@ -101,16 +109,16 @@ class RouterFacadeTest {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.STEERING)
facade.onUserInput(sessionId = SessionId("session-xyz"), input = "steer this way")
assertEquals(3, mockStore.appendedEvents.size)
val userEvent = mockStore.appendedEvents[0].payload
val routerEvent = mockStore.appendedEvents[1].payload
val steeringEvent = mockStore.appendedEvents[2].payload
assertTrue(userEvent is com.correx.core.events.events.ChatTurnEvent)
assertTrue(routerEvent is com.correx.core.events.events.ChatTurnEvent)
assertTrue(steeringEvent is com.correx.core.events.events.SteeringNoteAddedEvent)
assertEquals("steer this way", (userEvent as com.correx.core.events.events.ChatTurnEvent).content)
assertEquals("inference response", (routerEvent as com.correx.core.events.events.ChatTurnEvent).content)
assertEquals("inference response", (steeringEvent as com.correx.core.events.events.SteeringNoteAddedEvent).content)
val payloads = mockStore.appendedEvents.map { it.payload }
val chatEvents = payloads.filterIsInstance<com.correx.core.events.events.ChatTurnEvent>()
val steeringEvents = payloads.filterIsInstance<com.correx.core.events.events.SteeringNoteAddedEvent>()
assertEquals(2, chatEvents.size)
assertEquals(1, steeringEvents.size)
assertEquals("steer this way", chatEvents[0].content)
assertEquals(com.correx.core.events.events.ChatTurnRole.USER, chatEvents[0].role)
assertEquals("inference response", chatEvents[1].content)
assertEquals(com.correx.core.events.events.ChatTurnRole.ROUTER, chatEvents[1].role)
assertEquals("inference response", steeringEvents[0].content)
}
@Test
@@ -136,10 +144,10 @@ class RouterFacadeTest {
l3MemoryStore = InMemoryL3MemoryStore(),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!", mode = ChatMode.STEERING)
assertEquals(3, mockStore.appendedEvents.size)
val steeringEvent = mockStore.appendedEvents[2].payload
assertTrue(steeringEvent is com.correx.core.events.events.SteeringNoteAddedEvent)
assertEquals("steering response", (steeringEvent as com.correx.core.events.events.SteeringNoteAddedEvent).content)
val payloads = mockStore.appendedEvents.map { it.payload }
val steeringEvents = payloads.filterIsInstance<com.correx.core.events.events.SteeringNoteAddedEvent>()
assertEquals(1, steeringEvents.size)
assertEquals("steering response", steeringEvents[0].content)
}
// --------------------------------------------------------------------------
@@ -503,7 +511,9 @@ class RouterFacadeTest {
assertNotNull(response)
assertEquals("inference response", response.content)
assertTrue(response.steeringEmitted)
assertEquals(3, mockStore.appendedEvents.size)
val payloads = mockStore.appendedEvents.map { it.payload }
assertEquals(2, payloads.filterIsInstance<com.correx.core.events.events.ChatTurnEvent>().size)
assertEquals(1, payloads.filterIsInstance<com.correx.core.events.events.SteeringNoteAddedEvent>().size)
}
@Test
@@ -643,6 +653,200 @@ class RouterFacadeTest {
override fun capabilities(): Set<CapabilityScore> = emptySet()
}
// --------------------------------------------------------------------------
// L3 retrieval slice 2a
// --------------------------------------------------------------------------
@Test
fun `retrieval emits L3MemoryRetrievedEvent with cross-session hit`(): Unit = runBlocking {
val dimension = 8
val knownVector = FloatArray(dimension) { if (it == 0) 1f else 0f }
val stubEmbedder = object : Embedder {
override val dimension: Int = dimension
override suspend fun embed(text: String): FloatArray = knownVector.copyOf()
}
val l3Store = InMemoryL3MemoryStore()
val otherSessionId = SessionId("other-session")
val crossSessionEntryId = "cross-entry-id"
val crossSessionTurnId = "cross-turn-id"
l3Store.store(
L3MemoryEntry(
id = crossSessionEntryId,
sessionId = otherSessionId,
turnId = crossSessionTurnId,
text = "memory from another session",
vector = knownVector.copyOf(),
timestampMs = 1000L,
)
)
val eventStore = mockEventStore()
val replayer = DefaultEventReplayer<RouterState>(
store = eventStore,
projection = RouterProjector(DefaultRouterReducer()),
)
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
replayer.rebuild(sessionId)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("router reply"),
eventStore = eventStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000), retrievalK = 5),
embedder = stubEmbedder,
l3MemoryStore = l3Store,
)
facade.onUserInput(sessionId = SessionId("current-session"), input = "hello")
val l3Events = eventStore.appendedEvents.filter { it.payload is L3MemoryRetrievedEvent }
assertEquals(1, l3Events.size)
val retrieved = l3Events[0].payload as L3MemoryRetrievedEvent
assertEquals(1, retrieved.hits.size)
assertEquals(crossSessionEntryId, retrieved.hits[0].entryId)
assertEquals(otherSessionId, retrieved.hits[0].sourceSessionId)
assertEquals(crossSessionTurnId, retrieved.hits[0].sourceTurnId)
assertEquals("memory from another session", retrieved.hits[0].text)
}
@Test
fun `dedup filters hit whose turnId is already in current session history`(): Unit = runBlocking {
val dimension = 8
val knownVector = FloatArray(dimension) { if (it == 0) 1f else 0f }
val stubEmbedder = object : Embedder {
override val dimension: Int = dimension
override suspend fun embed(text: String): FloatArray = knownVector.copyOf()
}
val l3Store = InMemoryL3MemoryStore()
val currentSessionId = SessionId("dedup-session")
val inSessionTurnId = "in-session-turn-id"
// Pre-populate L3 with an entry whose turnId matches an in-session turn
l3Store.store(
L3MemoryEntry(
id = inSessionTurnId,
sessionId = currentSessionId,
turnId = inSessionTurnId,
text = "already in session history",
vector = knownVector.copyOf(),
timestampMs = 500L,
)
)
val eventStore = mockEventStore()
val replayer = DefaultEventReplayer<RouterState>(
store = eventStore,
projection = RouterProjector(DefaultRouterReducer()),
)
// Bootstrap session with a ChatTurnEvent whose turnId matches the L3 entry
eventStore.append(
NewEvent(
metadata = EventMetadata(
eventId = EventId("seed-event-id"),
sessionId = currentSessionId,
timestamp = kotlinx.datetime.Clock.System.now(),
schemaVersion = 1,
causationId = null,
correlationId = null,
),
payload = ChatTurnEvent(
sessionId = currentSessionId,
turnId = inSessionTurnId,
role = ChatTurnRole.USER,
content = "already in session history",
timestampMs = 500L,
),
)
)
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
replayer.rebuild(sessionId)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("router reply"),
eventStore = eventStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000), retrievalK = 5),
embedder = stubEmbedder,
l3MemoryStore = l3Store,
)
facade.onUserInput(sessionId = currentSessionId, input = "a new message")
val l3Events = eventStore.appendedEvents.filter { it.payload is L3MemoryRetrievedEvent }
// The only hit was in-session, so deduped to empty → no L3MemoryRetrievedEvent
assertTrue(l3Events.isEmpty(), "Expected no L3RetrievedEvent when all hits are deduplicated")
}
@Test
fun `retrieval failure is non-fatal - CHAT turn still completes`(): Unit = runBlocking {
val throwingEmbedder = object : Embedder {
override val dimension: Int = 8
override suspend fun embed(text: String): FloatArray =
throw RuntimeException("embedding service unavailable")
}
val eventStore = mockEventStore()
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("response"),
eventStore = eventStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
embedder = throwingEmbedder,
l3MemoryStore = InMemoryL3MemoryStore(),
)
// Should not throw despite embedder failure
val response = facade.onUserInput(sessionId = SessionId("fail-session"), input = "test")
assertEquals("response", response.content)
val payloads = eventStore.appendedEvents.map { it.payload }
val chatEvents = payloads.filterIsInstance<com.correx.core.events.events.ChatTurnEvent>()
assertEquals(2, chatEvents.size, "USER and ROUTER ChatTurnEvents must be present")
val l3Events = payloads.filterIsInstance<L3MemoryRetrievedEvent>()
assertTrue(l3Events.isEmpty(), "No L3RetrievedEvent expected on retrieval failure")
}
@Test
fun `replay of L3MemoryRetrievedEvent populates lastRetrievedMemory in RouterState`(): Unit = runBlocking {
val reducer = DefaultRouterReducer()
val projector = RouterProjector(reducer)
val sessionId = SessionId("replay-session")
val hits = listOf(
L3RetrievedHit(
entryId = "entry-1",
sourceSessionId = SessionId("source-session"),
sourceTurnId = "source-turn-1",
text = "remembered text",
score = 0.95f,
)
)
val event = EventFixtures.stored(
payload = L3MemoryRetrievedEvent(
sessionId = sessionId,
queryTurnId = "query-turn-id",
hits = hits,
timestampMs = 9999L,
)
)
val state = projector.apply(projector.initial(), event)
assertEquals(1, state.lastRetrievedMemory.size)
assertEquals("entry-1", state.lastRetrievedMemory[0].entryId)
assertEquals("remembered text", state.lastRetrievedMemory[0].text)
assertEquals(0.95f, state.lastRetrievedMemory[0].score)
}
private fun emptyContextPack(): ContextPack = ContextPack(
id = ContextPackId("empty"),
sessionId = SessionId("unknown"),