fix(middleware): lossless limiter handoff, generation-aware cap, burst-rate CB gate

Address the open P1/P2 review findings on #4294:

- P1 #1 (cancellation handoff): reserve the permit for a specific waiter at
  dequeue time (grant-at-dequeue, _AsyncWaiter.granted) so a waiter cancelled
  in the post-dequeue / pre-reacquire window hands its reservation to the next
  waiter (_handoff_granted_permit_locked) instead of stranding it. No
  cancellation window remains.
- P1 #2 (hot-reload generation): move cap updates out of the per-attempt path;
  give the limiter one generation-aware owner (set_limit_if_newer with a
  monotonic instance seq proxy for config freshness) so a stale in-flight run
  cannot rewrite a freshly-lowered cap. max_concurrent_calls is now genuinely
  hot-reloadable, resolving the reload-boundary inconsistency by option (b) -
  no STARTUP_ONLY_FIELDS change (retry params truly hot-reload).
- P2 (circuit breaker): gate _record_failure on reason != "burst_rate" so
  burst-rate (limit_burst_rate) exhaustion - a transient slope-throttle, not
  "provider down" - does not trip the CB and fast-fail ALL calls.
- P3: clamp the jitter window to the cap before drawing (uniform spread
  instead of piling at the cap); document the per-process / GATEWAY_WORKERS
  cap semantics in config + the field description.

Tests: add the reviewer-requested regressions (cancel-after-dequeue handoff;
stale-instance-doesn't-overwrite-lowered-cap across sync + isolated-loop async;
burst_rate-exhaustion-doesn't-trip-CB sync + async). Each is red on the prior
buggy logic and green on the fix. _build_middleware now routes llm_call knobs
through AppConfig so __init__ applies the cap. 71 middleware tests pass; 212
across the blast radius (1 pre-existing skip).

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Willem Jiang
2026-07-19 17:06:51 +08:00
co-authored by Claude
parent e0161aaff5
commit acfc761744
4 changed files with 567 additions and 116 deletions
@@ -132,6 +132,41 @@ _STREAM_DROP_EXCEPTIONS: frozenset[str] = frozenset(
# the *slope* of the request rate, so the cap must bound aggregate in-flight
# calls process-wide - a per-loop cap (which is what asyncio.Semaphore would
# give) is defeated the moment subagent fan-out runs on a second loop.
#
# Correctness invariants the design below preserves:
# * P1 #1 (lossless waiter handoff): a permit handed to a waiter is *reserved*
# for that waiter at dequeue time (``granted=True``). If the waiter is
# cancelled before it wakes, the reserved permit is re-handed to the next
# waiter (or freed) - so a cancellation in the post-dequeue/pre-reacquire
# window never strands the next waiter with capacity idle.
# * P1 #2 (generation-aware cap): the cap is set by a single owner
# (``_apply_configured_cap`` from each middleware ``__init__``) carrying a
# monotonic instance seq (a proxy for config freshness). Per-attempt
# callers only acquire/release - they never rewrite the cap - so a stale
# in-flight run cannot bump a freshly-lowered cap back up.
class _AsyncWaiter:
"""A parked async caller awaiting a transferred permit.
``granted`` is flipped to ``True`` (under the limiter lock) at the exact
moment a permit is reserved for this waiter - by ``release`` handing off a
returning permit, by ``set_limit_if_newer`` admitting queued callers on a
cap raise, or by another cancelling waiter handing off its reserved permit.
The reservation is atomic with the dequeue, so the invariant
``granted is True <=> not in _async_waiters`` always holds: once granted,
the permit is already counted in ``_in_flight`` and the waiter need only
wake and return. A cancelled waiter therefore knows from ``granted``
whether it owes a handoff (granted) or is merely unregistering (not yet
granted).
"""
__slots__ = ("loop", "event", "granted")
def __init__(self, loop: asyncio.AbstractEventLoop, event: asyncio.Event) -> None:
self.loop = loop
self.event = event
self.granted = False
class _ProcessWideLimiter:
@@ -143,42 +178,56 @@ class _ProcessWideLimiter:
limiter is built on ``threading`` primitives (not loop-bound): every call
path shares one in-flight counter and one cap.
The cap is mutable via ``set_limit``: changing it updates the cap in place
- the limiter is never recreated, so in-flight permits are never abandoned
(an old limiter is never orphaned holding permits a caller will later
release). A lower cap takes effect for the next acquire; in-flight callers
keep their permit until they release. Permits are released in a ``finally``
and async waiters unregister on cancellation, so a cancelled caller never
leaks capacity.
The cap has a single generation-aware owner: ``set_limit_if_newer`` accepts
an ``owner_seq`` (a monotonic proxy for config freshness - see
``_next_instance_seq``) and ignores writes from an older instance, so a
stale in-flight run cannot rewrite a cap set by a newer config. Per-attempt
callers (``acquire_sync``/``acquire_async``/``release``) never touch the
cap. A lower cap takes effect for subsequent acquires; in-flight permits are
preserved (holders keep theirs until release). Permits are released in a
``finally`` and an async waiter that is cancelled after its permit was
reserved hands the reservation to the next waiter, so capacity never leaks
and a cancellation never strands a later waiter.
"""
def __init__(self, limit: int) -> None:
def __init__(self, limit: int, owner_seq: int) -> None:
self._lock = threading.Lock()
self._cond = threading.Condition(self._lock)
self._in_flight = 0
self._limit = max(0, limit)
# FIFO of (owner_loop, event) for async callers waiting on capacity.
# Each event lives on its caller's loop; release() wakes one across
# loops via call_soon_threadsafe so the wakeup runs on the right loop.
self._async_waiters: deque[tuple[asyncio.AbstractEventLoop, asyncio.Event]] = deque()
self._owner_seq = owner_seq
# FIFO of async callers waiting on capacity. Each waiter lives on its
# caller's loop; release/handoff wakes one across loops via
# call_soon_threadsafe so the wakeup runs on the right loop.
self._async_waiters: deque[_AsyncWaiter] = deque()
@property
def limit(self) -> int:
return self._limit
def set_limit(self, limit: int) -> None:
"""Update the cap in place; never recreates the limiter.
@property
def in_flight(self) -> int:
return self._in_flight
A higher cap may now admit waiting callers, so wake them. A lower cap
takes effect for subsequent acquires; in-flight permits are preserved.
No-op when the limit is unchanged.
def set_limit_if_newer(self, limit: int, owner_seq: int) -> None:
"""Update the cap iff *owner_seq* is at least as new as the current owner.
Generation-aware ownership: a stale middleware instance (older config)
captured an older cap; its ``__init__`` still calls this, but its
``owner_seq`` is lower than the current owner's, so the write is ignored
and a freshly-lowered cap survives. A raise admits queued callers; a
lower takes effect for subsequent acquires while in-flight permits are
preserved. No-op when the limit is unchanged.
"""
new_limit = max(0, limit)
with self._cond:
if owner_seq < self._owner_seq:
return # stale owner: cannot rewrite the active cap
self._owner_seq = owner_seq
if new_limit == self._limit:
return
self._limit = new_limit
self._wake_waiters_locked(wake_all=True)
self._grant_to_queued_locked()
def acquire_sync(self) -> None:
"""Block the calling thread until a permit is available, then take one."""
@@ -187,42 +236,67 @@ class _ProcessWideLimiter:
self._cond.wait()
def release(self) -> None:
"""Return one permit and wake a single waiter (async preferred)."""
"""Return one permit, handing it to a waiter if one is queued.
If an async waiter is queued, the returning permit *transfers* to it
(ownership moves; ``_in_flight`` is unchanged) and its event is set so
it wakes already owning a permit. Otherwise the permit returns to the
free pool (``_in_flight -= 1``) and one sync waiter is notified to grab
it on its next ``_try_acquire_locked`` re-check.
"""
with self._cond:
if self._async_waiters:
waiter = self._async_waiters.popleft()
waiter.granted = True
if not self._wake_locked(waiter):
# Owner loop closed: the transferred permit is stranded;
# hand it to the next waiter or free it.
self._handoff_granted_permit_locked()
return
if self._in_flight > 0:
self._in_flight -= 1
self._wake_waiters_locked(wake_all=False)
self._cond.notify()
async def acquire_async(self) -> None:
"""Acquire a permit without blocking the event loop.
Capacity available -> returns immediately. Otherwise registers an
``asyncio.Event`` on the caller's loop and awaits it; ``release`` wakes
one waiter across loops via ``call_soon_threadsafe``. On cancellation
the registration is removed, so no later wakeup leaks a permit to a
dead coroutine - and no permit was taken while waiting, so there is
nothing to release.
Free capacity -> take one immediately. Otherwise park on an
``asyncio.Event``; ``release`` / a cap-raise transfers a permit to us
(``granted=True``) and sets the event. On cancellation, if a permit was
already reserved for us, hand it to the next waiter (or free it) so the
reservation is never lost; if we were still queued (not yet granted),
just unregister - no permit was reserved for us, so there is nothing to
release.
"""
loop = asyncio.get_running_loop()
while True:
event = asyncio.Event()
waiter = _AsyncWaiter(loop=loop, event=asyncio.Event())
with self._cond:
if self._try_acquire_locked():
return
self._async_waiters.append((loop, event))
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
"LLM call parking on process-wide limiter (in_flight=%d, limit=%d, queued=%d)",
self._in_flight,
self._limit,
len(self._async_waiters) + 1,
)
self._async_waiters.append(waiter)
try:
await event.wait()
await waiter.event.wait()
except asyncio.CancelledError:
with self._cond:
try:
self._async_waiters.remove((loop, event))
except ValueError:
# Already popped by a release(); that wakeup did not
# reserve a permit for us (we only take a permit inside
# _try_acquire_locked, which we hadn't re-run), so
# there is nothing to release.
pass
if waiter.granted:
# A permit was reserved for us but we're cancelling
# before waking. Pass the reservation to the next
# waiter (or free it) so it is not stranded.
self._handoff_granted_permit_locked()
else:
# Still queued, never granted (granted is set only when
# dequeued, under the lock): just unregister.
self._async_waiters.remove(waiter)
raise
return # woken => granted => we own a permit (already in _in_flight)
def _try_acquire_locked(self) -> bool:
if self._in_flight < self._limit:
@@ -230,49 +304,107 @@ class _ProcessWideLimiter:
return True
return False
def _wake_waiters_locked(self, *, wake_all: bool) -> None:
# Wake async waiters first (one per release, or all on a cap change);
# remaining wakeups go to sync waiters via the condition. The woken
# async waiter re-checks capacity on its own loop and either takes the
# permit or re-registers, so waking is always safe (never over-grants).
while self._async_waiters:
loop, event = self._async_waiters.popleft()
try:
loop.call_soon_threadsafe(event.set)
except RuntimeError:
continue # owner loop closed; drop the dead waiter
if not wake_all:
return
if wake_all:
def _grant_to_queued_locked(self) -> None:
"""Admit queued callers to newly-available capacity (cap raise).
Reserves a permit (counts it in ``_in_flight``) for each queued async
waiter that fits under the current cap, then notifies sync waiters for
any remaining free capacity so they re-check ``_try_acquire_locked``.
"""
while self._in_flight < self._limit and self._async_waiters:
waiter = self._async_waiters.popleft()
waiter.granted = True
self._in_flight += 1 # commit the reservation to this waiter
if not self._wake_locked(waiter):
# Owner loop closed: reclaim the reservation and try the next.
self._in_flight -= 1
if self._in_flight < self._limit:
self._cond.notify_all()
else:
self._cond.notify()
def _handoff_granted_permit_locked(self) -> None:
"""Transfer an already-reserved permit to the next queued waiter, or free it.
Used when a waiter that had a permit reserved cancels before waking, or
when a reservation target's loop is dead. The permit is already counted
in ``_in_flight``; transferring keeps it counted (ownership moves to the
next waiter), freeing returns it to the pool. Either way ``_in_flight``
stays correct and the reservation is never lost.
"""
while self._async_waiters:
waiter = self._async_waiters.popleft()
waiter.granted = True
if self._wake_locked(waiter):
return # ownership transferred; _in_flight unchanged
# dead loop; try the next waiter
# No async waiter to take it: free the permit and wake a sync waiter.
if self._in_flight > 0:
self._in_flight -= 1
self._cond.notify()
def _wake_locked(self, waiter: _AsyncWaiter) -> bool:
"""Schedule ``event.set`` on the waiter's loop. False if the loop is dead."""
try:
waiter.loop.call_soon_threadsafe(waiter.event.set)
return True
except RuntimeError:
return False # owner loop closed: the wakeup cannot land
_LIMITER_LOCK = threading.Lock()
_PROCESS_LIMITER: _ProcessWideLimiter | None = None
# Monotonic counter of middleware construction order - a proxy for config
# generation: ``get_app_config()`` reloads on file mtime change, so a later
# ``LLMErrorHandlingMiddleware`` always captured config at least as fresh as an
# earlier one. The cap owner (``_apply_configured_cap``) passes this seq to
# ``set_limit_if_newer`` so a stale instance cannot overwrite a cap set by a
# newer config (P1 #2).
_INSTANCE_SEQ_LOCK = threading.Lock()
_INSTANCE_SEQ: int = 0
def _get_process_limiter(limit: int) -> _ProcessWideLimiter | None:
"""Return the process-wide LLM-call limiter, or ``None`` when disabled.
Created once (double-checked under a lock so concurrent first-calls don't
create two). A subsequent change to ``limit`` updates the cap in place via
``set_limit`` - the limiter is never replaced, so in-flight permits are
never abandoned. ``limit <= 0`` disables the cap entirely (passthrough).
def _next_instance_seq() -> int:
global _INSTANCE_SEQ
with _INSTANCE_SEQ_LOCK:
_INSTANCE_SEQ += 1
return _INSTANCE_SEQ
def _get_process_limiter() -> _ProcessWideLimiter | None:
"""Return the process-wide LLM-call limiter, or ``None`` until first configured.
Per-attempt callers use this to acquire/release only - it never changes the
cap. The cap is owned solely by ``_apply_configured_cap`` (called once per
middleware ``__init__``), which makes updates generation-aware so a stale
in-flight run cannot rewrite the active cap.
"""
return _PROCESS_LIMITER
def _apply_configured_cap(limit: int, owner_seq: int) -> None:
"""Set the process-wide limiter cap from the latest config (single owner).
Called once per middleware ``__init__`` (which captures the live ``llm_call``
config). Creates the limiter on first use, then updates the cap in place via
``set_limit_if_newer`` - a stale instance (older ``owner_seq``) cannot
overwrite a cap set by a newer one. ``limit <= 0`` disables the cap: a prior
positive cap keeps serving its in-flight callers (which release normally),
while new callers see ``max_concurrent_llm_calls <= 0`` and skip the limiter
entirely.
"""
global _PROCESS_LIMITER
if limit <= 0:
return None
return # disabled: callers short-circuit before consulting the limiter
limiter = _PROCESS_LIMITER
if limiter is None:
with _LIMITER_LOCK:
if _PROCESS_LIMITER is None:
_PROCESS_LIMITER = _ProcessWideLimiter(limit)
_PROCESS_LIMITER = _ProcessWideLimiter(limit, owner_seq)
limiter = _PROCESS_LIMITER
elif limiter.limit != limit:
limiter.set_limit(limit)
return limiter
# Apply our cap. No-op if we just created the limiter with this exact cap;
# ignored by the generation guard if a newer owner already set a different
# cap (e.g. a racing thread created it with a fresher seq).
limiter.set_limit_if_newer(limit, owner_seq)
class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
@@ -293,6 +425,12 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
def __init__(self, *, app_config: AppConfig, **kwargs: Any) -> None:
super().__init__(**kwargs)
# Monotonic construction-order id - a proxy for config generation (a
# later instance captured config at least as fresh as an earlier one).
# Passed to the process limiter so cap updates are generation-aware and
# a stale in-flight instance cannot rewrite a cap set by a newer config.
self._instance_seq = _next_instance_seq()
self.circuit_failure_threshold = app_config.circuit_breaker.failure_threshold
self.circuit_recovery_timeout_sec = app_config.circuit_breaker.recovery_timeout_sec
@@ -306,6 +444,12 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
self.burst_retry_base_delay_ms = llm_call.burst_retry_base_delay_ms
self.max_concurrent_llm_calls = llm_call.max_concurrent_calls
# Apply the configured cap to the process-wide limiter. This is the
# SINGLE owner of the cap (per-attempt callers only acquire/release);
# the instance seq makes it generation-aware so a stale in-flight run
# cannot overwrite a cap set by a newer config (P1 #2).
_apply_configured_cap(self.max_concurrent_llm_calls, self._instance_seq)
# Circuit Breaker state
self._circuit_lock = threading.Lock()
self._circuit_failure_count = 0
@@ -453,8 +597,8 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
(cap disabled) is a direct passthrough. Permits release on any exit
(return or raise) via ``finally`` so a raised handler never leaks a slot.
"""
limiter = _get_process_limiter(self.max_concurrent_llm_calls)
if limiter is None:
limiter = _get_process_limiter()
if limiter is None or self.max_concurrent_llm_calls <= 0:
return handler(request)
limiter.acquire_sync()
try:
@@ -476,8 +620,8 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
cancellation) via ``finally``; ``acquire_async`` separately cleans up if
cancelled while waiting, so capacity never leaks.
"""
limiter = _get_process_limiter(self.max_concurrent_llm_calls)
if limiter is None:
limiter = _get_process_limiter()
if limiter is None or self.max_concurrent_llm_calls <= 0:
return await handler(request)
await limiter.acquire_async()
try:
@@ -492,9 +636,11 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
jitter) - the server told us exactly when to come back, and for a
burst-rate 429 this is strongly preferred over any computed delay.
Otherwise AWS-style "decorrelated jitter" is applied:
``delay = min(cap, random(base, seed * 3))`` where ``seed`` is the
previous delay, or the reason-specific base on the first retry
(``prev_delay_ms is None``). ``reason="burst_rate"`` swaps in
``delay = random(base, min(cap, max(base, seed * 3)))`` where ``seed``
is the previous delay, or the reason-specific base on the first retry
(``prev_delay_ms is None``). The window is clamped to the cap *before*
drawing (not after) so the distribution stays uniform up to the cap
rather than piling up at it. ``reason="burst_rate"`` swaps in
``burst_retry_base_delay_ms`` (longer than the normal base) so the
single burst retry lands after the throttle window subsides.
@@ -503,8 +649,9 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
non-degenerate: with the normal base (1000ms) the burst window would
collapse to ``randint(5000, max(5000, 1000*3)) = randint(5000, 5000)``
and every concurrent burst failure would realign on the same 5s tick.
Seeding from 5000ms gives ``randint(5000, 15000)`` (capped at
``retry_cap_delay_ms``), so a fleet that failed together spreads out.
Seeding from 5000ms gives ``randint(5000, min(8000, 15000)) =
randint(5000, 8000)`` with defaults, so a fleet that failed together
spreads out across the whole window.
Deterministic exponential backoff (``base * 2^(attempt-1)``) makes
every concurrent retryer realign on the same backoff ticks; when a
@@ -519,9 +666,15 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
base = self.burst_retry_base_delay_ms if reason == "burst_rate" else self.retry_base_delay_ms
cap = self.retry_cap_delay_ms
seed = base if prev_delay_ms is None else prev_delay_ms
high = max(base, seed * 3)
delay = random.randint(base, high)
return min(delay, cap)
# Clamp the window to the cap *before* drawing so the jitter spreads
# uniformly across [base, min(cap, seed*3)] instead of concentrating at
# the cap: with defaults seed*3 (=15000) >> cap (=8000), drawing
# randint(base, seed*3) then min(delay, cap) would put ~70% of draws at
# exactly cap, re-clustering a fleet that the jitter is meant to spread.
high = min(cap, max(base, seed * 3))
if high < base:
return cap # base exceeds cap (misconfiguration): the cap wins
return random.randint(base, high)
def _build_retry_message(self, attempt: int, wait_ms: int, reason: str) -> str:
seconds = max(1, round(wait_ms / 1000))
@@ -652,10 +805,14 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
_extract_error_detail(exc),
exc_info=exc,
)
if retriable:
if retriable and reason != "burst_rate":
self._record_failure()
else:
# Non-retriable: release the probe without recording a failure.
# Non-retriable, OR burst_rate (a transient provider
# slope-throttle, not "provider down"): release the half-open
# probe without recording a failure so the circuit doesn't
# trip and fast-fail ALL calls for the recovery window - the
# exact self-inflicted outage #4290 is trying to prevent.
self._release_half_open_probe()
return self._build_user_fallback_message(exc, reason)
@@ -707,10 +864,14 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
_extract_error_detail(exc),
exc_info=exc,
)
if retriable:
if retriable and reason != "burst_rate":
self._record_failure()
else:
# Non-retriable: release the probe without recording a failure.
# Non-retriable, OR burst_rate (a transient provider
# slope-throttle, not "provider down"): release the half-open
# probe without recording a failure so the circuit doesn't
# trip and fast-fail ALL calls for the recovery window - the
# exact self-inflicted outage #4290 is trying to prevent.
self._release_half_open_probe()
return self._build_user_fallback_message(exc, reason)
@@ -81,7 +81,12 @@ class LlmCallConfig(BaseModel):
"Process-wide cap on concurrently in-flight LLM calls. 0 disables "
"the cap (default, preserving existing behavior). Set to a positive "
"int to smooth provider burst-rate (limit_burst_rate) spikes by "
"bounding the request-rate slope at the morning peak."
"bounding the request-rate slope at the morning peak. Per-process, "
"not per-cluster: with GATEWAY_WORKERS > 1 the aggregate cap is "
"effectively max_concurrent_calls * GATEWAY_WORKERS (and a "
"multi-node rollout multiplies it further), so size the per-process "
"value accordingly and pair it with an nginx limit_req at the ingress "
"for a true cluster-wide slope cap."
),
)
retry_max_attempts: int = Field(
@@ -21,17 +21,21 @@ from deerflow.config.sandbox_config import SandboxConfig
@pytest.fixture(autouse=True)
def _reset_process_limiter() -> Iterator[None]:
"""Reset the module-global limiter between tests.
"""Reset the module-global limiter + instance seq between tests.
The limiter is a process singleton shared across all middleware instances,
so cap / in-flight state from one test would otherwise bleed into the next
(notably the limit-change test, which asserts the singleton is reused).
(notably the limit-change test, which asserts the singleton is reused). The
instance seq is reset too so the generation ordering in the stale-instance
regression tests is deterministic within each test.
"""
from deerflow.agents.middlewares import llm_error_handling_middleware as mod
mod._PROCESS_LIMITER = None
mod._INSTANCE_SEQ = 0
yield
mod._PROCESS_LIMITER = None
mod._INSTANCE_SEQ = 0
def _make_app_config() -> AppConfig:
@@ -56,10 +60,30 @@ class FakeError(Exception):
self.response = SimpleNamespace(status_code=status_code, headers=headers or {}) if status_code is not None or headers else None
# Middleware-level attribute -> ``LlmCallConfig`` field. ``llm_call`` knobs are
# routed through ``AppConfig`` so ``__init__`` applies the cap to the process
# limiter (the single generation-aware owner). Circuit-breaker knobs are read
# per-call from ``self`` (not via the limiter), so setattr after ``__init__``
# still works for them.
_LLM_CALL_ATTR_MAP: dict[str, str] = {
"max_concurrent_llm_calls": "max_concurrent_calls",
"retry_max_attempts": "retry_max_attempts",
"retry_base_delay_ms": "retry_base_delay_ms",
"retry_cap_delay_ms": "retry_cap_delay_ms",
"burst_retry_base_delay_ms": "burst_retry_base_delay_ms",
}
def _build_middleware(**attrs: int) -> LLMErrorHandlingMiddleware:
middleware = LLMErrorHandlingMiddleware(app_config=_make_app_config())
llm_call_fields = {_LLM_CALL_ATTR_MAP[key]: value for key, value in attrs.items() if key in _LLM_CALL_ATTR_MAP}
app_config = AppConfig(
sandbox=SandboxConfig(use="test"),
llm_call=LlmCallConfig(**llm_call_fields),
)
middleware = LLMErrorHandlingMiddleware(app_config=app_config)
for key, value in attrs.items():
setattr(middleware, key, value)
if key not in _LLM_CALL_ATTR_MAP:
setattr(middleware, key, value)
return middleware
@@ -1117,7 +1141,8 @@ def test_max_concurrent_llm_calls_defaults_to_disabled() -> None:
def test_max_concurrent_llm_calls_wired_from_config() -> None:
"""llm_call.max_concurrent_calls flows through AppConfig into the middleware
attribute that _get_process_limiter reads.
attribute that ``__init__`` applies to the process limiter via
``_apply_configured_cap`` (the single generation-aware cap owner).
"""
app_config = AppConfig(
sandbox=SandboxConfig(use="test"),
@@ -1291,6 +1316,65 @@ async def test_async_burst_rate_uses_tight_budget_and_longer_base(
assert result.additional_kwargs.get("error_reason") == "burst_rate"
def test_burst_rate_exhaustion_does_not_trip_circuit_breaker(monkeypatch: pytest.MonkeyPatch) -> None:
"""P2: burst-rate (limit_burst_rate) is a transient provider slope-throttle,
not "provider down". Exhausting its retry budget must NOT count toward the
circuit breaker - otherwise N consecutive burst failures flip the CB open
and fast-fail ALL calls for the recovery window, the exact self-inflicted
outage #4290 is trying to prevent. The burst reason is already distinctively
classified by this PR, so it is the natural place to exclude it.
"""
monkeypatch.setattr("time.sleep", lambda _d: None)
middleware = _build_middleware(circuit_failure_threshold=3, retry_max_attempts=3)
def handler(_request) -> AIMessage:
raise FakeError("rate increased too quickly", status_code=429, code="limit_burst_rate")
# Exceed the failure threshold with burst-rate failures: the CB must stay
# closed (burst-rate is transient-by-design, not a provider outage).
for _ in range(5):
result = middleware.wrap_model_call(SimpleNamespace(), handler)
assert result.additional_kwargs.get("error_reason") == "burst_rate"
assert middleware._circuit_failure_count == 0
assert middleware._circuit_state == "closed"
assert middleware._check_circuit() is False # still admitting calls
# A subsequent successful call still goes through (CB never opened).
def ok_handler(_request) -> AIMessage:
return AIMessage(content="ok")
result = middleware.wrap_model_call(SimpleNamespace(), ok_handler)
assert result.content == "ok"
@pytest.mark.anyio
async def test_async_burst_rate_exhaustion_does_not_trip_circuit_breaker(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Async mirror: burst-rate exhaustion on the async path also stays out of
the circuit breaker (the gate lives in both ``wrap_model_call`` and
``awrap_model_call``).
"""
async def _noop_sleep(_d: float) -> None:
return None
monkeypatch.setattr(asyncio, "sleep", _noop_sleep)
middleware = _build_middleware(circuit_failure_threshold=3, retry_max_attempts=3)
async def handler(_request) -> AIMessage:
raise FakeError("rate increased too quickly", status_code=429, code="limit_burst_rate")
for _ in range(5):
result = await middleware.awrap_model_call(SimpleNamespace(), handler)
assert result.additional_kwargs.get("error_reason") == "burst_rate"
assert middleware._circuit_failure_count == 0
assert middleware._circuit_state == "closed"
assert middleware._check_circuit() is False
# ---------- Retry params config wiring ----------
@@ -1455,17 +1539,18 @@ def test_limiter_caps_concurrent_sync_calls() -> None:
assert len(results) == 2
def test_limit_change_does_not_recreate_limiter_or_abandon_permits() -> None:
"""Changing the configured cap updates the limiter in place - same instance,
in-flight permit preserved (not abandoned). In-flight never exceeds the
effective cap during the transition. Regression for the recreate-on-change
permit-abandonment the reviewer found.
def test_cap_raise_via_newer_instance_updates_in_place_and_preserves_in_flight() -> None:
"""A newer middleware instance (higher instance seq = newer config) raises
the cap in place - same limiter instance, in-flight permits preserved (not
abandoned). Replaces the old recreate-on-change behavior with
generation-aware in-place updates.
"""
from deerflow.agents.middlewares.llm_error_handling_middleware import _get_process_limiter
middleware = _build_middleware(max_concurrent_llm_calls=1)
limiter_at_cap1 = _get_process_limiter(1)
assert limiter_at_cap1 is not None
old_mw = _build_middleware(max_concurrent_llm_calls=1)
limiter = _get_process_limiter()
assert limiter is not None
assert limiter.limit == 1
gate = threading.Event()
in_flight = 0
@@ -1482,32 +1567,148 @@ def test_limit_change_does_not_recreate_limiter_or_abandon_permits() -> None:
in_flight -= 1
return AIMessage(content="ok")
holder = threading.Thread(target=lambda: middleware.wrap_model_call(SimpleNamespace(), handler))
holder = threading.Thread(target=lambda: old_mw.wrap_model_call(SimpleNamespace(), handler))
holder.start()
# Wait until the holder has taken the single permit (cap=1).
for _ in range(400):
try:
# Wait until the holder has taken the single permit (cap=1).
for _ in range(400):
with lock:
if max_in_flight >= 1:
break
time.sleep(0.005)
with lock:
if max_in_flight >= 1:
break
time.sleep(0.005)
with lock:
assert max_in_flight == 1
assert max_in_flight == 1
assert limiter.in_flight == 1 # holder holds the single permit
# Simulate a config reload raising the cap: the next limiter lookup with a
# new limit updates the cap in place (a fresh middleware instance would do
# the same on its first model call).
limiter_at_cap3 = _get_process_limiter(3)
# Same instance, cap updated in place - NOT recreated.
assert limiter_at_cap3 is limiter_at_cap1
assert limiter_at_cap3.limit == 3
# The holder's permit is still counted against the new cap (not abandoned).
assert limiter_at_cap3._in_flight == 1
# Release the holder; it returns its permit to the SAME limiter.
gate.set()
holder.join(timeout=5)
# A newer instance raises the cap to 3: same limiter, cap updated in place.
_build_middleware(max_concurrent_llm_calls=3)
assert _get_process_limiter() is limiter # same instance, not recreated
assert limiter.limit == 3
assert limiter.in_flight == 1 # holder's permit preserved, not abandoned
finally:
# Always release the holder so a failing assertion doesn't leak the thread.
gate.set()
holder.join(timeout=5)
assert not holder.is_alive()
assert limiter_at_cap3._in_flight == 0 # permit returned, not leaked
assert limiter.in_flight == 0 # permit returned, not leaked
def test_stale_instance_does_not_overwrite_lowered_cap() -> None:
"""P1 #2 regression: a stale middleware instance (older config, captured
cap 3) making calls after a newer instance lowered the cap to 1 must NOT
rewrite it back to 3, and aggregate in-flight calls must never exceed 1.
The per-attempt path only acquires/releases - it never rewrites the cap -
so a stale in-flight run cannot bump a freshly-lowered cap back up.
"""
from deerflow.agents.middlewares.llm_error_handling_middleware import _get_process_limiter
# Old instance created first (lower instance seq = older config generation).
old_mw = _build_middleware(max_concurrent_llm_calls=3)
# New instance created later (higher seq) lowers the cap to 1.
_build_middleware(max_concurrent_llm_calls=1)
limiter = _get_process_limiter()
assert limiter is not None
assert limiter.limit == 1 # newer instance won
gate = threading.Event()
in_flight = 0
max_in_flight = 0
lock = threading.Lock()
def handler(_request) -> AIMessage:
nonlocal in_flight, max_in_flight
with lock:
in_flight += 1
max_in_flight = max(max_in_flight, in_flight)
gate.wait()
with lock:
in_flight -= 1
return AIMessage(content="ok")
# The OLD instance (captured cap 3) makes overlapping calls; the live cap
# is 1, so only one may be in-flight at a time.
threads = [threading.Thread(target=lambda: old_mw.wrap_model_call(SimpleNamespace(), handler)) for _ in range(3)]
for t in threads:
t.start()
try:
for _ in range(400):
with lock:
if max_in_flight >= 1:
break
time.sleep(0.005)
# Give the queued calls a chance to (incorrectly) exceed the lowered cap.
time.sleep(0.05)
with lock:
assert max_in_flight == 1 # never exceeded the lowered cap
assert limiter.limit == 1 # stale instance did NOT rewrite it back to 3
finally:
# Always release the holder threads so a failing assertion doesn't hang.
gate.set()
for t in threads:
t.join(timeout=5)
assert all(not t.is_alive() for t in threads)
assert limiter.in_flight == 0
@pytest.mark.anyio
async def test_stale_instance_respects_lowered_cap_across_isolated_loop() -> None:
"""P1 #2 regression (cross-loop): a stale instance (captured cap 3) making
async calls on an isolated event loop after a newer instance lowered the cap
to 1 respects the lowered cap across loops - it cannot rewrite it back to 3,
and aggregate in-flight calls never exceed 1 across the lead loop and the
isolated subagent loop.
"""
from deerflow.agents.middlewares.llm_error_handling_middleware import _get_process_limiter
old_mw = _build_middleware(max_concurrent_llm_calls=3)
_build_middleware(max_concurrent_llm_calls=1) # newer instance lowers cap
limiter = _get_process_limiter()
assert limiter is not None
assert limiter.limit == 1 # newer instance lowered it; stale one didn't rewrite
gate = threading.Event()
counter_lock = threading.Lock()
in_flight = 0
max_in_flight = 0
async def handler(_request) -> AIMessage:
nonlocal in_flight, max_in_flight
with counter_lock:
in_flight += 1
max_in_flight = max(max_in_flight, in_flight)
try:
await asyncio.to_thread(gate.wait)
finally:
with counter_lock:
in_flight -= 1
return AIMessage(content="ok")
# Lead-loop call + isolated-loop call, both on the OLD (stale, cap 3) instance.
task_lead = asyncio.create_task(old_mw.awrap_model_call(SimpleNamespace(), handler))
fut_iso = _run_on_isolated_loop(lambda: old_mw.awrap_model_call(SimpleNamespace(), handler))
try:
for _ in range(100):
await asyncio.sleep(0)
with counter_lock:
if max_in_flight >= 1:
break
await asyncio.sleep(0.05) # let the second call queue on the limiter
with counter_lock:
assert max_in_flight == 1 # never exceeded the lowered cap across loops
assert limiter.limit == 1 # still 1 - stale instance did not rewrite it
finally:
# Always release the parked calls so a failing assertion doesn't hang.
gate.set()
try:
await asyncio.wait_for(task_lead, timeout=5)
except TimeoutError:
pass
try:
await asyncio.wait_for(asyncio.wrap_future(fut_iso), timeout=5)
except TimeoutError:
pass
with counter_lock:
assert in_flight == 0
@pytest.mark.anyio
@@ -1555,6 +1756,83 @@ async def test_limiter_cancellation_does_not_leak_capacity() -> None:
assert result_c.content == "b-ok"
@pytest.mark.anyio
async def test_limiter_cancellation_after_dequeue_hands_off_to_next_waiter(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""P1 #1 regression: cap=1, A holds the permit; B and C queue. A releases
(reserving/granting B); B is cancelled in the post-dequeue / pre-reacquire
handoff window; C must still complete WITHOUT another release - the reserved
permit is handed off to C rather than stranded with capacity idle.
``test_limiter_cancellation_does_not_leak_capacity`` cancels B while it is
still purely queued (before A releases); this test cancels B *after* it has
been dequeued+granted but *before* it wakes, which is the window the prior
limiter stranded the next waiter in.
"""
from deerflow.agents.middlewares.llm_error_handling_middleware import _get_process_limiter
middleware = _build_middleware(max_concurrent_llm_calls=1)
a_started = asyncio.Event()
gate = threading.Event()
async def handler_a(_request) -> AIMessage:
a_started.set()
await asyncio.to_thread(gate.wait)
return AIMessage(content="a-ok")
async def handler_ok(_request) -> AIMessage:
return AIMessage(content="ok")
task_a = asyncio.create_task(middleware.awrap_model_call(SimpleNamespace(), handler_a))
await asyncio.wait_for(a_started.wait(), timeout=2)
# B and C queue on the limiter (no permit available).
task_b = asyncio.create_task(middleware.awrap_model_call(SimpleNamespace(), handler_ok))
task_c = asyncio.create_task(middleware.awrap_model_call(SimpleNamespace(), handler_ok))
await asyncio.sleep(0.05) # let B and C register as waiters
assert not task_b.done() and not task_c.done()
limiter = _get_process_limiter()
assert limiter is not None
# Patch the wake so the FIRST granted waiter (B) does NOT wake yet - this
# is the post-dequeue / pre-reacquire handoff window. Later wakes (C, handed
# off from B's cancellation) proceed via the real path. ``_wake_locked`` is
# called under the limiter lock and must not block, so we only toggle which
# branch runs; B's event is simply never set, leaving B parked until cancel.
real_wake = type(limiter)._wake_locked
state = {"first_done": False}
def patched_wake(waiter: Any) -> bool:
if not state["first_done"]:
state["first_done"] = True
return True # pretend the loop is alive; do NOT set B's event
return real_wake(limiter, waiter)
monkeypatch.setattr(limiter, "_wake_locked", patched_wake)
# Release A: reserves the permit for B (dequeues B, granted=True) but B does
# not wake (patched). B sits in the handoff window.
gate.set()
await task_a
await asyncio.sleep(0.02) # let A's release grant B
assert state["first_done"] # B was dequeued+granted, never woken
# Cancel B while it is granted-but-not-yet-awake. Its cancellation must hand
# the reserved permit to C (the next waiter), waking C via the real path.
task_b.cancel()
with pytest.raises(asyncio.CancelledError):
await task_b
# C completes WITHOUT another release - no stranded permit.
result_c = await asyncio.wait_for(task_c, timeout=2)
assert result_c.content == "ok"
# No permit leaked: in_flight is back to 0 after C completes and releases.
assert limiter.in_flight == 0
# ---------- Burst-rate first-retry jitter (review P1) ----------
+7
View File
@@ -1986,6 +1986,13 @@ authorization:
# middleware, and ideally an nginx `limit_req` at the ingress.
#
# 0 disables the cap (default) - existing deployments see no behavior change.
#
# Per-process, not per-cluster: the cap bounds in-flight LLM calls within ONE
# gateway process. With GATEWAY_WORKERS > 1 the aggregate cap across the
# deployment is effectively `max_concurrent_calls * GATEWAY_WORKERS`, and a
# multi-node rollout multiplies it further - size the per-process value with
# that in mind (and pair it with an nginx `limit_req` at the ingress for a
# true cluster-wide slope cap).
# llm_call:
# # Max concurrently in-flight LLM calls across the whole process (default: 0 = disabled)