From c5289420a1b9bf785c31ccbbc0fc4a42d7f94114 Mon Sep 17 00:00:00 2001 From: kami Date: Tue, 21 Jul 2026 01:57:04 +0400 Subject: [PATCH] fix(inference): tolerate a briefly-absent provider instead of killing the session (#299) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A transient provider connection drop (e.g. OOM-killed llama.cpp mid-request) collapsed into a hard NoEligibleProviderException that escaped runInference's try/catch (the route() call sat outside it), propagated past the stage retry loop, and landed in ServerModule's session-level catch — failing the WHOLE session even though the failure was retryable and the provider recovered seconds later. - SessionOrchestrator.runInference: wrap router.route() in try/catch so routing failures become InferenceResult.Failed and flow through the normal retryable/backoff/exhaustion machinery instead of escaping. - DefaultInferenceRouter.route: distinguish "capability never configured on any provider" (fail fast, waiting can't help) from "configured but currently unhealthy" (bounded wait/backoff — default 3 attempts x 2s — re-checking health before declaring NoEligibleProvider terminal). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01GeyGFXczJb8RUWGBKmkm6G --- .../core/inference/DefaultInferenceRouter.kt | 36 +++++++++++- .../orchestration/SessionOrchestrator.kt | 10 +++- .../inference/DefaultInferenceRouterTest.kt | 56 +++++++++++++++++++ .../SessionOrchestratorIntegrationTest.kt | 37 ++++++++++++ 4 files changed, 137 insertions(+), 2 deletions(-) diff --git a/core/inference/src/main/kotlin/com/correx/core/inference/DefaultInferenceRouter.kt b/core/inference/src/main/kotlin/com/correx/core/inference/DefaultInferenceRouter.kt index 3fa819d1..b48dd369 100644 --- a/core/inference/src/main/kotlin/com/correx/core/inference/DefaultInferenceRouter.kt +++ b/core/inference/src/main/kotlin/com/correx/core/inference/DefaultInferenceRouter.kt @@ -2,6 +2,7 @@ package com.correx.core.inference import com.correx.core.events.types.ProviderId import com.correx.core.events.types.StageId +import kotlinx.coroutines.delay import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock import kotlin.time.Duration @@ -21,6 +22,12 @@ class DefaultInferenceRouter( private val strategy: RoutingStrategy, private val cacheTtl: Duration = 5.seconds, private val timeSource: TimeSource = TimeSource.Monotonic, + // A provider that briefly drops (crash + qa-stack restart) shouldn't collapse into a hard + // NoEligibleProvider abort — give it bounded time to come back before giving up. This only + // applies when the capability IS configured on some provider but that provider is currently + // unhealthy; a capability nobody ever declared fails immediately (see routeCapabilityCandidates). + private val unavailableRetryAttempts: Int = 3, + private val unavailableRetryDelay: Duration = 2.seconds, ) : InferenceRouter { private val cache = mutableMapOf() @@ -45,12 +52,39 @@ class DefaultInferenceRouter( } } + private suspend fun refreshedHealth(provider: InferenceProvider): ProviderHealth = + lockFor(provider.id).withLock { + val fresh = provider.healthCheck() + mapMutex.withLock { cache[provider.id] = HealthEntry(fresh, timeSource.markNow()) } + fresh + } + override suspend fun route(stageId: StageId, requiredCapabilities: Set): InferenceProvider { val candidates = requiredCapabilities .flatMap { registry.resolve(it) } .distinctBy { it.id } .ifEmpty { registry.listAll() } - val healthy = candidates.filter { cachedHealth(it) !is ProviderHealth.Unavailable } + + // Nobody was ever configured with this capability set — no amount of waiting fixes that, + // fail fast instead of burning the bounded-wait budget below. + if (requiredCapabilities.isNotEmpty() && + candidates.none { it.capabilities().map { c -> c.capability }.toSet().containsAll(requiredCapabilities) } + ) { + throw NoEligibleProviderException(stageId, requiredCapabilities) + } + + var healthy = candidates.filter { cachedHealth(it) !is ProviderHealth.Unavailable } + var attempt = 0 + while (healthy.isEmpty() && attempt < unavailableRetryAttempts) { + attempt++ + log.warn( + "route: capability {} configured but all candidates unhealthy for stage={}" + + " — waiting {} (attempt {}/{}) before declaring NoEligibleProvider", + requiredCapabilities, stageId.value, unavailableRetryDelay, attempt, unavailableRetryAttempts, + ) + delay(unavailableRetryDelay) + healthy = candidates.filter { refreshedHealth(it) !is ProviderHealth.Unavailable } + } val selected = strategy.select(healthy, requiredCapabilities) // Post-selection re-check closes the TOCTOU window between initial filter and dispatch. when (val postHealth = selected.healthCheck()) { 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 a928bf21..b98c6ac7 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 @@ -300,7 +300,15 @@ abstract class SessionOrchestrator( // ToolCalling score and routes the stage to the best tool-caller. val requiredCapabilities = stageConfig.requiredCapabilities + if (withTools && stageConfig.allowedTools.isNotEmpty()) setOf(ModelCapability.ToolCalling) else emptySet() - val provider = inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId) + // Routing itself can fail transiently (provider mid-crash-recovery) — it must be retryable + // like any other inference failure, not escape and kill the whole session (see #299). + val provider = try { + inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId) + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + return InferenceResult.Failed(e.message ?: "routing failed") + } log.debug( "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}", sessionId.value, stageId.value, provider.id.value, timeoutMs, diff --git a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/DefaultInferenceRouterTest.kt b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/DefaultInferenceRouterTest.kt index 5e405efb..d6284870 100644 --- a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/DefaultInferenceRouterTest.kt +++ b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/DefaultInferenceRouterTest.kt @@ -13,8 +13,10 @@ import com.correx.core.inference.RoutingStrategy import com.correx.testing.fixtures.inference.MockInferenceProvider import kotlinx.coroutines.runBlocking import org.junit.jupiter.api.Assertions.assertSame +import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import org.junit.jupiter.api.assertThrows +import kotlin.time.Duration.Companion.milliseconds class DefaultInferenceRouterTest { @@ -157,4 +159,58 @@ class DefaultInferenceRouterTest { val result = router.route(stage, setOf(ModelCapability.General), "llama-cpp:sick") assertSame(healthy, result) } + + // ── bounded wait for a briefly-absent provider (#299) ───────────────────── + + @Test + fun `waits for the sole capable provider to recover instead of failing immediately`(): Unit = runBlocking { + var healthChecks = 0 + val recovering = object : InferenceProvider by provider("a", ModelCapability.ToolCalling) { + override suspend fun healthCheck(): ProviderHealth { + healthChecks++ + return if (healthChecks < 3) ProviderHealth.Unavailable("connection dropped") else ProviderHealth.Healthy + } + } + val router = DefaultInferenceRouter( + registryOf(recovering), + firstStrategy(), + unavailableRetryAttempts = 5, + unavailableRetryDelay = 5.milliseconds, + ) + val result = router.route(stage, setOf(ModelCapability.ToolCalling)) + assertSame(recovering, result) + assertTrue(healthChecks >= 3) { "expected at least 3 health checks, got $healthChecks" } + } + + @Test + fun `still throws NoEligibleProviderException if the sole provider never recovers within the bound`() { + val neverRecovers = MockInferenceProvider( + id = ProviderId("a"), + declaredCapabilities = setOf(CapabilityScore(ModelCapability.ToolCalling, 1.0)), + health = ProviderHealth.Unavailable("still down"), + ) + val router = DefaultInferenceRouter( + registryOf(neverRecovers), + throwingStrategy(), + unavailableRetryAttempts = 2, + unavailableRetryDelay = 5.milliseconds, + ) + assertThrows { + runBlocking { router.route(stage, setOf(ModelCapability.ToolCalling)) } + } + } + + @Test + fun `fails fast without waiting when the capability was never configured on any provider`(): Unit = runBlocking { + val p = provider("a", ModelCapability.General) // does not declare ToolCalling + val router = DefaultInferenceRouter( + registryOf(p), + throwingStrategy(), + unavailableRetryAttempts = 5, + unavailableRetryDelay = 10_000.milliseconds, // would time the test out if the wait loop ran + ) + assertThrows { + router.route(stage, setOf(ModelCapability.ToolCalling)) + } + } } diff --git a/testing/integration/src/test/kotlin/SessionOrchestratorIntegrationTest.kt b/testing/integration/src/test/kotlin/SessionOrchestratorIntegrationTest.kt index 39f2454d..cd833972 100644 --- a/testing/integration/src/test/kotlin/SessionOrchestratorIntegrationTest.kt +++ b/testing/integration/src/test/kotlin/SessionOrchestratorIntegrationTest.kt @@ -46,6 +46,7 @@ import com.correx.core.inference.InferenceResponse import com.correx.core.inference.InferenceRouter import com.correx.core.inference.InferenceState 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.TokenUsage @@ -216,6 +217,42 @@ class SessionOrchestratorIntegrationTest { assertTrue(failed.retryExhausted) } + @Test + fun `transient NoEligibleProviderException on a retry is tolerated, not a session-killing crash (#299)`(): Unit = + runBlocking { + val sessionId = SessionId("s2b") + val config = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0), + ) + val graph = threeStageGraph() + var routeCalls = 0 + + val recoveringOrchestrator = DefaultSessionOrchestrator( + repositories = repositories, + engines = engines.copy( + inferenceRouter = object : InferenceRouter { + override suspend fun route(stageId: StageId, requiredCapabilities: Set): InferenceProvider { + routeCalls++ + // First attempt hits the provider mid-crash — subsequent retries find it back. + if (routeCalls == 1) throw NoEligibleProviderException(stageId, requiredCapabilities) + return MockInferenceProvider() + } + }, + ), + retryCoordinator = retryCoordinator, + artifactStore = artifactStore, + decisionJournalRepository = decisionJournalRepository, + ) + recoveringOrchestrator.run(sessionId, graph, config) + + val events = eventStore.read(sessionId) + // A crash would have skipped straight to an unhandled failure with no WorkflowCompletedEvent + // and no orderly retry — assert the session survived and made it through via the retry path. + assertTrue(routeCalls > 1) { "expected routing to be retried, only called $routeCalls time(s)" } + assertNotNull(events.find { it.payload is WorkflowCompletedEvent }) + assertTrue(events.none { it.payload is WorkflowFailedEvent }) + } + @Test fun `stage that matches no transition retries through the budget before failing`(): Unit = runBlocking { val sessionId = SessionId("s-no-transition-retry")