diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 8498fb31..f37b585e 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -81,6 +81,8 @@ fun main() { val firstProvider: InferenceProvider val infraRegistry: DefaultProviderRegistry 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()) { val settings = correxConfig.modelsSettings @@ -103,7 +105,9 @@ fun main() { }.forEach { infraRegistry.register(it) } // 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. - 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 Runtime.getRuntime().addShutdownHook(Thread { log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id) @@ -238,6 +242,7 @@ fun main() { approvalRepository = repositories.approvalRepository, toolRegistry = toolRegistry, sessionUndoService = sessionUndoService, + modelSwapper = modelSwapper, ) module.start() log.info("==============================") diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt b/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt index 04dc471f..dcc89c5b 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt @@ -49,6 +49,9 @@ class ServerModule( val approvalRepository: DefaultApprovalRepository, val toolRegistry: ToolRegistry, 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. // SupervisorJob so one failure doesn't kill the whole module; // Dispatchers.Default since the work is non-blocking and CPU-light. diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/bridge/DomainEventMapper.kt b/apps/server/src/main/kotlin/com/correx/apps/server/bridge/DomainEventMapper.kt index 40aead26..6e6e6479 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/bridge/DomainEventMapper.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/bridge/DomainEventMapper.kt @@ -14,6 +14,8 @@ import com.correx.core.events.events.InferenceCompletedEvent import com.correx.core.events.events.WorkflowStartedEvent import com.correx.core.events.events.InferenceStartedEvent 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.OrchestrationResumedEvent import com.correx.core.events.events.StageCompletedEvent @@ -190,6 +192,16 @@ suspend fun domainEventToServerMessage( sequence = seq, 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 -> { log.debug( "DomainEventMapper: unmapped payload type={} sessionId={} sequence={}", diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ClientMessage.kt b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ClientMessage.kt index ebf959b6..ee3b6b4a 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ClientMessage.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ClientMessage.kt @@ -62,4 +62,12 @@ sealed class ClientMessage { @Serializable 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() } diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt index 095e31c2..08ca69a8 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt @@ -268,6 +268,21 @@ sealed interface ServerMessage { override val sessionSequence: Long? = null, ) : 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 @SerialName("protocol_error") data class ProtocolError( diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt index d233ece4..dbc0779b 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt @@ -161,9 +161,38 @@ class GlobalStreamHandler(private val module: ServerModule) { 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( msg: ClientMessage.ApprovalResponse, sendFrame: suspend (ServerMessage) -> Unit, diff --git a/apps/server/src/test/kotlin/com/correx/apps/server/bridge/DomainEventMapperTest.kt b/apps/server/src/test/kotlin/com/correx/apps/server/bridge/DomainEventMapperTest.kt index 9176a88d..058f8190 100644 --- a/apps/server/src/test/kotlin/com/correx/apps/server/bridge/DomainEventMapperTest.kt +++ b/apps/server/src/test/kotlin/com/correx/apps/server/bridge/DomainEventMapperTest.kt @@ -490,6 +490,39 @@ class DomainEventMapperTest { 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 fun `unmapped event returns null`(): Unit = runTest { val event = storedEvent(SteeringNoteAddedEvent(sessionId = sessionId, content = "test")) diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt index fe375147..ae1e3123 100644 --- a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt @@ -30,6 +30,13 @@ class ManagedInferenceRouter( private val byId: Map = descriptors.associateBy { it.modelId } private val ordered: List = 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( stageId: StageId, requiredCapabilities: Set, @@ -40,7 +47,8 @@ class ManagedInferenceRouter( requiredCapabilities: Set, modelId: String?, ): InferenceProvider { - val targetId = modelId + val targetId = pinnedModelId + ?: modelId ?: capabilityMatch(requiredCapabilities) ?: defaultModelId val descriptor = byId[targetId] @@ -48,6 +56,29 @@ class ManagedInferenceRouter( 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, * preferring the highest summed score over those required capabilities. Returns null when no diff --git a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt index ba83f9f6..ad2580b1 100644 --- a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt +++ b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt @@ -14,6 +14,7 @@ import com.correx.infrastructure.inference.commons.ResidencyMode import com.correx.testing.fixtures.inference.MockInferenceProvider import kotlinx.coroutines.runBlocking 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.assertThrows @@ -124,6 +125,53 @@ class ManagedInferenceRouterTest { 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 { runBlocking { router.swap("ghost") } } + assertNull(router.pinnedModel()) + } + @Test fun `unknown explicit modelId throws NoEligibleProviderException`() { val manager = FakeModelManager()