feat: correx-managed model lifecycle slice 2 — per-stage model selection

StageConfig.modelId; InferenceRouter gains a backward-compatible model-aware
route() overload (default ignores modelId, so static routers and test fakes are
unaffected). New ManagedInferenceRouter (infra commons) resolves a target model
per stage (stage.modelId > capability match > default) and ensureLoaded()s it
before routing, returning the live managed provider — its id changes with the
resident model, so there is no stale health cache to invalidate on swap.
ModelManager.ensureLoaded default delegates to the idempotent load(). Main wires
the managed router on the [[models]] path; the static path keeps DefaultInferenceRouter.

Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 2 of 5).

apps/server/src/main/kotlin/com/correx/apps/server/Main.kt

@@ -16,6 +16,7 @@ import com.correx.core.inference.CapabilityScore
 import com.correx.core.inference.DefaultInferenceRouter
 import com.correx.core.inference.InferenceProjector
 import com.correx.core.inference.InferenceRepository
+import com.correx.core.inference.InferenceRouter
 import com.correx.core.inference.ModelCapability
 import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
 import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
@@ -47,6 +48,7 @@ import com.correx.core.toolintent.rules.PathContainmentRule
 import com.correx.infrastructure.InfrastructureModule
 import com.correx.infrastructure.inference.DefaultProviderRegistry
 import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
+import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
 import com.correx.core.inference.InferenceProvider
 import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
 import com.correx.infrastructure.tools.FileEditConfig
@@ -78,10 +80,12 @@ fun main() {
     // Static path:  [[models]] absent  → legacy [[providers]] / env-var behavior
     val firstProvider: InferenceProvider
     val infraRegistry: DefaultProviderRegistry
+    val inferenceRouter: InferenceRouter
 
     if (correxConfig.models.isNotEmpty()) {
         val settings = correxConfig.modelsSettings
         val modelManager = InfrastructureModule.createModelManager(settings, eventStore, eventDispatcher)
+        val descriptors = correxConfig.models.map { InfrastructureModule.modelConfigToDescriptor(it) }
         val targetId = settings.defaultModel ?: correxConfig.models.first().id
         val targetModelConfig = correxConfig.models.firstOrNull { it.id == targetId }
             ?: correxConfig.models.first()
@@ -97,6 +101,9 @@ fun main() {
                 else -> { log.error("Unknown provider type: {}", providerConfig.type); null }
             }
         }.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)
         // Shutdown hook to kill the managed llama-server
         Runtime.getRuntime().addShutdownHook(Thread {
             log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
@@ -110,6 +117,7 @@ fun main() {
         require(staticProviders.isNotEmpty()) { "At least one provider must be configured" }
         firstProvider = staticProviders.first()
         infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders)
+        inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())
     }
 
     val modelId = firstProvider.id.value
@@ -161,7 +169,6 @@ fun main() {
         ),
     )
 
-    val inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())
     val engines = OrchestratorEngines(
         transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
         contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()),

core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt

@@ -27,6 +27,18 @@ interface InferenceRouter {
         stageId: StageId,
         requiredCapabilities: Set<ModelCapability>,
     ): InferenceProvider
+
+    /**
+     * Resolves and selects a provider for a stage, optionally pinned to a specific model id.
+     * Implementations that own model lifecycle (e.g. a managed router) use [modelId] to select
+     * and load the target model before routing. The default implementation ignores [modelId] and
+     * falls back to capability-based selection, so static routers remain unaffected.
+     */
+    suspend fun route(
+        stageId: StageId,
+        requiredCapabilities: Set<ModelCapability>,
+        modelId: String?,
+    ): InferenceProvider = route(stageId, requiredCapabilities)
 }
 
 class NoEligibleProviderException(

core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt

@@ -753,7 +753,7 @@ abstract class SessionOrchestrator(
         timeoutMs: Long,
         responseFormat: ResponseFormat = ResponseFormat.Text,
     ): InferenceResult {
-        val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities)
+        val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities, stageConfig.modelId)
         log.debug(
             "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}",
             sessionId.value, stageId.value, provider.id.value, timeoutMs,

core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt

@@ -7,6 +7,7 @@ import com.correx.core.inference.ModelCapability
 
 data class StageConfig(
     val requiredCapabilities: Set<ModelCapability> = emptySet(),
+    val modelId: String? = null,
     val tokenBudget: Int = 4096,
     val generationConfig: GenerationConfig = GenerationConfig(
         temperature = 0.7,

infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt

@@ -0,0 +1,70 @@
+package com.correx.infrastructure.inference.commons
+
+import com.correx.core.events.types.StageId
+import com.correx.core.inference.InferenceProvider
+import com.correx.core.inference.InferenceRouter
+import com.correx.core.inference.ModelCapability
+import com.correx.core.inference.NoEligibleProviderException
+
+/**
+ * An [InferenceRouter] that owns the local model lifecycle: it resolves a single target model
+ * per stage and ensures it is resident (via [ModelManager.ensureLoaded]) before returning the
+ * managed provider to run inference against.
+ *
+ * Resolution precedence (per the model-lifecycle spec):
+ * ```
+ * stage.modelId  >  capability match over configured models  >  defaultModelId
+ * ```
+ * The manual-swap pin (a higher-precedence override) layers in on top of this in a later slice.
+ *
+ * Because the returned provider is the live managed provider (its [InferenceProvider.id] changes
+ * with the resident model), there is no stale router-level health cache to invalidate on swap —
+ * selection always reflects the model the manager just made resident.
+ */
+class ManagedInferenceRouter(
+    private val manager: ModelManager,
+    descriptors: List<ModelDescriptor>,
+    private val defaultModelId: String,
+) : InferenceRouter {
+
+    private val byId: Map<String, ModelDescriptor> = descriptors.associateBy { it.modelId }
+    private val ordered: List<ModelDescriptor> = descriptors
+
+    override suspend fun route(
+        stageId: StageId,
+        requiredCapabilities: Set<ModelCapability>,
+    ): InferenceProvider = route(stageId, requiredCapabilities, null)
+
+    override suspend fun route(
+        stageId: StageId,
+        requiredCapabilities: Set<ModelCapability>,
+        modelId: String?,
+    ): InferenceProvider {
+        val targetId = modelId
+            ?: capabilityMatch(requiredCapabilities)
+            ?: defaultModelId
+        val descriptor = byId[targetId]
+            ?: throw NoEligibleProviderException(stageId, requiredCapabilities)
+        return manager.ensureLoaded(descriptor)
+    }
+
+    /**
+     * Selects the configured model whose declared capabilities cover every required capability,
+     * preferring the highest summed score over those required capabilities. Returns null when no
+     * capabilities are required or none of the models satisfy them (callers fall back to default).
+     */
+    private fun capabilityMatch(required: Set<ModelCapability>): String? {
+        if (required.isEmpty()) return null
+        return ordered
+            .filter { descriptor ->
+                val declared = descriptor.capabilities.map { it.capability }.toSet()
+                declared.containsAll(required)
+            }
+            .maxByOrNull { descriptor ->
+                descriptor.capabilities
+                    .filter { it.capability in required }
+                    .sumOf { it.score }
+            }
+            ?.modelId
+    }
+}

infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt

@@ -16,6 +16,13 @@ interface ModelManager {
      */
     suspend fun load(descriptor: ModelDescriptor): InferenceProvider
 
+    /**
+     * Ensures the model described by [descriptor] is the resident model, returning its provider.
+     * A no-op load if that model is already current; otherwise loads it (evicting the previous
+     * model). [load] is idempotent for the already-loaded model, so the default delegates to it.
+     */
+    suspend fun ensureLoaded(descriptor: ModelDescriptor): InferenceProvider = load(descriptor)
+
     /**
      * Unloads the currently loaded model.
      * 

testing/contracts/build.gradle

@@ -15,6 +15,7 @@ dependencies {
     testImplementation(project(":testing:fixtures"))
     testImplementation(project(":core:inference"))
     testImplementation(project(":core:context"))
+    testImplementation(project(":infrastructure:inference:commons"))
     testFixturesImplementation(project(":core:events"))
     testFixturesImplementation(project(":core:sessions"))
     testFixturesImplementation(project(":core:artifacts-store"))

testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt

@@ -0,0 +1,140 @@
+package com.correx.testing.contracts.model.inference
+
+import com.correx.core.events.types.ProviderId
+import com.correx.core.events.types.StageId
+import com.correx.core.inference.CapabilityScore
+import com.correx.core.inference.InferenceProvider
+import com.correx.core.inference.ModelCapability
+import com.correx.core.inference.NoEligibleProviderException
+import com.correx.core.inference.ProviderHealth
+import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
+import com.correx.infrastructure.inference.commons.ModelDescriptor
+import com.correx.infrastructure.inference.commons.ModelManager
+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.Test
+import org.junit.jupiter.api.assertThrows
+
+class ManagedInferenceRouterTest {
+
+    private val stage = StageId("s1")
+
+    /** Records every (idempotent) load/ensureLoaded so we can assert which model was made resident. */
+    private class FakeModelManager : ModelManager {
+        val loads = mutableListOf<String>()
+        private var current: ModelDescriptor? = null
+
+        override suspend fun load(descriptor: ModelDescriptor): InferenceProvider {
+            loads.add(descriptor.modelId)
+            current = descriptor
+            return MockInferenceProvider(id = ProviderId("managed:${descriptor.modelId}"))
+        }
+
+        override suspend fun unload(modelId: String) { current = null }
+        override fun currentModel(): ModelDescriptor? = current
+        override fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
+    }
+
+    private fun desc(id: String, vararg caps: Pair<ModelCapability, Double>): ModelDescriptor =
+        ModelDescriptor(
+            modelId = id,
+            modelPath = "/models/$id",
+            residencyMode = ResidencyMode.PERSISTENT,
+            capabilities = caps.map { CapabilityScore(it.first, it.second) }.toSet(),
+        )
+
+    @Test
+    fun `stage modelId takes precedence over capability and default`(): 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",
+        )
+
+        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = "coder")
+
+        assertEquals("managed:coder", provider.id.value)
+        assertEquals(listOf("coder"), manager.loads)
+    }
+
+    @Test
+    fun `capability match selects a model that covers the required capabilities`(): Unit = runBlocking {
+        val manager = FakeModelManager()
+        val router = ManagedInferenceRouter(
+            manager = manager,
+            descriptors = listOf(
+                desc("generalist", ModelCapability.General to 1.0),
+                desc("coder", ModelCapability.Coding to 1.0),
+            ),
+            defaultModelId = "generalist",
+        )
+
+        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null)
+
+        assertEquals("managed:coder", provider.id.value)
+        assertEquals(listOf("coder"), manager.loads)
+    }
+
+    @Test
+    fun `capability match prefers the highest summed score`(): Unit = runBlocking {
+        val manager = FakeModelManager()
+        val router = ManagedInferenceRouter(
+            manager = manager,
+            descriptors = listOf(
+                desc("weak", ModelCapability.Coding to 0.3),
+                desc("strong", ModelCapability.Coding to 0.9),
+            ),
+            defaultModelId = "weak",
+        )
+
+        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null)
+
+        assertEquals("managed:strong", provider.id.value)
+    }
+
+    @Test
+    fun `falls back to default when no modelId and no required capabilities`(): 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",
+        )
+
+        val provider = router.route(stage, emptySet(), modelId = null)
+
+        assertEquals("managed:default-model", provider.id.value)
+        assertEquals(listOf("default-model"), manager.loads)
+    }
+
+    @Test
+    fun `falls back to default when no model satisfies the required capabilities`(): Unit = runBlocking {
+        val manager = FakeModelManager()
+        val router = ManagedInferenceRouter(
+            manager = manager,
+            descriptors = listOf(desc("default-model")),
+            defaultModelId = "default-model",
+        )
+
+        val provider = router.route(stage, setOf(ModelCapability.Reasoning), modelId = null)
+
+        assertEquals("managed:default-model", provider.id.value)
+    }
+
+    @Test
+    fun `unknown explicit modelId throws NoEligibleProviderException`() {
+        val manager = FakeModelManager()
+        val router = ManagedInferenceRouter(
+            manager = manager,
+            descriptors = listOf(desc("default-model")),
+            defaultModelId = "default-model",
+        )
+
+        assertThrows<NoEligibleProviderException> {
+            runBlocking { router.route(stage, emptySet(), modelId = "ghost") }
+        }
+    }
+}