feat: correx-managed model lifecycle slice 3 — manual swap + pin
ManagedInferenceRouter owns a live @Volatile pin (decision D1): swap(modelId) makes a model resident and pins it; clearPin() releases it. Pin is highest precedence in route() (pin > stage.modelId > capability > default). New SwapModel/ClearModelPin client messages handled in GlobalStreamHandler via ServerModule.modelSwapper (null on the static path). ServerMessage.ModelChanged (model.changed) mapped from ModelLoadedEvent/ModelUnloadedEvent — the swap's load event surfaces to clients through streamGlobal, so the handler doesn't push it directly. Tests: pin precedence + swap/clear, Model* -> ModelChanged. Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 3 of 5).
This commit is contained in:
@@ -81,6 +81,8 @@ fun main() {
|
|||||||
val firstProvider: InferenceProvider
|
val firstProvider: InferenceProvider
|
||||||
val infraRegistry: DefaultProviderRegistry
|
val infraRegistry: DefaultProviderRegistry
|
||||||
val inferenceRouter: InferenceRouter
|
val inferenceRouter: InferenceRouter
|
||||||
|
// Non-null only on the managed path; backs manual model swap/pin from clients.
|
||||||
|
var modelSwapper: ManagedInferenceRouter? = null
|
||||||
|
|
||||||
if (correxConfig.models.isNotEmpty()) {
|
if (correxConfig.models.isNotEmpty()) {
|
||||||
val settings = correxConfig.modelsSettings
|
val settings = correxConfig.modelsSettings
|
||||||
@@ -103,7 +105,9 @@ fun main() {
|
|||||||
}.forEach { infraRegistry.register(it) }
|
}.forEach { infraRegistry.register(it) }
|
||||||
// Managed router owns per-stage model selection: it ensures the resolved model is resident
|
// Managed router owns per-stage model selection: it ensures the resolved model is resident
|
||||||
// before routing (stage.modelId > capability match > default), evicting the prior model.
|
// before routing (stage.modelId > capability match > default), evicting the prior model.
|
||||||
inferenceRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
|
val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
|
||||||
|
inferenceRouter = managedRouter
|
||||||
|
modelSwapper = managedRouter
|
||||||
// Shutdown hook to kill the managed llama-server
|
// Shutdown hook to kill the managed llama-server
|
||||||
Runtime.getRuntime().addShutdownHook(Thread {
|
Runtime.getRuntime().addShutdownHook(Thread {
|
||||||
log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
|
log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
|
||||||
@@ -238,6 +242,7 @@ fun main() {
|
|||||||
approvalRepository = repositories.approvalRepository,
|
approvalRepository = repositories.approvalRepository,
|
||||||
toolRegistry = toolRegistry,
|
toolRegistry = toolRegistry,
|
||||||
sessionUndoService = sessionUndoService,
|
sessionUndoService = sessionUndoService,
|
||||||
|
modelSwapper = modelSwapper,
|
||||||
)
|
)
|
||||||
module.start()
|
module.start()
|
||||||
log.info("==============================")
|
log.info("==============================")
|
||||||
|
|||||||
@@ -49,6 +49,9 @@ class ServerModule(
|
|||||||
val approvalRepository: DefaultApprovalRepository,
|
val approvalRepository: DefaultApprovalRepository,
|
||||||
val toolRegistry: ToolRegistry,
|
val toolRegistry: ToolRegistry,
|
||||||
val sessionUndoService: com.correx.apps.server.undo.SessionUndoService,
|
val sessionUndoService: com.correx.apps.server.undo.SessionUndoService,
|
||||||
|
// The managed-model router when correx owns the local model process ([[models]] configured);
|
||||||
|
// null on the static-provider path. Backs the manual SwapModel / ClearModelPin operations.
|
||||||
|
val modelSwapper: com.correx.infrastructure.inference.commons.ManagedInferenceRouter? = null,
|
||||||
// Long-lived scope owned by the module — backs the event-store subscription.
|
// Long-lived scope owned by the module — backs the event-store subscription.
|
||||||
// SupervisorJob so one failure doesn't kill the whole module;
|
// SupervisorJob so one failure doesn't kill the whole module;
|
||||||
// Dispatchers.Default since the work is non-blocking and CPU-light.
|
// Dispatchers.Default since the work is non-blocking and CPU-light.
|
||||||
|
|||||||
@@ -14,6 +14,8 @@ import com.correx.core.events.events.InferenceCompletedEvent
|
|||||||
import com.correx.core.events.events.WorkflowStartedEvent
|
import com.correx.core.events.events.WorkflowStartedEvent
|
||||||
import com.correx.core.events.events.InferenceStartedEvent
|
import com.correx.core.events.events.InferenceStartedEvent
|
||||||
import com.correx.core.events.events.InferenceTimeoutEvent
|
import com.correx.core.events.events.InferenceTimeoutEvent
|
||||||
|
import com.correx.core.events.events.ModelLoadedEvent
|
||||||
|
import com.correx.core.events.events.ModelUnloadedEvent
|
||||||
import com.correx.core.events.events.OrchestrationPausedEvent
|
import com.correx.core.events.events.OrchestrationPausedEvent
|
||||||
import com.correx.core.events.events.OrchestrationResumedEvent
|
import com.correx.core.events.events.OrchestrationResumedEvent
|
||||||
import com.correx.core.events.events.StageCompletedEvent
|
import com.correx.core.events.events.StageCompletedEvent
|
||||||
@@ -190,6 +192,16 @@ suspend fun domainEventToServerMessage(
|
|||||||
sequence = seq,
|
sequence = seq,
|
||||||
sessionSequence = sessionSequence,
|
sessionSequence = sessionSequence,
|
||||||
)
|
)
|
||||||
|
is ModelLoadedEvent -> ServerMessage.ModelChanged(
|
||||||
|
modelId = p.modelId,
|
||||||
|
providerId = p.providerId.value,
|
||||||
|
loaded = true,
|
||||||
|
)
|
||||||
|
is ModelUnloadedEvent -> ServerMessage.ModelChanged(
|
||||||
|
modelId = p.modelId,
|
||||||
|
providerId = p.providerId.value,
|
||||||
|
loaded = false,
|
||||||
|
)
|
||||||
else -> {
|
else -> {
|
||||||
log.debug(
|
log.debug(
|
||||||
"DomainEventMapper: unmapped payload type={} sessionId={} sequence={}",
|
"DomainEventMapper: unmapped payload type={} sessionId={} sequence={}",
|
||||||
|
|||||||
@@ -62,4 +62,12 @@ sealed class ClientMessage {
|
|||||||
|
|
||||||
@Serializable
|
@Serializable
|
||||||
data class ChatInput(val sessionId: SessionId, val text: String, val mode: ChatMode) : ClientMessage()
|
data class ChatInput(val sessionId: SessionId, val text: String, val mode: ChatMode) : ClientMessage()
|
||||||
|
|
||||||
|
/** Operator request to make [modelId] resident and pin it (beats per-stage selection). */
|
||||||
|
@Serializable
|
||||||
|
data class SwapModel(val modelId: String) : ClientMessage()
|
||||||
|
|
||||||
|
/** Operator request to clear the manual-swap pin; per-stage selection resumes. */
|
||||||
|
@Serializable
|
||||||
|
data object ClearModelPin : ClientMessage()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -268,6 +268,21 @@ sealed interface ServerMessage {
|
|||||||
override val sessionSequence: Long? = null,
|
override val sessionSequence: Long? = null,
|
||||||
) : ServerMessage, NonEventMessage
|
) : ServerMessage, NonEventMessage
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The resident local model changed (load or unload). Derived from `ModelLoadedEvent` /
|
||||||
|
* `ModelUnloadedEvent`; a global control message (no session cursors). `loaded = false`
|
||||||
|
* signals the model was unloaded.
|
||||||
|
*/
|
||||||
|
@Serializable
|
||||||
|
@SerialName("model.changed")
|
||||||
|
data class ModelChanged(
|
||||||
|
val modelId: String,
|
||||||
|
val providerId: String,
|
||||||
|
val loaded: Boolean,
|
||||||
|
override val sequence: Long? = null,
|
||||||
|
override val sessionSequence: Long? = null,
|
||||||
|
) : ServerMessage, NonEventMessage
|
||||||
|
|
||||||
@Serializable
|
@Serializable
|
||||||
@SerialName("protocol_error")
|
@SerialName("protocol_error")
|
||||||
data class ProtocolError(
|
data class ProtocolError(
|
||||||
|
|||||||
@@ -161,9 +161,38 @@ class GlobalStreamHandler(private val module: ServerModule) {
|
|||||||
sendFrame(ServerMessage.ProtocolError("Router error: ${it.message}"))
|
sendFrame(ServerMessage.ProtocolError("Router error: ${it.message}"))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
is ClientMessage.SwapModel -> handleSwapModel(msg, sendFrame)
|
||||||
|
is ClientMessage.ClearModelPin -> {
|
||||||
|
val swapper = module.modelSwapper
|
||||||
|
if (swapper == null) {
|
||||||
|
sendFrame(errorResponse("Model management is not enabled"))
|
||||||
|
} else {
|
||||||
|
swapper.clearPin()
|
||||||
|
log.info("model pin cleared")
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private suspend fun handleSwapModel(
|
||||||
|
msg: ClientMessage.SwapModel,
|
||||||
|
sendFrame: suspend (ServerMessage) -> Unit,
|
||||||
|
) {
|
||||||
|
val swapper = module.modelSwapper
|
||||||
|
if (swapper == null) {
|
||||||
|
sendFrame(errorResponse("Model management is not enabled"))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// swap() makes the model resident; the resulting ModelLoadedEvent surfaces the change to
|
||||||
|
// clients via streamGlobal, so we don't push ModelChanged here.
|
||||||
|
runCatching { swapper.swap(msg.modelId) }
|
||||||
|
.onSuccess { log.info("swapped to model={}", msg.modelId) }
|
||||||
|
.onFailure {
|
||||||
|
log.error("model swap failed for '{}': {}", msg.modelId, it.message)
|
||||||
|
sendFrame(errorResponse("Model swap failed: ${it.message}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
private suspend fun handleApprovalResponse(
|
private suspend fun handleApprovalResponse(
|
||||||
msg: ClientMessage.ApprovalResponse,
|
msg: ClientMessage.ApprovalResponse,
|
||||||
sendFrame: suspend (ServerMessage) -> Unit,
|
sendFrame: suspend (ServerMessage) -> Unit,
|
||||||
|
|||||||
@@ -490,6 +490,39 @@ class DomainEventMapperTest {
|
|||||||
assertNull(result.reason)
|
assertNull(result.reason)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `ModelLoadedEvent maps to ModelChanged loaded=true`(): Unit = runTest {
|
||||||
|
val event = storedEvent(
|
||||||
|
com.correx.core.events.events.ModelLoadedEvent(
|
||||||
|
modelId = "qwen-coder",
|
||||||
|
providerId = ProviderId("llama-cpp:qwen-coder"),
|
||||||
|
sessionId = sessionId,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
val result = domainEventToServerMessage(event, noopStore, sessionSequence = 0L)
|
||||||
|
assertEquals(
|
||||||
|
ServerMessage.ModelChanged("qwen-coder", "llama-cpp:qwen-coder", loaded = true),
|
||||||
|
result,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `ModelUnloadedEvent maps to ModelChanged loaded=false`(): Unit = runTest {
|
||||||
|
val event = storedEvent(
|
||||||
|
com.correx.core.events.events.ModelUnloadedEvent(
|
||||||
|
modelId = "qwen-coder",
|
||||||
|
providerId = ProviderId("llama-cpp"),
|
||||||
|
sessionId = sessionId,
|
||||||
|
cancellationReason = null,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
val result = domainEventToServerMessage(event, noopStore, sessionSequence = 0L)
|
||||||
|
assertEquals(
|
||||||
|
ServerMessage.ModelChanged("qwen-coder", "llama-cpp", loaded = false),
|
||||||
|
result,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun `unmapped event returns null`(): Unit = runTest {
|
fun `unmapped event returns null`(): Unit = runTest {
|
||||||
val event = storedEvent(SteeringNoteAddedEvent(sessionId = sessionId, content = "test"))
|
val event = storedEvent(SteeringNoteAddedEvent(sessionId = sessionId, content = "test"))
|
||||||
|
|||||||
+32
-1
@@ -30,6 +30,13 @@ class ManagedInferenceRouter(
|
|||||||
private val byId: Map<String, ModelDescriptor> = descriptors.associateBy { it.modelId }
|
private val byId: Map<String, ModelDescriptor> = descriptors.associateBy { it.modelId }
|
||||||
private val ordered: List<ModelDescriptor> = descriptors
|
private val ordered: List<ModelDescriptor> = descriptors
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Manual-swap pin — global operator intent. Live runtime state (decision D1): not event-sourced,
|
||||||
|
* so a restart reverts to the default model. When set, it overrides per-stage selection.
|
||||||
|
*/
|
||||||
|
@Volatile
|
||||||
|
private var pinnedModelId: String? = null
|
||||||
|
|
||||||
override suspend fun route(
|
override suspend fun route(
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
requiredCapabilities: Set<ModelCapability>,
|
requiredCapabilities: Set<ModelCapability>,
|
||||||
@@ -40,7 +47,8 @@ class ManagedInferenceRouter(
|
|||||||
requiredCapabilities: Set<ModelCapability>,
|
requiredCapabilities: Set<ModelCapability>,
|
||||||
modelId: String?,
|
modelId: String?,
|
||||||
): InferenceProvider {
|
): InferenceProvider {
|
||||||
val targetId = modelId
|
val targetId = pinnedModelId
|
||||||
|
?: modelId
|
||||||
?: capabilityMatch(requiredCapabilities)
|
?: capabilityMatch(requiredCapabilities)
|
||||||
?: defaultModelId
|
?: defaultModelId
|
||||||
val descriptor = byId[targetId]
|
val descriptor = byId[targetId]
|
||||||
@@ -48,6 +56,29 @@ class ManagedInferenceRouter(
|
|||||||
return manager.ensureLoaded(descriptor)
|
return manager.ensureLoaded(descriptor)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Manual swap from the operator: make [modelId] resident and pin it so it beats per-stage
|
||||||
|
* selection for all later work until [clearPin]. The load emits a `ModelLoadedEvent` (when the
|
||||||
|
* resident model actually changes), which surfaces to clients via the event stream.
|
||||||
|
*
|
||||||
|
* @throws NoEligibleProviderException if [modelId] is not a configured model.
|
||||||
|
*/
|
||||||
|
suspend fun swap(modelId: String): InferenceProvider {
|
||||||
|
val descriptor = byId[modelId]
|
||||||
|
?: throw NoEligibleProviderException(StageId("manual-swap"), emptySet())
|
||||||
|
val provider = manager.ensureLoaded(descriptor)
|
||||||
|
pinnedModelId = modelId
|
||||||
|
return provider
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Clear the manual-swap pin; later stages resume per-stage selection. */
|
||||||
|
fun clearPin() {
|
||||||
|
pinnedModelId = null
|
||||||
|
}
|
||||||
|
|
||||||
|
/** The currently pinned model id, or null when no manual swap is active. */
|
||||||
|
fun pinnedModel(): String? = pinnedModelId
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Selects the configured model whose declared capabilities cover every required capability,
|
* Selects the configured model whose declared capabilities cover every required capability,
|
||||||
* preferring the highest summed score over those required capabilities. Returns null when no
|
* preferring the highest summed score over those required capabilities. Returns null when no
|
||||||
|
|||||||
+48
@@ -14,6 +14,7 @@ import com.correx.infrastructure.inference.commons.ResidencyMode
|
|||||||
import com.correx.testing.fixtures.inference.MockInferenceProvider
|
import com.correx.testing.fixtures.inference.MockInferenceProvider
|
||||||
import kotlinx.coroutines.runBlocking
|
import kotlinx.coroutines.runBlocking
|
||||||
import org.junit.jupiter.api.Assertions.assertEquals
|
import org.junit.jupiter.api.Assertions.assertEquals
|
||||||
|
import org.junit.jupiter.api.Assertions.assertNull
|
||||||
import org.junit.jupiter.api.Test
|
import org.junit.jupiter.api.Test
|
||||||
import org.junit.jupiter.api.assertThrows
|
import org.junit.jupiter.api.assertThrows
|
||||||
|
|
||||||
@@ -124,6 +125,53 @@ class ManagedInferenceRouterTest {
|
|||||||
assertEquals("managed:default-model", provider.id.value)
|
assertEquals("managed:default-model", provider.id.value)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `swap pins a model that then beats per-stage selection`(): Unit = runBlocking {
|
||||||
|
val manager = FakeModelManager()
|
||||||
|
val router = ManagedInferenceRouter(
|
||||||
|
manager = manager,
|
||||||
|
descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
|
||||||
|
defaultModelId = "default-model",
|
||||||
|
)
|
||||||
|
|
||||||
|
router.swap("coder")
|
||||||
|
assertEquals("coder", router.pinnedModel())
|
||||||
|
|
||||||
|
// A stage explicitly asking for a different model is overridden by the pin.
|
||||||
|
val provider = router.route(stage, emptySet(), modelId = "default-model")
|
||||||
|
assertEquals("managed:coder", provider.id.value)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `clearPin restores per-stage selection`(): Unit = runBlocking {
|
||||||
|
val manager = FakeModelManager()
|
||||||
|
val router = ManagedInferenceRouter(
|
||||||
|
manager = manager,
|
||||||
|
descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
|
||||||
|
defaultModelId = "default-model",
|
||||||
|
)
|
||||||
|
|
||||||
|
router.swap("coder")
|
||||||
|
router.clearPin()
|
||||||
|
assertNull(router.pinnedModel())
|
||||||
|
|
||||||
|
val provider = router.route(stage, emptySet(), modelId = "default-model")
|
||||||
|
assertEquals("managed:default-model", provider.id.value)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `swap to an unknown model throws and leaves the pin unset`() {
|
||||||
|
val manager = FakeModelManager()
|
||||||
|
val router = ManagedInferenceRouter(
|
||||||
|
manager = manager,
|
||||||
|
descriptors = listOf(desc("default-model")),
|
||||||
|
defaultModelId = "default-model",
|
||||||
|
)
|
||||||
|
|
||||||
|
assertThrows<NoEligibleProviderException> { runBlocking { router.swap("ghost") } }
|
||||||
|
assertNull(router.pinnedModel())
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun `unknown explicit modelId throws NoEligibleProviderException`() {
|
fun `unknown explicit modelId throws NoEligibleProviderException`() {
|
||||||
val manager = FakeModelManager()
|
val manager = FakeModelManager()
|
||||||
|
|||||||
Reference in New Issue
Block a user