feat(inference): capability-aware routing for tool-heavy stages

Add CapabilityAwareRoutingStrategy — it hard-filters to providers that declare every
required capability (like FirstAvailable) but ranks the matches by summed required-capability
score; ties keep list order, so it is a strict superset of first-available. The server now
wires it as the routing policy.

The orchestrator augments a stage's required capabilities with ModelCapability.ToolCalling
when the stage grants tools, so tool-heavy stages route to the best tool-calling model rather
than whichever healthy provider comes first.

Scores come from each provider's declared capabilities(); observed per-model reliability
(GET /metrics/tool-reliability) can feed in later as an additional weight.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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

@@ -76,7 +76,7 @@ import com.correx.core.events.types.SessionId
 import com.correx.infrastructure.InfrastructureModule
 import com.correx.infrastructure.inference.DefaultProviderRegistry
 import com.correx.infrastructure.workflow.ExecutionPlanCompiler
-import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
+import com.correx.infrastructure.inference.CapabilityAwareRoutingStrategy
 import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
 import com.correx.infrastructure.inference.commons.AmdResourceProbe
 import com.correx.infrastructure.inference.commons.NvidiaResourceProbe
@@ -170,7 +170,7 @@ fun main() {
         require(staticProviders.isNotEmpty()) { "At least one provider must be configured" }
         firstProvider = staticProviders.first()
         infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders)
-        inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())
+        inferenceRouter = DefaultInferenceRouter(infraRegistry, CapabilityAwareRoutingStrategy())
         llamaBaseUrl = correxConfig.providers.firstOrNull()?.url
             ?: System.getenv("CORREX_LLAMA_URL")
             ?: "http://127.0.0.1:10000"

core/kernel/AGENTS.md

@@ -14,6 +14,7 @@ CORREX kernel team. This is the integration point for all other `core/` modules.
 - `OrchestrationState` / `OrchestrationReducer` (`DefaultOrchestrationReducer`) / `OrchestrationProjector` / `OrchestrationRepository` — standard event-sourcing stack for orchestration state.
 - `RetryCoordinator` / `DefaultRetryCoordinator` — manages retry logic per `RetryPolicy`.
 - On a recoverable tool failure `dispatchToolCalls` feeds the failing tool's argument schema back into context alongside the error (`toolArgsHint`), so the model self-corrects a malformed call instead of repeating it — the contract stays strict; the feedback is what loosens.
+- Stages that grant tools (`allowedTools` non-empty) request `ModelCapability.ToolCalling` on top of their declared capabilities when routing, so the capability-aware strategy steers them to the best tool-calling model.
 - `ApprovalGateway` — kernel-side approval bridge; calls `core:approvals` engine before executing gated operations. Per-tool gating in `dispatchToolCalls` builds the `ApprovalContext` mode from the session's bound operator profile (`boundProfile.approvalMode`, mapped by `approvalModeFor`): unset/`prompt` keeps a human in the loop (default), `auto` auto-approves up to T2, `yolo` all tiers, `deny` blocks above T0. The engine is always consulted (Invariant #4 holds); policy/plane-2 BLOCK stays terminal regardless of mode.
 - `ReplayOrchestrator` / `ReplayInferenceProvider` / `ReplayStrategy` — deterministic replay of a session from its event log. `ReplayInferenceProvider` returns recorded responses — no live LLM (Hard Invariant #8).
 - `SubagentRunner` / `InSessionSubagentRunner` — runs sub-agent invocations within an active session.

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

@@ -106,6 +106,7 @@ import com.correx.core.inference.ResponseFormat
 import com.correx.core.inference.Tokenizer
 import com.correx.core.inference.ToolCallRequest
 import com.correx.core.inference.ToolDefinition
+import com.correx.core.inference.ModelCapability
 import com.correx.core.sessions.ApprovalMode
 import com.correx.core.inference.ToolFunction
 import com.correx.core.kernel.execution.WorkflowResult
@@ -1753,7 +1754,12 @@ abstract class SessionOrchestrator(
         responseFormat: ResponseFormat = ResponseFormat.Text,
         effectives: RunEffectives = RunEffectives(toolRegistry, toolExecutor, workspacePolicy),
     ): InferenceResult {
-        val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities, stageConfig.modelId)
+        // A stage that grants tools needs a tool-calling model, so request ToolCalling on top of any
+        // declared capabilities — the capability-aware strategy then ranks eligible providers by their
+        // ToolCalling score and routes the stage to the best tool-caller.
+        val requiredCapabilities = stageConfig.requiredCapabilities +
+            if (stageConfig.allowedTools.isNotEmpty()) setOf(ModelCapability.ToolCalling) else emptySet()
+        val provider = inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId)
         log.debug(
             "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}",
             sessionId.value, stageId.value, provider.id.value, timeoutMs,

infrastructure/inference/AGENTS.md

@@ -11,7 +11,7 @@ Adapter for LLM inference backends. Implements `core:inference` interfaces. No d
 ## Local Contracts
 
 - `DefaultProviderRegistry` implements `ProviderRegistry` from `core:inference`.
-- `FirstAvailableRoutingStrategy` is the default routing policy; extend only in `core:inference`, not here.
+- Two `RoutingStrategy` implementations (both hard-filter to providers declaring every required capability): `FirstAvailableRoutingStrategy` picks the first match; `CapabilityAwareRoutingStrategy` ranks the matches by summed required-capability score (ties keep list order, so it's a superset of first-available). The server wires `CapabilityAwareRoutingStrategy` so tool-heavy stages (which request `ModelCapability.ToolCalling`) route to the best tool-caller. Extend the `RoutingStrategy` interface only in `core:inference`, not here.
 - All LLM responses are proposals — they must be validated by the core before affecting state (invariant #7). Adapters return raw responses; they do not validate.
 - Network calls to inference backends are environment observations; results must be recorded as events by callers if replay must reproduce them (invariant #9).
 - Model lifecycle (spawn/own the llama-server process) lives in `llama_cpp/`; the autonomous scheduler is intentionally unbuilt (see root CLAUDE.md).

infrastructure/inference/src/main/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategy.kt

@@ -0,0 +1,38 @@
+package com.correx.infrastructure.inference
+
+import com.correx.core.events.types.StageId
+import com.correx.core.inference.InferenceProvider
+import com.correx.core.inference.ModelCapability
+import com.correx.core.inference.NoEligibleProviderException
+import com.correx.core.inference.RoutingStrategy
+
+/**
+ * Like [FirstAvailableRoutingStrategy] it hard-filters to providers that declare every
+ * [requiredCapabilities] — but among those it picks the one with the highest summed capability
+ * **score** for the required set (falling back to total declared score when nothing is required).
+ * Ties keep list order, so behaviour is identical to first-available when scores are equal.
+ *
+ * This is the seam the orchestrator relies on to steer tool-heavy stages (which request
+ * `ModelCapability.ToolCalling`) to the model that calls tools best, rather than whichever healthy
+ * provider happens to come first. The scores come from each provider's declared
+ * `capabilities()` (config-supplied); observed per-model reliability (the `/metrics/tool-reliability`
+ * groundwork) can later feed in as an additional weight.
+ */
+class CapabilityAwareRoutingStrategy : RoutingStrategy {
+    override fun select(
+        candidates: List<InferenceProvider>,
+        requiredCapabilities: Set<ModelCapability>,
+    ): InferenceProvider {
+        val eligible = candidates.filter { provider ->
+            provider.capabilities().map { it.capability }.toSet().containsAll(requiredCapabilities)
+        }
+        // maxByOrNull keeps the first element on ties, preserving first-available order.
+        return eligible.maxByOrNull { score(it, requiredCapabilities) }
+            ?: throw NoEligibleProviderException(StageId("routing"), requiredCapabilities)
+    }
+
+    private fun score(provider: InferenceProvider, required: Set<ModelCapability>): Double {
+        val scores = provider.capabilities().associate { it.capability to it.score }
+        return if (required.isEmpty()) scores.values.sum() else required.sumOf { scores[it] ?: 0.0 }
+    }
+}

infrastructure/inference/src/test/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategyTest.kt

@@ -0,0 +1,73 @@
+package com.correx.infrastructure.inference
+
+import com.correx.core.events.types.ProviderId
+import com.correx.core.inference.CapabilityScore
+import com.correx.core.inference.InferenceProvider
+import com.correx.core.inference.InferenceRequest
+import com.correx.core.inference.InferenceResponse
+import com.correx.core.inference.ModelCapability
+import com.correx.core.inference.NoEligibleProviderException
+import com.correx.core.inference.ProviderHealth
+import com.correx.core.inference.Token
+import com.correx.core.inference.Tokenizer
+import com.correx.core.utils.TypeId
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+
+class CapabilityAwareRoutingStrategyTest {
+
+    private val strategy = CapabilityAwareRoutingStrategy()
+
+    private fun fakeProvider(id: String, vararg caps: Pair<ModelCapability, Double>): InferenceProvider =
+        object : InferenceProvider {
+            override val id: ProviderId = TypeId(id)
+            override val name: String = id
+            override val tokenizer: Tokenizer = object : Tokenizer {
+                override suspend fun tokenize(text: String): List<Token> = emptyList()
+                override suspend fun countTokens(text: String): Int = 0
+            }
+            override suspend fun infer(request: InferenceRequest): InferenceResponse = error("unused")
+            override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
+            override fun capabilities(): Set<CapabilityScore> =
+                caps.map { CapabilityScore(it.first, it.second) }.toSet()
+        }
+
+    @Test
+    fun `picks the highest-scoring provider for the required capability`() {
+        val weak = fakeProvider("weak", ModelCapability.ToolCalling to 0.4, ModelCapability.General to 1.0)
+        val strong = fakeProvider("strong", ModelCapability.ToolCalling to 0.95, ModelCapability.General to 1.0)
+        val selected = strategy.select(listOf(weak, strong), setOf(ModelCapability.ToolCalling))
+        assertEquals("strong", selected.id.value)
+    }
+
+    @Test
+    fun `still hard-filters to providers declaring the required capability`() {
+        val noTools = fakeProvider("noTools", ModelCapability.General to 1.0)
+        val tools = fakeProvider("tools", ModelCapability.ToolCalling to 0.3)
+        val selected = strategy.select(listOf(noTools, tools), setOf(ModelCapability.ToolCalling))
+        assertEquals("tools", selected.id.value)
+    }
+
+    @Test
+    fun `throws when no candidate declares the required capability`() {
+        val general = fakeProvider("general", ModelCapability.General to 1.0)
+        assertThrows<NoEligibleProviderException> {
+            strategy.select(listOf(general), setOf(ModelCapability.ToolCalling))
+        }
+    }
+
+    @Test
+    fun `keeps list order on a score tie (first-available parity)`() {
+        val a = fakeProvider("a", ModelCapability.ToolCalling to 0.5)
+        val b = fakeProvider("b", ModelCapability.ToolCalling to 0.5)
+        assertEquals("a", strategy.select(listOf(a, b), setOf(ModelCapability.ToolCalling)).id.value)
+    }
+
+    @Test
+    fun `with no required capabilities ranks by total declared score`() {
+        val lean = fakeProvider("lean", ModelCapability.General to 0.5)
+        val rich = fakeProvider("rich", ModelCapability.General to 1.0, ModelCapability.Coding to 0.9)
+        assertEquals("rich", strategy.select(listOf(lean, rich), emptySet()).id.value)
+    }
+}