From 82674e85c49c886550deaf0edd6b933ee07af074 Mon Sep 17 00:00:00 2001 From: kami Date: Mon, 29 Jun 2026 11:06:20 +0000 Subject: [PATCH] feat(inference): capability-aware routing for tool-heavy stages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../kotlin/com/correx/apps/server/Main.kt | 4 +- core/kernel/AGENTS.md | 1 + .../orchestration/SessionOrchestrator.kt | 8 +- infrastructure/inference/AGENTS.md | 2 +- .../CapabilityAwareRoutingStrategy.kt | 38 ++++++++++ .../CapabilityAwareRoutingStrategyTest.kt | 73 +++++++++++++++++++ 6 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 infrastructure/inference/src/main/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategy.kt create mode 100644 infrastructure/inference/src/test/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategyTest.kt 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 a2b27520..924e4b96 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 @@ -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" diff --git a/core/kernel/AGENTS.md b/core/kernel/AGENTS.md index 1e1de4f9..7bba11ce 100644 --- a/core/kernel/AGENTS.md +++ b/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. diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index 0f0c4ea0..e7ed7be5 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/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, diff --git a/infrastructure/inference/AGENTS.md b/infrastructure/inference/AGENTS.md index 2c3187c2..c1317836 100644 --- a/infrastructure/inference/AGENTS.md +++ b/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). diff --git a/infrastructure/inference/src/main/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategy.kt b/infrastructure/inference/src/main/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategy.kt new file mode 100644 index 00000000..86b20ff6 --- /dev/null +++ b/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, + requiredCapabilities: Set, + ): 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): 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 } + } +} diff --git a/infrastructure/inference/src/test/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategyTest.kt b/infrastructure/inference/src/test/kotlin/com/correx/infrastructure/inference/CapabilityAwareRoutingStrategyTest.kt new file mode 100644 index 00000000..b19d9d16 --- /dev/null +++ b/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): 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 = 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 = + 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 { + 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) + } +}