test(server): real llama-server WebSocket lifecycle tests
This commit is contained in:
@@ -54,6 +54,9 @@ dependencies {
|
||||
implementation "io.ktor:ktor-server-call-logging:$ktor_version"
|
||||
|
||||
testImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-test:$kotlinx_coroutines_version"
|
||||
testImplementation "io.ktor:ktor-server-test-host:$ktor_version"
|
||||
testImplementation "io.ktor:ktor-client-websockets:$ktor_version"
|
||||
testImplementation "io.ktor:ktor-client-mock:$ktor_version"
|
||||
|
||||
implementation "org.apache.logging.log4j:log4j-core:2.24.1"
|
||||
implementation "org.apache.logging.log4j:log4j-slf4j2-impl:2.24.1"
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
package com.correx.apps.server.lifecycle
|
||||
|
||||
import com.correx.apps.server.ServerModule
|
||||
import com.correx.apps.server.registry.ProviderRegistry
|
||||
import com.correx.apps.server.registry.WorkflowRegistry
|
||||
import com.correx.apps.server.registry.WorkflowSummary
|
||||
import com.correx.apps.server.undo.SessionUndoService
|
||||
import com.correx.core.approvals.domain.DefaultApprovalEngine
|
||||
import com.correx.core.artifactstore.ArtifactStore
|
||||
import com.correx.core.context.builder.DefaultContextPackBuilder
|
||||
import com.correx.core.context.compression.DefaultContextCompressor
|
||||
import com.correx.core.events.EventDispatcher
|
||||
import com.correx.core.events.stores.EventStore
|
||||
import com.correx.core.events.types.ArtifactId
|
||||
import com.correx.core.inference.InferenceProjector
|
||||
import com.correx.core.inference.InferenceProvider
|
||||
import com.correx.core.inference.InferenceRepository
|
||||
import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
|
||||
import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
|
||||
import com.correx.core.kernel.orchestration.OrchestrationConfig
|
||||
import com.correx.core.kernel.orchestration.OrchestrationProjector
|
||||
import com.correx.core.kernel.orchestration.OrchestrationRepository
|
||||
import com.correx.core.kernel.orchestration.OrchestratorEngines
|
||||
import com.correx.core.kernel.orchestration.OrchestratorRepositories
|
||||
import com.correx.core.kernel.retry.DefaultRetryCoordinator
|
||||
import com.correx.core.risk.DefaultRiskAssessor
|
||||
import com.correx.core.router.model.RouterConfig
|
||||
import com.correx.core.sessions.DefaultSessionReducer
|
||||
import com.correx.core.sessions.DefaultSessionRepository
|
||||
import com.correx.core.sessions.SessionProjector
|
||||
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
|
||||
import com.correx.core.transitions.graph.WorkflowGraph
|
||||
import com.correx.core.transitions.resolution.DefaultTransitionResolver
|
||||
import com.correx.core.validation.pipeline.ValidationPipeline
|
||||
import com.correx.infrastructure.InfrastructureModule
|
||||
import com.correx.infrastructure.inference.DefaultProviderRegistry
|
||||
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
|
||||
import com.correx.infrastructure.inference.commons.ResourceProbe
|
||||
import com.correx.infrastructure.inference.commons.UnavailableProbe
|
||||
import com.correx.infrastructure.inference.llama.cpp.DefaultModelManager
|
||||
import com.correx.infrastructure.persistence.InMemoryEventStore
|
||||
import com.correx.infrastructure.tools.FileEditConfig
|
||||
import com.correx.infrastructure.tools.FileReadConfig
|
||||
import com.correx.infrastructure.tools.FileWriteConfig
|
||||
import com.correx.infrastructure.tools.ShellConfig
|
||||
import com.correx.infrastructure.tools.ToolConfig
|
||||
import com.correx.infrastructure.tools.filesystem.FileMutationReverser
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.engine.mock.MockEngine
|
||||
import io.ktor.client.engine.mock.respond
|
||||
import io.ktor.http.ContentType
|
||||
import io.ktor.http.HttpHeaders
|
||||
import io.ktor.http.headersOf
|
||||
import java.nio.file.Path
|
||||
|
||||
/**
|
||||
* Builds a fully-wired [ServerModule] on the managed-model path, using an in-memory event store
|
||||
* and a no-op artifact store. The [modelSwapper] and [resourceProbe] are injected externally so
|
||||
* tests can supply stubs or real implementations.
|
||||
*
|
||||
* The heavy infrastructure (SQLite, CAS, real tool executors) is deliberately avoided: the
|
||||
* lifecycle tests only drive model-swap WS messages and never start sessions or invoke tools.
|
||||
*/
|
||||
fun buildTestServerModule(
|
||||
eventStore: EventStore,
|
||||
modelSwapper: ManagedInferenceRouter,
|
||||
resourceProbe: ResourceProbe = UnavailableProbe,
|
||||
tempDir: Path,
|
||||
): ServerModule {
|
||||
val artifactStore = noopArtifactStore()
|
||||
val eventDispatcher = EventDispatcher(eventStore)
|
||||
|
||||
val firstProvider: InferenceProvider = InfrastructureModule.createLlamaCppProvider(
|
||||
modelId = modelSwapper.availableModelIds().first(),
|
||||
modelPath = "/dev/null",
|
||||
baseUrl = "http://127.0.0.1:1",
|
||||
)
|
||||
|
||||
val infraRegistry: DefaultProviderRegistry = InfrastructureModule.createProviderRegistry(listOf(firstProvider))
|
||||
|
||||
val toolConfig = ToolConfig(
|
||||
shell = ShellConfig(enabled = false, allowedExecutables = emptySet(), workingDir = tempDir),
|
||||
fileRead = FileReadConfig(enabled = false, allowedPaths = setOf(tempDir)),
|
||||
fileWrite = FileWriteConfig(enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir),
|
||||
fileEdit = FileEditConfig(enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir),
|
||||
)
|
||||
val toolRegistry = InfrastructureModule.createToolRegistry(toolConfig)
|
||||
val toolExecutor = InfrastructureModule.createToolExecutor(
|
||||
registry = toolRegistry,
|
||||
eventDispatcher = eventDispatcher,
|
||||
workDir = tempDir,
|
||||
artifactStore = null,
|
||||
)
|
||||
|
||||
val approvalEngine = DefaultApprovalEngine()
|
||||
|
||||
val engines = OrchestratorEngines(
|
||||
transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
|
||||
contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()),
|
||||
inferenceRouter = modelSwapper,
|
||||
validationPipeline = ValidationPipeline(validators = emptyList()),
|
||||
approvalEngine = approvalEngine,
|
||||
riskAssessor = DefaultRiskAssessor(),
|
||||
toolRegistry = toolRegistry,
|
||||
toolExecutor = toolExecutor,
|
||||
)
|
||||
|
||||
val repositories = buildRepositories(eventStore)
|
||||
|
||||
val orchestrator = DefaultSessionOrchestrator(
|
||||
repositories = repositories,
|
||||
engines = engines,
|
||||
retryCoordinator = DefaultRetryCoordinator(eventStore),
|
||||
artifactStore = artifactStore,
|
||||
tokenizer = firstProvider.tokenizer,
|
||||
)
|
||||
|
||||
val routerFacade = InfrastructureModule.createRouterFacade(
|
||||
eventStore = eventStore,
|
||||
inferenceRouter = modelSwapper,
|
||||
config = RouterConfig(),
|
||||
tokenizer = firstProvider.tokenizer,
|
||||
)
|
||||
|
||||
val sessionUndoService = SessionUndoService(
|
||||
eventStore = eventStore,
|
||||
reverser = FileMutationReverser(
|
||||
artifactStore = artifactStore,
|
||||
allowedRoots = setOf(tempDir),
|
||||
),
|
||||
)
|
||||
|
||||
val providerRegistry: ProviderRegistry = object : ProviderRegistry {
|
||||
override fun listAll() = infraRegistry.listAll()
|
||||
override suspend fun healthCheckAll() = infraRegistry.healthCheckAll()
|
||||
}
|
||||
|
||||
val noopWorkflowRegistry = object : WorkflowRegistry {
|
||||
override fun listAll(): List<WorkflowSummary> = emptyList()
|
||||
override fun find(workflowId: String): WorkflowGraph? = null
|
||||
}
|
||||
|
||||
return ServerModule(
|
||||
orchestrator = orchestrator,
|
||||
eventStore = eventStore,
|
||||
artifactStore = artifactStore,
|
||||
sessionRepository = repositories.sessionRepository,
|
||||
workflowRegistry = noopWorkflowRegistry,
|
||||
providerRegistry = providerRegistry,
|
||||
defaultOrchestrationConfig = OrchestrationConfig(sandboxRoot = tempDir),
|
||||
routerFacade = routerFacade,
|
||||
orchestrationRepository = repositories.orchestrationRepository,
|
||||
approvalRepository = repositories.approvalRepository,
|
||||
toolRegistry = toolRegistry,
|
||||
sessionUndoService = sessionUndoService,
|
||||
modelSwapper = modelSwapper,
|
||||
resourceProbe = resourceProbe,
|
||||
)
|
||||
}
|
||||
|
||||
private fun buildRepositories(eventStore: EventStore): OrchestratorRepositories =
|
||||
OrchestratorRepositories(
|
||||
eventStore = eventStore,
|
||||
inferenceRepository = InferenceRepository(DefaultEventReplayer(eventStore, InferenceProjector())),
|
||||
orchestrationRepository = OrchestrationRepository(
|
||||
DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())),
|
||||
),
|
||||
sessionRepository = DefaultSessionRepository(
|
||||
DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())),
|
||||
),
|
||||
artifactRepository = InfrastructureModule.createArtifactRepository(eventStore),
|
||||
approvalRepository = InfrastructureModule.createApprovalRepository(eventStore),
|
||||
)
|
||||
|
||||
private fun noopArtifactStore(): ArtifactStore = object : ArtifactStore {
|
||||
override suspend fun put(bytes: ByteArray): ArtifactId = ArtifactId("noop")
|
||||
override suspend fun get(id: ArtifactId): ByteArray? = null
|
||||
override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
|
||||
}
|
||||
|
||||
fun buildMockHttpClient(): HttpClient {
|
||||
val engine = MockEngine {
|
||||
respond(
|
||||
content = """{"status":"healthy"}""",
|
||||
headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
|
||||
)
|
||||
}
|
||||
return HttpClient(engine)
|
||||
}
|
||||
|
||||
fun buildTestModelManager(
|
||||
llamaBin: String,
|
||||
host: String,
|
||||
port: Int,
|
||||
healthTimeoutMs: Long,
|
||||
eventStore: EventStore,
|
||||
httpClient: HttpClient,
|
||||
): DefaultModelManager {
|
||||
val dispatcher = EventDispatcher(eventStore)
|
||||
return DefaultModelManager(
|
||||
llamaServerBin = llamaBin,
|
||||
host = host,
|
||||
port = port,
|
||||
healthTimeoutMs = healthTimeoutMs,
|
||||
eventStore = eventStore,
|
||||
httpClient = httpClient,
|
||||
eventDispatcher = dispatcher,
|
||||
)
|
||||
}
|
||||
|
||||
fun buildInMemoryEventStore(): InMemoryEventStore = InMemoryEventStore()
|
||||
+193
@@ -0,0 +1,193 @@
|
||||
package com.correx.apps.server.lifecycle
|
||||
|
||||
import com.correx.apps.server.configureServer
|
||||
import com.correx.apps.server.protocol.ClientMessage
|
||||
import com.correx.apps.server.protocol.ServerMessage
|
||||
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
|
||||
import com.correx.infrastructure.inference.commons.ModelDescriptor
|
||||
import com.correx.infrastructure.inference.commons.NvidiaResourceProbe
|
||||
import com.correx.infrastructure.inference.commons.ResidencyMode
|
||||
import com.correx.infrastructure.inference.commons.UnavailableProbe
|
||||
import io.ktor.client.engine.cio.CIO
|
||||
import io.ktor.client.plugins.websocket.WebSockets
|
||||
import io.ktor.client.plugins.websocket.webSocket
|
||||
import io.ktor.websocket.Frame
|
||||
import io.ktor.websocket.readText
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import kotlinx.serialization.encodeToString
|
||||
import kotlinx.serialization.json.Json
|
||||
import org.junit.jupiter.api.AfterEach
|
||||
import org.junit.jupiter.api.Assumptions
|
||||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable
|
||||
import org.junit.jupiter.api.io.TempDir
|
||||
import java.nio.file.Path
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertTrue
|
||||
import io.ktor.server.testing.testApplication
|
||||
|
||||
/**
|
||||
* Live integration test: spawns a real llama-server process and exercises the correx-managed
|
||||
* model lifecycle WebSocket protocol end-to-end.
|
||||
*
|
||||
* Gated on the [CORREX_TEST_LLAMA_BIN] environment variable; skipped automatically in CI.
|
||||
*
|
||||
* **Run command:**
|
||||
* ```
|
||||
* CORREX_TEST_LLAMA_BIN=/path/to/llama-server \
|
||||
* CORREX_TEST_MODEL=/path/to/model-a.gguf \
|
||||
* CORREX_TEST_MODEL_2=/path/to/model-b.gguf \
|
||||
* ./gradlew :apps:server:test --tests '*ModelLifecycleLiveTest*' --rerun-tasks
|
||||
* ```
|
||||
*
|
||||
* [CORREX_TEST_MODEL_2] is optional. When absent, the swap assertion is skipped via
|
||||
* [Assumptions.assumeTrue] and only the load + telemetry assertions run.
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "CORREX_TEST_LLAMA_BIN", matches = ".+")
|
||||
class ModelLifecycleLiveTest {
|
||||
|
||||
private val llamaBin: String = System.getenv("CORREX_TEST_LLAMA_BIN")
|
||||
private val modelAPath: String = System.getenv("CORREX_TEST_MODEL") ?: ""
|
||||
private val modelBPath: String? = System.getenv("CORREX_TEST_MODEL_2")
|
||||
private val modelPort: Int = System.getenv("CORREX_TEST_MODEL_PORT")?.toIntOrNull() ?: 10000
|
||||
|
||||
private var modelManager: com.correx.infrastructure.inference.llama.cpp.DefaultModelManager? = null
|
||||
private var currentModelId: String? = null
|
||||
|
||||
private val protocolJson = Json {
|
||||
classDiscriminator = "type"
|
||||
ignoreUnknownKeys = true
|
||||
}
|
||||
|
||||
private fun decodeServerMessage(text: String): ServerMessage =
|
||||
protocolJson.decodeFromString(text)
|
||||
|
||||
private fun encodeClientMessage(msg: ClientMessage): String =
|
||||
protocolJson.encodeToString(msg)
|
||||
|
||||
@AfterEach
|
||||
fun teardown() {
|
||||
val mgr = modelManager ?: return
|
||||
val id = currentModelId ?: return
|
||||
runBlocking {
|
||||
runCatching { mgr.unload(id) }
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `lifecycle protocol over websocket — live llama-server`(@TempDir tempDir: Path) {
|
||||
val eventStore = buildInMemoryEventStore()
|
||||
val httpClient = io.ktor.client.HttpClient(CIO)
|
||||
|
||||
val descriptorA = ModelDescriptor(
|
||||
modelId = "live-model-a",
|
||||
modelPath = modelAPath,
|
||||
residencyMode = ResidencyMode.PERSISTENT,
|
||||
contextSize = 512,
|
||||
)
|
||||
|
||||
val descriptors = if (modelBPath != null) {
|
||||
listOf(
|
||||
descriptorA,
|
||||
ModelDescriptor(
|
||||
modelId = "live-model-b",
|
||||
modelPath = modelBPath,
|
||||
residencyMode = ResidencyMode.PERSISTENT,
|
||||
contextSize = 512,
|
||||
),
|
||||
)
|
||||
} else {
|
||||
listOf(descriptorA)
|
||||
}
|
||||
|
||||
val mgr = buildTestModelManager(
|
||||
llamaBin = llamaBin,
|
||||
host = "127.0.0.1",
|
||||
port = modelPort,
|
||||
healthTimeoutMs = 120_000L,
|
||||
eventStore = eventStore,
|
||||
httpClient = httpClient,
|
||||
)
|
||||
modelManager = mgr
|
||||
currentModelId = "live-model-a"
|
||||
|
||||
val modelSwapper = ManagedInferenceRouter(mgr, descriptors, defaultModelId = "live-model-a")
|
||||
|
||||
val resourceProbe = if (NvidiaResourceProbe.isAvailable()) {
|
||||
NvidiaResourceProbe(pidSupplier = { mgr.currentPid() })
|
||||
} else {
|
||||
UnavailableProbe
|
||||
}
|
||||
|
||||
val module = buildTestServerModule(
|
||||
eventStore = eventStore,
|
||||
modelSwapper = modelSwapper,
|
||||
resourceProbe = resourceProbe,
|
||||
tempDir = tempDir,
|
||||
)
|
||||
|
||||
runBlocking { mgr.load(descriptorA) }
|
||||
|
||||
testApplication {
|
||||
application { configureServer(module) }
|
||||
val client = createClient { install(WebSockets) }
|
||||
|
||||
client.webSocket("/stream") {
|
||||
var modelList: ServerMessage.ModelList? = null
|
||||
var resourceStatus: ServerMessage.ResourceStatus? = null
|
||||
|
||||
withTimeout(10_000L) {
|
||||
while (modelList == null || resourceStatus == null) {
|
||||
val frame = incoming.receive()
|
||||
if (frame !is Frame.Text) continue
|
||||
val raw = frame.readText()
|
||||
when (val msg = decodeServerMessage(raw)) {
|
||||
is ServerMessage.ModelList -> modelList = msg
|
||||
is ServerMessage.ResourceStatus -> resourceStatus = msg
|
||||
else -> Unit
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
val ml = assertNotNull(modelList, "model.list must arrive in initial snapshot")
|
||||
assertTrue(ml.models.contains("live-model-a"), "model-a must be listed")
|
||||
assertEquals("live-model-a", ml.current)
|
||||
|
||||
assertNotNull(resourceStatus, "resource.status must arrive in initial snapshot")
|
||||
|
||||
// Swap assertion: only when a second model is configured
|
||||
Assumptions.assumeTrue(
|
||||
modelBPath != null,
|
||||
"CORREX_TEST_MODEL_2 not set — skipping swap assertion",
|
||||
)
|
||||
|
||||
currentModelId = "live-model-b"
|
||||
val swapMsg = encodeClientMessage(ClientMessage.SwapModel("live-model-b"))
|
||||
send(Frame.Text(swapMsg))
|
||||
|
||||
var modelChanged: ServerMessage.ModelChanged? = null
|
||||
withTimeout(120_000L) {
|
||||
while (modelChanged == null) {
|
||||
val frame = incoming.receive()
|
||||
if (frame !is Frame.Text) continue
|
||||
val raw = frame.readText()
|
||||
when (val msg = decodeServerMessage(raw)) {
|
||||
is ServerMessage.ModelChanged -> modelChanged = msg
|
||||
is ServerMessage.ProtocolError -> error("Unexpected protocol error during swap: ${msg.message}")
|
||||
else -> Unit
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
val mc = assertNotNull(modelChanged, "model.changed must arrive after SwapModel")
|
||||
assertEquals("live-model-b", mc.modelId)
|
||||
assertTrue(mc.loaded, "loaded must be true after swap")
|
||||
|
||||
val clearMsg = encodeClientMessage(ClientMessage.ClearModelPin)
|
||||
send(Frame.Text(clearMsg))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+184
@@ -0,0 +1,184 @@
|
||||
package com.correx.apps.server.lifecycle
|
||||
|
||||
import com.correx.apps.server.configureServer
|
||||
import com.correx.apps.server.protocol.ClientMessage
|
||||
import com.correx.apps.server.protocol.ServerMessage
|
||||
import kotlinx.serialization.encodeToString
|
||||
import kotlinx.serialization.json.Json
|
||||
import com.correx.infrastructure.inference.commons.GpuStatus
|
||||
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
|
||||
import com.correx.infrastructure.inference.commons.ModelDescriptor
|
||||
import com.correx.infrastructure.inference.commons.ResidencyMode
|
||||
import com.correx.infrastructure.inference.commons.ResourceProbe
|
||||
import com.correx.infrastructure.inference.commons.ResourceSnapshot
|
||||
import io.ktor.client.plugins.websocket.WebSockets
|
||||
import io.ktor.client.plugins.websocket.webSocket
|
||||
import io.ktor.websocket.Frame
|
||||
import io.ktor.websocket.readText
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.io.TempDir
|
||||
import java.nio.file.Path
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertTrue
|
||||
import io.ktor.server.testing.testApplication
|
||||
|
||||
/**
|
||||
* CI test: exercises the correx-managed model lifecycle WebSocket protocol against a
|
||||
* real Ktor [testApplication] server. No actual llama-server binary is spawned — the
|
||||
* [DefaultModelManager] uses the system "sleep" binary + a MockEngine health client
|
||||
* (mirrors the seam used by DefaultModelManagerTest).
|
||||
*
|
||||
* Assertions:
|
||||
* 1. Initial snapshot contains model.list with [model-a, model-b], current=model-a.
|
||||
* 2. Initial snapshot contains resource.status with the fake probe values.
|
||||
* 3. SwapModel("model-b") → model.changed arrives with modelId=model-b, loaded=true.
|
||||
* 4. ClearModelPin does not produce an error frame.
|
||||
*/
|
||||
class ModelLifecycleWiringTest {
|
||||
|
||||
private val protocolJson = Json {
|
||||
classDiscriminator = "type"
|
||||
ignoreUnknownKeys = true
|
||||
}
|
||||
|
||||
private fun decodeServerMessage(text: String): ServerMessage =
|
||||
protocolJson.decodeFromString(text)
|
||||
|
||||
private fun encodeClientMessage(msg: ClientMessage): String =
|
||||
protocolJson.encodeToString(msg)
|
||||
|
||||
@Test
|
||||
fun `lifecycle protocol over websocket`(@TempDir tempDir: Path): Unit {
|
||||
val eventStore = buildInMemoryEventStore()
|
||||
val httpClient = buildMockHttpClient()
|
||||
|
||||
val descriptorA = ModelDescriptor(
|
||||
modelId = "model-a",
|
||||
modelPath = "/dev/null",
|
||||
residencyMode = ResidencyMode.PERSISTENT,
|
||||
contextSize = 512,
|
||||
)
|
||||
val descriptorB = ModelDescriptor(
|
||||
modelId = "model-b",
|
||||
modelPath = "/dev/null",
|
||||
residencyMode = ResidencyMode.PERSISTENT,
|
||||
contextSize = 512,
|
||||
)
|
||||
|
||||
val modelManager = buildTestModelManager(
|
||||
llamaBin = "sleep",
|
||||
host = "127.0.0.1",
|
||||
port = 19888,
|
||||
healthTimeoutMs = 5_000L,
|
||||
eventStore = eventStore,
|
||||
httpClient = httpClient,
|
||||
)
|
||||
|
||||
// Pre-load model-a so that currentModelId() returns "model-a" in the initial snapshot,
|
||||
// mirroring Main's runBlocking { modelManager.load(descriptor) } call.
|
||||
runBlocking { modelManager.load(descriptorA) }
|
||||
|
||||
val descriptors = listOf(descriptorA, descriptorB)
|
||||
val modelSwapper = ManagedInferenceRouter(modelManager, descriptors, defaultModelId = "model-a")
|
||||
|
||||
val fakeProbe = ResourceProbe {
|
||||
ResourceSnapshot(
|
||||
gpu = GpuStatus(memoryUsedMb = 2048L, memoryTotalMb = 8192L, utilizationPct = 42),
|
||||
processRssBytes = 512L * 1024L * 1024L,
|
||||
)
|
||||
}
|
||||
|
||||
val module = buildTestServerModule(
|
||||
eventStore = eventStore,
|
||||
modelSwapper = modelSwapper,
|
||||
resourceProbe = fakeProbe,
|
||||
tempDir = tempDir,
|
||||
)
|
||||
|
||||
try {
|
||||
testApplication {
|
||||
application { configureServer(module) }
|
||||
val client = createClient { install(WebSockets) }
|
||||
|
||||
client.webSocket("/stream") {
|
||||
// Collect frames until we have both model.list and resource.status
|
||||
var modelList: ServerMessage.ModelList? = null
|
||||
var resourceStatus: ServerMessage.ResourceStatus? = null
|
||||
|
||||
// Pass 1: collect initial snapshot messages
|
||||
withTimeout(5_000L) {
|
||||
while (modelList == null || resourceStatus == null) {
|
||||
val frame = incoming.receive()
|
||||
if (frame !is Frame.Text) continue
|
||||
val raw = frame.readText()
|
||||
when (val msg = decodeServerMessage(raw)) {
|
||||
is ServerMessage.ModelList -> modelList = msg
|
||||
is ServerMessage.ResourceStatus -> resourceStatus = msg
|
||||
else -> Unit
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
val ml = assertNotNull(modelList, "model.list must arrive in initial snapshot")
|
||||
assertEquals(listOf("model-a", "model-b"), ml.models)
|
||||
assertEquals("model-a", ml.current)
|
||||
|
||||
val rs = assertNotNull(resourceStatus, "resource.status must arrive in initial snapshot")
|
||||
assertEquals(2048L, rs.gpuMemoryUsedMb)
|
||||
assertEquals(8192L, rs.gpuMemoryTotalMb)
|
||||
assertEquals(42, rs.gpuUtilizationPct)
|
||||
assertNotNull(rs.processRssMb)
|
||||
|
||||
// Send SwapModel("model-b")
|
||||
val swapMsg = encodeClientMessage(ClientMessage.SwapModel("model-b"))
|
||||
send(Frame.Text(swapMsg))
|
||||
|
||||
// Wait for model.changed with modelId=model-b, loaded=true
|
||||
var modelChanged: ServerMessage.ModelChanged? = null
|
||||
withTimeout(8_000L) {
|
||||
while (modelChanged == null) {
|
||||
val frame = incoming.receive()
|
||||
if (frame !is Frame.Text) continue
|
||||
val raw = frame.readText()
|
||||
when (val msg = decodeServerMessage(raw)) {
|
||||
is ServerMessage.ModelChanged -> modelChanged = msg
|
||||
is ServerMessage.ProtocolError -> error("Unexpected protocol error: ${msg.message}")
|
||||
else -> Unit
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
val mc = assertNotNull(modelChanged, "model.changed must arrive after SwapModel")
|
||||
assertEquals("model-b", mc.modelId)
|
||||
assertTrue(mc.loaded, "loaded must be true after swap")
|
||||
|
||||
// Send ClearModelPin — the server responds with nothing (just clears the pin);
|
||||
// verify no ProtocolError arrives within a brief window.
|
||||
val clearMsg = encodeClientMessage(ClientMessage.ClearModelPin)
|
||||
send(Frame.Text(clearMsg))
|
||||
|
||||
// Give the server a moment to process ClearModelPin; if it sends back a
|
||||
// ProtocolError it would arrive almost immediately (well before 2500ms resource push).
|
||||
delay(100)
|
||||
val pending = incoming.tryReceive()
|
||||
if (pending.isSuccess) {
|
||||
val f = pending.getOrThrow()
|
||||
if (f is Frame.Text) {
|
||||
val msg = decodeServerMessage(f.readText())
|
||||
assertTrue(msg !is ServerMessage.ProtocolError, "ClearModelPin produced error: $msg")
|
||||
}
|
||||
}
|
||||
// If no frame arrived — which is also fine (no error sent).
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
runBlocking {
|
||||
runCatching { modelManager.unload(modelSwapper.currentModelId() ?: "model-b") }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user