diff --git a/core/events/src/main/kotlin/com/correx/core/events/events/RouterEvents.kt b/core/events/src/main/kotlin/com/correx/core/events/events/RouterEvents.kt index 8f51a5b0..a2f79567 100644 --- a/core/events/src/main/kotlin/com/correx/core/events/events/RouterEvents.kt +++ b/core/events/src/main/kotlin/com/correx/core/events/events/RouterEvents.kt @@ -1,6 +1,7 @@ package com.correx.core.events.events import com.correx.core.events.types.SessionId +import com.correx.core.inference.TokenUsage import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable @@ -12,6 +13,8 @@ data class ChatTurnEvent( val role: ChatTurnRole, val content: String, val timestampMs: Long, + val latencyMs: Long? = null, + val tokensUsed: TokenUsage? = null, ) : EventPayload @Serializable diff --git a/core/router/src/main/kotlin/com/correx/core/router/RouterFacade.kt b/core/router/src/main/kotlin/com/correx/core/router/RouterFacade.kt index d7f2f1a0..6889aa4b 100644 --- a/core/router/src/main/kotlin/com/correx/core/router/RouterFacade.kt +++ b/core/router/src/main/kotlin/com/correx/core/router/RouterFacade.kt @@ -18,6 +18,7 @@ import com.correx.core.inference.InferenceRequest import com.correx.core.inference.InferenceRouter import com.correx.core.inference.ModelCapability 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.L3MemoryStore import com.correx.core.router.l3.L3Query @@ -148,7 +149,11 @@ class DefaultRouterFacade( val content = inferenceResponse.text // Emit ROUTER turn event - emitChatTurn(sessionId, content, ChatTurnRole.ROUTER) + emitChatTurn( + sessionId, content, ChatTurnRole.ROUTER, + latencyMs = inferenceResponse.latencyMs, + tokensUsed = inferenceResponse.tokensUsed, + ) if (mode == ChatMode.STEERING) { val validationError = validateSteering?.invoke(content) @@ -160,7 +165,13 @@ class DefaultRouterFacade( 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 now = Clock.System.now() eventStore.append( @@ -179,6 +190,8 @@ class DefaultRouterFacade( role = role, content = content, timestampMs = now.toEpochMilliseconds(), + latencyMs = latencyMs, + tokensUsed = tokensUsed, ), ), ) diff --git a/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt b/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt index c3603aef..fdc6bbcc 100644 --- a/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt +++ b/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt @@ -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() + 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 = 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): 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() + 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(legacyJson) + assertEquals(SessionId("s1"), event.sessionId) + assertEquals("hello", event.content) + assertNull(event.latencyMs) + assertNull(event.tokensUsed) + } + // -------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------