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
@@ -52,9 +52,11 @@ import com.correx.testing.fixtures.inference.MockTokenizer
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Instant
import kotlinx.serialization.decodeFromString
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
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.Test
@@ -544,6 +546,75 @@ class RouterFacadeTest {
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
// --------------------------------------------------------------------------