Files
correx/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt
T

620 lines
28 KiB
Kotlin

import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.TokenUsage
import com.correx.core.router.ChatMode
import com.correx.core.router.DefaultRouterFacade
import com.correx.core.router.RouterContextBuilder
import com.correx.core.router.RouterFacade
import com.correx.core.router.RouterRepository
import com.correx.core.router.model.RouterConfig
import com.correx.core.router.model.RouterResponse
import com.correx.core.router.model.RouterState
import com.correx.core.router.model.TurnRole
import com.correx.core.router.model.WorkflowStatus
import com.correx.testing.fixtures.inference.MockTokenizer
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking
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.assertTrue
import org.junit.jupiter.api.Test
class RouterFacadeTest {
// --------------------------------------------------------------------------
// CHAT mode
// --------------------------------------------------------------------------
@Test
fun `CHAT mode returns inference response content`(): Unit = runBlocking {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.CHAT)
val response = facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello, world!")
assertEquals("inference response", response.content)
}
@Test
fun `CHAT mode sets steeringEmitted to false`(): Unit = runBlocking {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.CHAT)
val response = facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertFalse(response.steeringEmitted)
}
@Test
fun `CHAT mode does not append to EventStore`(): Unit = runBlocking {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.CHAT)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertTrue(mockStore.appendedEvents.isEmpty())
}
// --------------------------------------------------------------------------
// STEERING mode
// --------------------------------------------------------------------------
@Test
fun `STEERING mode sets steeringEmitted to true`(): Unit = runBlocking {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.STEERING)
val response = facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertTrue(response.steeringEmitted)
}
@Test
fun `STEERING mode does not append events to store`(): Unit = runBlocking {
// SteeringNoteAddedEvent emission is deferred to P4-1 (see P0-6).
// The steeringEmitted flag still informs the TUI that steering mode was used.
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.STEERING)
facade.onUserInput(sessionId = SessionId("session-xyz"), input = "steer this way")
assertTrue(mockStore.appendedEvents.isEmpty())
}
@Test
fun `STEERING mode does not append events regardless of state stageId`(): Unit = runBlocking {
val mockStore = mockEventStore()
val stageId = StageId("stage-A")
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = stageId,
)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("response"),
eventStore = mockStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!", mode = ChatMode.STEERING)
assertTrue(mockStore.appendedEvents.isEmpty())
}
// --------------------------------------------------------------------------
// In-memory conversation history
// --------------------------------------------------------------------------
@Test
fun `conversation history grows per call - user and router turns appended`(): Unit = runBlocking {
val capturedStates = mutableListOf<RouterState>()
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = StageId("s1"),
)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedStates.add(state)
return emptyContextPack()
}
},
inferenceRouter = mockInferenceRouter("router reply"),
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
val sessionId = SessionId("history-session")
facade.onUserInput(sessionId = sessionId, input = "first message")
facade.onUserInput(sessionId = sessionId, input = "second message")
// Second call's context builder sees two user turns + one router turn from the first call
val stateOnSecondCall = capturedStates[1]
assertEquals(3, stateOnSecondCall.conversationHistory.size)
assertEquals(TurnRole.USER, stateOnSecondCall.conversationHistory[0].role)
assertEquals("first message", stateOnSecondCall.conversationHistory[0].content)
assertEquals(TurnRole.ROUTER, stateOnSecondCall.conversationHistory[1].role)
assertEquals("router reply", stateOnSecondCall.conversationHistory[1].content)
assertEquals(TurnRole.USER, stateOnSecondCall.conversationHistory[2].role)
assertEquals("second message", stateOnSecondCall.conversationHistory[2].content)
}
@Test
fun `conversation history is session-scoped - different sessions do not share history`(): Unit = runBlocking {
val capturedStates = mutableListOf<RouterState>()
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
RouterState(sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedStates.add(state)
return emptyContextPack()
}
},
inferenceRouter = mockInferenceRouter("response"),
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("session-A"), input = "message A")
facade.onUserInput(sessionId = SessionId("session-B"), input = "message B")
val stateA = capturedStates[0]
val stateB = capturedStates[1]
assertEquals(1, stateA.conversationHistory.size)
assertEquals(1, stateB.conversationHistory.size)
assertEquals("message A", stateA.conversationHistory[0].content)
assertEquals("message B", stateB.conversationHistory[0].content)
}
// --------------------------------------------------------------------------
// Orchestration
// --------------------------------------------------------------------------
@Test
fun `state is passed through to context builder`(): Unit = runBlocking {
val capturedState = mutableListOf<RouterState>()
val mockContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedState.add(state)
return emptyContextPack()
}
}
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = StageId("s1"),
)
},
routerContextBuilder = mockContextBuilder,
inferenceRouter = mockInferenceRouter("inference response"),
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertEquals(1, capturedState.size)
assertEquals(SessionId("test-session"), capturedState[0].sessionId)
assertEquals(WorkflowStatus.RUNNING, capturedState[0].workflowStatus)
}
@Test
fun `budget is passed through to context builder`(): Unit = runBlocking {
val capturedBudget = mutableListOf<TokenBudget>()
val mockContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedBudget.add(budget)
return emptyContextPack()
}
}
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = mockContextBuilder,
inferenceRouter = mockInferenceRouter("response"),
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 4200)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertEquals(1, capturedBudget.size)
assertEquals(4200, capturedBudget[0].limit)
}
@Test
fun `stage ID uses state currentStageId`(): Unit = runBlocking {
val capturedStageId = mutableListOf<StageId>()
val mockInferenceRouter = object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
capturedStageId.add(stageId)
return mockProvider("response")
}
}
val stateStageId = StageId("state-stage")
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = stateStageId,
)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertEquals(stateStageId, capturedStageId[0])
}
@Test
fun `stage ID falls back to StageId none when state has no currentStageId`(): Unit = runBlocking {
val capturedStageId = mutableListOf<StageId>()
val mockInferenceRouter = object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
capturedStageId.add(stageId)
return mockProvider("response")
}
}
val facade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.IDLE,
currentStageId = null,
)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertEquals(StageId("none"), capturedStageId[0])
}
@Test
fun `new InferenceRequestId per call`(): Unit = runBlocking {
val capturedRequestIds = mutableListOf<InferenceRequestId>()
val mockInferenceRouter = object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
return mockProviderWithCapture("response", capturedRequestIds)
}
}
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,
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("s1"), input = "first")
facade.onUserInput(sessionId = SessionId("s1"), input = "second")
assertEquals(2, capturedRequestIds.size)
assertFalse(capturedRequestIds[0] == capturedRequestIds[1])
}
@Test
fun `GenerationConfig defaults temperature 0 7 topP 0 9 maxTokens 512`(): Unit = runBlocking {
val capturedRequests = mutableListOf<InferenceRequest>()
val mockInferenceRouter = object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
return mockProviderWithRequestCapture("response", capturedRequests)
}
}
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,
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
val req = capturedRequests[0]
assertEquals(0.7, req.generationConfig.temperature)
assertEquals(0.9, req.generationConfig.topP)
assertEquals(512, req.generationConfig.maxTokens)
}
@Test
fun `context pack is passed to inference provider`(): Unit = runBlocking {
val capturedContextPacks = mutableListOf<ContextPack>()
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 {
val pack = ContextPack(
id = ContextPackId("test-pack"),
sessionId = state.sessionId ?: SessionId("unknown"),
stageId = StageId("test"),
layers = emptyMap(),
budgetUsed = 0,
budgetLimit = budget.limit,
)
capturedContextPacks.add(pack)
return pack
}
},
inferenceRouter = object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
return object : InferenceProvider {
override val id = ProviderId("mock")
override val name = "Mock"
override val tokenizer = MockTokenizer()
override suspend fun infer(request: InferenceRequest): InferenceResponse {
assertEquals(1, capturedContextPacks.size)
assertEquals(capturedContextPacks[0].id, request.contextPack.id)
return InferenceResponse(
requestId = request.requestId,
text = "response",
finishReason = FinishReason.Stop,
tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
latencyMs = 0,
)
}
override suspend fun healthCheck(): ProviderHealth =
ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> = emptySet()
}
}
},
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
}
@Test
fun `responseFormat defaults to Text`(): Unit = runBlocking {
val capturedRequests = mutableListOf<InferenceRequest>()
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 {
return mockProviderWithRequestCapture("response", capturedRequests)
}
},
eventStore = mockEventStore(),
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
)
facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
val req = capturedRequests[0]
assertTrue(req.responseFormat is ResponseFormat.Text)
}
@Test
fun `onUserInput returns RouterResponse with content and mode flag`(): Unit = runBlocking {
val mockStore = mockEventStore()
val facade = facadeWithMocks(eventStore = mockStore, chatMode = ChatMode.STEERING)
val response = facade.onUserInput(sessionId = SessionId("test-session"), input = "Hello!")
assertNotNull(response)
assertEquals("inference response", response.content)
assertTrue(response.steeringEmitted)
// Steering events are deferred to P4-1; no events appended yet.
assertTrue(mockStore.appendedEvents.isEmpty())
}
// --------------------------------------------------------------------------
// Helpers
// --------------------------------------------------------------------------
private fun mockEventStore(): MapBackedEventStore = MapBackedEventStore()
private fun facadeWithMocks(
eventStore: EventStore,
chatMode: ChatMode = ChatMode.CHAT,
): RouterFacade = DefaultRouterFacade(
routerRepository = object : RouterRepository {
override suspend fun getRouterState(sessionId: SessionId): RouterState =
RouterState(
sessionId = sessionId,
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = StageId("s1"),
)
},
routerContextBuilder = object : RouterContextBuilder {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("inference response"),
eventStore = eventStore,
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
).let { impl ->
object : RouterFacade {
override suspend fun onUserInput(sessionId: SessionId, input: String, mode: ChatMode): RouterResponse =
impl.onUserInput(sessionId = sessionId, input = input, mode = chatMode)
}
}
private fun mockInferenceRouter(responseText: String): InferenceRouter =
object : InferenceRouter {
override suspend fun route(
stageId: StageId,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider =
mockProvider(responseText)
}
private fun mockProvider(responseText: String): InferenceProvider =
object : InferenceProvider {
override val id = ProviderId("mock")
override val name = "Mock"
override val tokenizer = MockTokenizer()
override suspend fun infer(request: InferenceRequest): InferenceResponse = InferenceResponse(
requestId = request.requestId,
text = responseText,
finishReason = FinishReason.Stop,
tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
latencyMs = 0,
)
override suspend fun healthCheck(): ProviderHealth =
ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> =
setOf(CapabilityScore(ModelCapability.General, 1.0))
}
private fun mockProviderWithCapture(
responseText: String,
requestIds: MutableList<InferenceRequestId>,
): InferenceProvider =
object : InferenceProvider {
override val id = ProviderId("mock")
override val name = "Mock"
override val tokenizer = MockTokenizer()
override suspend fun infer(request: InferenceRequest): InferenceResponse {
requestIds.add(request.requestId)
return InferenceResponse(
requestId = request.requestId,
text = responseText,
finishReason = FinishReason.Stop,
tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
latencyMs = 0,
)
}
override suspend fun healthCheck(): ProviderHealth =
ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> = emptySet()
}
private fun mockProviderWithRequestCapture(
responseText: String,
requests: MutableList<InferenceRequest>,
): InferenceProvider =
object : InferenceProvider {
override val id = ProviderId("mock")
override val name = "Mock"
override val tokenizer = MockTokenizer()
override suspend fun infer(request: InferenceRequest): InferenceResponse {
requests.add(request)
return InferenceResponse(
requestId = request.requestId,
text = responseText,
finishReason = FinishReason.Stop,
tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
latencyMs = 0,
)
}
override suspend fun healthCheck(): ProviderHealth =
ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> = emptySet()
}
private fun emptyContextPack(): ContextPack = ContextPack(
id = ContextPackId("empty"),
sessionId = SessionId("unknown"),
stageId = StageId("none"),
layers = emptyMap(),
budgetUsed = 0,
budgetLimit = 5000,
)
private class MapBackedEventStore : EventStore {
val appendedEvents = mutableListOf<NewEvent>()
private val storedEvents: MutableMap<EventId, StoredEvent> = mutableMapOf()
private var nextSequence = 1L
override suspend fun append(event: NewEvent): StoredEvent {
appendedEvents.add(event)
val stored = StoredEvent(
metadata = event.metadata,
sequence = nextSequence,
sessionSequence = nextSequence++,
payload = event.payload,
)
storedEvents[event.metadata.eventId] = stored
return stored
}
override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> =
events.map { append(it) }
override fun read(sessionId: SessionId): List<StoredEvent> =
storedEvents.values.filter { it.metadata.sessionId == sessionId }.toList()
override fun readFrom(
sessionId: SessionId,
fromSequence: Long,
): List<StoredEvent> =
read(sessionId).filter { it.sequence >= fromSequence }
override fun lastSequence(sessionId: SessionId): Long? =
read(sessionId).maxOfOrNull { it.sequence }
override fun subscribe(sessionId: SessionId): Flow<StoredEvent> =
throw UnsupportedOperationException("subscribe not implemented for mock")
override fun allEvents(): Sequence<StoredEvent> =
storedEvents.values.asSequence()
override fun allSessionIds(): Set<SessionId> = storedEvents.values.map { it.metadata.sessionId }.toSet()
override fun subscribeAll(): Flow<StoredEvent> = TODO("Not needed in this test context")
override suspend fun lastGlobalSequence(): Long = TODO("Not needed in this test context")
}
}