feat(router): capture latency + tokens on ROUTER ChatTurnEvent

ChatTurnEvent gains nullable latencyMs and tokensUsed fields (defaults preserve
backward-compat; legacy JSON without them deserializes cleanly). emitChatTurn
accepts optional metrics; the ROUTER emit site passes inferenceResponse values
through; the USER emit site leaves both null.
This commit is contained in:
2026-06-03 15:11:37 +04:00
parent e9a87febc4
commit 8d8b2914e0
3 changed files with 89 additions and 2 deletions
@@ -1,6 +1,7 @@
package com.correx.core.events.events package com.correx.core.events.events
import com.correx.core.events.types.SessionId import com.correx.core.events.types.SessionId
import com.correx.core.inference.TokenUsage
import kotlinx.serialization.SerialName import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable import kotlinx.serialization.Serializable
@@ -12,6 +13,8 @@ data class ChatTurnEvent(
val role: ChatTurnRole, val role: ChatTurnRole,
val content: String, val content: String,
val timestampMs: Long, val timestampMs: Long,
val latencyMs: Long? = null,
val tokensUsed: TokenUsage? = null,
) : EventPayload ) : EventPayload
@Serializable @Serializable
@@ -18,6 +18,7 @@ import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceRouter import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ResponseFormat import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.TokenUsage
import com.correx.core.router.l3.L3MemoryEntry import com.correx.core.router.l3.L3MemoryEntry
import com.correx.core.router.l3.L3MemoryStore import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query import com.correx.core.router.l3.L3Query
@@ -148,7 +149,11 @@ class DefaultRouterFacade(
val content = inferenceResponse.text val content = inferenceResponse.text
// Emit ROUTER turn event // Emit ROUTER turn event
emitChatTurn(sessionId, content, ChatTurnRole.ROUTER) emitChatTurn(
sessionId, content, ChatTurnRole.ROUTER,
latencyMs = inferenceResponse.latencyMs,
tokensUsed = inferenceResponse.tokensUsed,
)
if (mode == ChatMode.STEERING) { if (mode == ChatMode.STEERING) {
val validationError = validateSteering?.invoke(content) val validationError = validateSteering?.invoke(content)
@@ -160,7 +165,13 @@ class DefaultRouterFacade(
return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING)) return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING))
} }
private suspend fun emitChatTurn(sessionId: SessionId, content: String, role: ChatTurnRole): String { private suspend fun emitChatTurn(
sessionId: SessionId,
content: String,
role: ChatTurnRole,
latencyMs: Long? = null,
tokensUsed: TokenUsage? = null,
): String {
val turnId = UUID.randomUUID().toString() val turnId = UUID.randomUUID().toString()
val now = Clock.System.now() val now = Clock.System.now()
eventStore.append( eventStore.append(
@@ -179,6 +190,8 @@ class DefaultRouterFacade(
role = role, role = role,
content = content, content = content,
timestampMs = now.toEpochMilliseconds(), timestampMs = now.toEpochMilliseconds(),
latencyMs = latencyMs,
tokensUsed = tokensUsed,
), ),
), ),
) )
@@ -52,9 +52,11 @@ import com.correx.testing.fixtures.inference.MockTokenizer
import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Instant import kotlinx.datetime.Instant
import kotlinx.serialization.decodeFromString
import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertNotNull import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test import org.junit.jupiter.api.Test
@@ -544,6 +546,75 @@ class RouterFacadeTest {
assertEquals(2, hits.size) assertEquals(2, hits.size)
} }
// --------------------------------------------------------------------------
// Task 2.1 — latency + token metrics on ROUTER ChatTurnEvent
// --------------------------------------------------------------------------
@Test
fun `ROUTER ChatTurnEvent carries latencyMs and tokensUsed from inference and USER turn has both null`(): Unit = runBlocking {
val capturedRequests = mutableListOf<InferenceRequest>()
val knownLatencyMs = 42L
val knownTokenUsage = TokenUsage(promptTokens = 10, completionTokens = 20)
val metricsProvider = object : InferenceProvider {
override val id = ProviderId("mock")
override val name = "Mock"
override val tokenizer = MockTokenizer()
override suspend fun infer(request: InferenceRequest): InferenceResponse {
capturedRequests.add(request)
return InferenceResponse(
requestId = request.requestId,
text = "metrics response",
finishReason = FinishReason.Stop,
tokensUsed = knownTokenUsage,
latencyMs = knownLatencyMs,
)
}
override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> = setOf(CapabilityScore(ModelCapability.General, 1.0))
}
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 = object : InferenceRouter {
override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider =
metricsProvider
},
eventStore = eventStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
embedder = NoopEmbedder(dimension = 8),
l3MemoryStore = InMemoryL3MemoryStore(),
)
facade.onUserInput(sessionId = SessionId("metrics-session"), input = "hello metrics")
val chatEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<ChatTurnEvent>()
assertEquals(2, chatEvents.size)
val userEvent = chatEvents.first { it.role == ChatTurnRole.USER }
val routerEvent = chatEvents.first { it.role == ChatTurnRole.ROUTER }
assertNull(userEvent.latencyMs, "USER ChatTurnEvent must have null latencyMs")
assertNull(userEvent.tokensUsed, "USER ChatTurnEvent must have null tokensUsed")
assertEquals(knownLatencyMs, routerEvent.latencyMs)
assertNotNull(routerEvent.tokensUsed)
assertEquals(30, routerEvent.tokensUsed!!.totalTokens)
}
@Test
fun `ChatTurnEvent backward-compat - legacy JSON without latencyMs and tokensUsed deserializes with both null`() {
val legacyJson = """{"sessionId":"s1","turnId":"t1","role":"USER","content":"hello","timestampMs":1000}"""
val event = com.correx.core.events.serialization.eventJson.decodeFromString<ChatTurnEvent>(legacyJson)
assertEquals(SessionId("s1"), event.sessionId)
assertEquals("hello", event.content)
assertNull(event.latencyMs)
assertNull(event.tokensUsed)
}
// -------------------------------------------------------------------------- // --------------------------------------------------------------------------
// Helpers // Helpers
// -------------------------------------------------------------------------- // --------------------------------------------------------------------------