fix(router): preserve classifier breaker state under concurrency

litellm/router_strategy/complexity_router/complexity_router.py

@@ -836,31 +836,45 @@ class _ClassifierCircuitBreaker:
         self._clock = clock
         self._state = self.CLOSED
         self._opened_at: float | None = None
+        self._generation = 0
         self._lock = Lock()
 
-    def allow_request(self) -> bool:
-        """Allow ordinary calls while closed and exactly one probe after cooldown."""
+    def acquire_permit(self) -> int | None:
+        """Return a generation-scoped permit, or deny the call while the circuit is open.
+
+        Calls admitted together while closed share a generation. The first timeout advances it,
+        making every other in-flight completion stale so it cannot erase the new cooldown.
+        """
         with self._lock:
             if self._state == self.CLOSED:
-                return True
+                return self._generation
             if self._state == self.HALF_OPEN:
-                return False
+                return None
             opened_at: Final = self._opened_at
             if opened_at is not None and self._clock() - opened_at >= self._cooldown_seconds:
                 self._state = self.HALF_OPEN
-                return True
-            return False
+                return self._generation
+            return None
 
-    def record_success(self) -> None:
+    def record_success(self, permit: int) -> None:
+        """Close only when the current half-open recovery probe succeeds."""
         with self._lock:
+            if self._state != self.HALF_OPEN or permit != self._generation:
+                return
             self._state = self.CLOSED
             self._opened_at = None
 
-    def record_failure(self, *, is_timeout: bool) -> None:
+    def record_failure(self, permit: int, *, is_timeout: bool) -> None:
         """Open on a normal timeout, or reopen when the single recovery probe fails."""
         with self._lock:
-            if not is_timeout and self._state != self.HALF_OPEN:
+            if permit != self._generation:
                 return
+            if self._state == self.CLOSED:
+                if not is_timeout:
+                    return
+            elif self._state != self.HALF_OPEN:
+                return
+            self._generation += 1
             self._state = self.OPEN
             self._opened_at = self._clock()
 
@@ -1541,7 +1555,8 @@ class ComplexityRouter(CustomLogger):
         has. It is handed to the failure path so a classifier error does not re-run the scorer.
         """
         breaker: Final = self._classifier_circuit_breaker
-        if breaker is not None and not breaker.allow_request():
+        permit: Final = breaker.acquire_permit() if breaker is not None else None
+        if breaker is not None and permit is None:
             return self._classifier_failure_outcome(
                 "LLM classifier circuit is open",
                 prompt,
@@ -1551,8 +1566,8 @@ class ComplexityRouter(CustomLogger):
             )
         try:
             tier, classifier_cost = await self._classify_with_llm(prompt, system_prompt, request_kwargs, messages)
-            if breaker is not None:
-                breaker.record_success()
+            if breaker is not None and permit is not None:
+                breaker.record_success(permit)
             return ClassificationOutcome(
                 tier=tier,
                 score=None,
@@ -1560,9 +1575,13 @@ class ComplexityRouter(CustomLogger):
                 cause="llm_classifier",
                 classifier_cost=classifier_cost,
             )
+        except asyncio.CancelledError:
+            if breaker is not None and permit is not None:
+                breaker.record_failure(permit, is_timeout=False)
+            raise
         except Exception as e:  # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path
-            if breaker is not None:
-                breaker.record_failure(is_timeout=_is_classifier_timeout(e))
+            if breaker is not None and permit is not None:
+                breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e))
             return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored)
 
     def _classifier_failure_outcome(

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -2100,29 +2100,74 @@ class TestLLMClassifier:
         now = 100.0
         breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
 
-        assert breaker.allow_request() is True
-        breaker.record_failure(is_timeout=True)
-        assert breaker.allow_request() is False
+        initial_permit = breaker.acquire_permit()
+        assert initial_permit is not None
+        breaker.record_failure(initial_permit, is_timeout=True)
+        assert breaker.acquire_permit() is None
 
         now = 130.0
-        assert breaker.allow_request() is True
-        assert breaker.allow_request() is False
+        probe_permit = breaker.acquire_permit()
+        assert probe_permit is not None
+        assert breaker.acquire_permit() is None
 
-        breaker.record_success()
-        assert breaker.allow_request() is True
+        breaker.record_success(probe_permit)
+        assert breaker.acquire_permit() is not None
 
     def test_failed_classifier_probe_restarts_cooldown(self):
         now = 100.0
         breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
-        breaker.record_failure(is_timeout=True)
+        initial_permit = breaker.acquire_permit()
+        assert initial_permit is not None
+        breaker.record_failure(initial_permit, is_timeout=True)
 
         now = 130.0
-        assert breaker.allow_request() is True
-        breaker.record_failure(is_timeout=False)
-        assert breaker.allow_request() is False
+        probe_permit = breaker.acquire_permit()
+        assert probe_permit is not None
+        breaker.record_failure(probe_permit, is_timeout=False)
+        assert breaker.acquire_permit() is None
 
         now = 160.0
-        assert breaker.allow_request() is True
+        assert breaker.acquire_permit() is not None
+
+    def test_stale_success_cannot_close_circuit_opened_by_overlapping_timeout(self):
+        breaker = _ClassifierCircuitBreaker(30.0)
+        timeout_permit = breaker.acquire_permit()
+        stale_success_permit = breaker.acquire_permit()
+        assert timeout_permit is not None
+        assert stale_success_permit is not None
+
+        breaker.record_failure(timeout_permit, is_timeout=True)
+        breaker.record_success(stale_success_permit)
+
+        assert breaker.acquire_permit() is None
+
+    @pytest.mark.asyncio
+    async def test_cancelled_classifier_probe_restarts_cooldown(self, mock_router_instance, llm_classifier_config):
+        now = 100.0
+        mock_router_instance.acompletion = AsyncMock(
+            side_effect=[
+                TimeoutError("classifier timed out"),
+                asyncio.CancelledError(),
+                _llm_response('{"tier": "SIMPLE"}'),
+            ]
+        )
+        router = ComplexityRouter(
+            model_name="test-complexity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config=llm_classifier_config,
+        )
+        router._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
+
+        await router.aclassify("open the circuit")
+        now = 130.0
+        with pytest.raises(asyncio.CancelledError):
+            await router.aclassify("cancel the recovery probe")
+
+        outcome = await router.aclassify("stay in cooldown")
+
+        assert outcome.cause == "heuristic_scorer"
+        assert "classifier-circuit-open" in outcome.signals
+        assert mock_router_instance.acompletion.await_count == 2
 
     @pytest.mark.asyncio
     async def test_classifier_circuit_can_be_disabled(self, mock_router_instance, llm_classifier_config):
@@ -2146,8 +2191,10 @@ class TestLLMClassifier:
 
     def test_non_timeout_failure_does_not_open_closed_classifier_circuit(self):
         breaker = _ClassifierCircuitBreaker(30.0)
-        breaker.record_failure(is_timeout=False)
-        assert breaker.allow_request() is True
+        permit = breaker.acquire_permit()
+        assert permit is not None
+        breaker.record_failure(permit, is_timeout=False)
+        assert breaker.acquire_permit() is not None
 
     @pytest.mark.asyncio
     async def test_aclassify_classifier_cost_is_none_when_call_is_unpriced(