fix(llm): release owned recovery probe on cancellation (#5197)

* fix(llm): release owned recovery probe on cancellation

* docs: keep middleware guidance within chain budget

---------

Co-authored-by: zaoshangduziteng <309590849+zaoshangduziteng@users.noreply.github.com>

README.md

@@ -333,6 +333,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
 > After a run publishes its terminal stream marker, its process-local `RunRecord` remains available for the existing five-minute grace period before cleanup; durable run history remains available through `RunStore`, while the stream bridge retains its delivery tail on its separate cleanup schedule.
 >
 > Run cancellation may land on any Gateway worker. A non-owning worker now persists the interrupt or rollback request for the live owner, which observes it during lease renewal and performs the normal cancellation flow; load-balancer routing alone no longer produces a 409. The first accepted action wins even if a retry lands on the owner, and accepted cancellation competes atomically with owner completion. Dead owners still follow lease takeover and orphan recovery. Cancellation latency is therefore bounded by the lease heartbeat interval.
+
+> Cancelling a model recovery probe, including while it is queued or waiting to retry, lets the next call check whether the provider has recovered. Cancellation does not count as a provider failure or release another call's active recovery probe.
 >
 > With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation.
 >

backend/packages/harness/deerflow/agents/middlewares/AGENTS.md

@@ -2,16 +2,14 @@
 
 Persisted delegation verdicts are untrusted durable context; ledger rendering revalidates them and ignores malformed values.
 
-Lead-agent middlewares are assembled in strict order across three functions: the shared base in `packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py` (`_build_runtime_middlewares`, exposed via `build_lead_runtime_middlewares`), then the lead-only middlewares appended in `packages/harness/deerflow/agents/lead_agent/agent.py` (`build_middlewares`). Items marked *(optional)* are appended only when their config/runtime condition holds, so the live chain length varies.
+Assembly order: `tool_error_handling_middleware.py::_build_runtime_middlewares` (exposed as `build_lead_runtime_middlewares`), then `../lead_agent/agent.py::build_middlewares` appends lead-only entries. Optional entries require their config/runtime condition.
 
 **Message provenance.** A middleware that injects or rewrites a message stamps
 `additional_kwargs` with the neutral provenance keys from
 `deerflow_extension_api.provenance` (`deerflow_content_kind`,
 `deerflow_producer_kind`, and optionally `deerflow_producer_entity_id`) via
-`provenance_kwargs()`. The producer is not recoverable downstream — by the
-model-call boundary the message is indistinguishable from any other — so the
-fact is recorded where it is known. Stamping is unconditional: a fact whose
-presence depends on whether an observer is installed is not a fact. All three
+`provenance_kwargs()`. Stamp at injection/rewrite regardless of installed
+observers; downstream cannot recover the producer. All three
 keys are in `_SERVER_OWNED_MESSAGE_METADATA_KEYS`, so a caller cannot forge
 provenance on inbound messages. Currently stamped by: DynamicContext (reminder + memory),
 DurableContext (contract + data), SystemMessageCoalescing, ViewImage,
@@ -55,7 +53,7 @@ it to that middleware's declaration in the same change.
    their narrower discovery allowlists never rebuild the shared thread view or
    force eager sandbox acquisition.
 7. **DanglingToolCallMiddleware** - Injects placeholder ToolMessages for AIMessage tool_calls that lack responses (e.g., user interruption), preserving raw provider tool-call payloads in `additional_kwargs["tool_calls"]`; malformed tool-call names and arguments are sanitized in the model-bound request so strict OpenAI-compatible providers do not reject the next request
-8. **LLMErrorHandlingMiddleware** - Normalizes provider/model invocation failures into recoverable assistant-facing errors before later stages run
+8. **LLMErrorHandlingMiddleware** - Converts provider/model failures to recoverable assistant errors. Async cancellation at admission, provider execution, retry events, or backoff releases only the call's own half-open probe (ownership assigned under the circuit lock), then propagates unchanged, without retry or failure accounting.
 9. **Authorization / GuardrailMiddleware** - Up to two independent pre-tool-call gates run here. When `authorization.enabled`, the `AuthorizationProvider` instance already used for Layer 1 capability filtering is wrapped by `GuardrailAuthorizationAdapter` and reused for Layer 2 execution checks. A generated `tool_search` bypasses the adapter's second provider call only when the current build has a concrete deferred setup; its catalog was already filtered by Layer 1, and an ordinary same-named tool without that deferred setup receives no exemption. When `guardrails.enabled`, the explicitly configured `GuardrailProvider` is appended after authorization and still evaluates every call, including `tool_search`. Authorization therefore runs outermost and can deny before an external guardrail call; both use the existing middleware's fail-closed, audit, sync/async, and error-`ToolMessage` behavior. See the authorization RFC and [docs/GUARDRAILS.md](../../../../../docs/GUARDRAILS.md).
 
    Every guardrail decision path publishes a neutral

backend/packages/harness/deerflow/agents/middlewares/llm_error_handling_middleware.py

@@ -410,6 +410,7 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
         self._circuit_open_until = 0.0
         self._circuit_state = "closed"
         self._circuit_probe_in_flight = False
+        self._circuit_probe_token: object | None = None
 
     def _max_attempts_for(self, exc: BaseException, reason: str = "transient") -> int:
         """Return the effective max attempt count for this exception.
@@ -430,7 +431,7 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
             candidates.append(reason_override)
         return min(candidates)
 
-    def _check_circuit(self) -> bool:
+    def _check_circuit(self, *, probe_token: object | None = None) -> bool:
         """Returns True if circuit is OPEN (fast fail), False otherwise."""
         with self._circuit_lock:
             now = time.time()
@@ -440,11 +441,13 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
                     return True
                 self._circuit_state = "half_open"
                 self._circuit_probe_in_flight = False
+                self._circuit_probe_token = None
 
             if self._circuit_state == "half_open":
                 if self._circuit_probe_in_flight:
                     return True
                 self._circuit_probe_in_flight = True
+                self._circuit_probe_token = probe_token
                 return False
 
             return False
@@ -457,6 +460,7 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
             self._circuit_open_until = 0.0
             self._circuit_state = "closed"
             self._circuit_probe_in_flight = False
+            self._circuit_probe_token = None
 
     def _record_failure(self) -> None:
         with self._circuit_lock:
@@ -464,6 +468,7 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
                 self._circuit_open_until = time.time() + self.circuit_recovery_timeout_sec
                 self._circuit_state = "open"
                 self._circuit_probe_in_flight = False
+                self._circuit_probe_token = None
                 logger.error(
                     "Circuit breaker probe failed (Open). Will probe again after %ds.",
                     self.circuit_recovery_timeout_sec,
@@ -476,22 +481,27 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
                 if self._circuit_state != "open":
                     self._circuit_state = "open"
                     self._circuit_probe_in_flight = False
+                    self._circuit_probe_token = None
                     logger.error(
                         "Circuit breaker tripped (Open). Threshold reached (%d). Will probe after %ds.",
                         self.circuit_failure_threshold,
                         self.circuit_recovery_timeout_sec,
                     )
 
-    def _release_half_open_probe(self) -> None:
+    def _release_half_open_probe(self, *, probe_token: object | None = None) -> None:
         """Release the in-flight half-open probe without recording a failure.
 
         Used when something other than a classified success/failure consumes the probe (a
         GraphBubbleUp control-flow signal, or a non-retriable error), so the circuit can admit
-        the next probe instead of fast-failing forever.
+        the next probe instead of fast-failing forever. Cancellation supplies an
+        admission token so an older call cannot release a different call's probe.
         """
         with self._circuit_lock:
+            if probe_token is not None and self._circuit_probe_token is not probe_token:
+                return
             if self._circuit_state == "half_open":
                 self._circuit_probe_in_flight = False
+                self._circuit_probe_token = None
 
     def _classify_error(self, exc: BaseException) -> tuple[bool, str]:
         detail = _extract_error_detail(exc)
@@ -837,7 +847,8 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        if self._check_circuit():
+        probe_token = object()
+        if self._check_circuit(probe_token=probe_token):
             return self._build_error_fallback_message(
                 self._build_circuit_breaker_message(),
                 error_type="CircuitBreakerOpen",
@@ -847,48 +858,54 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]):
 
         attempt = 1
         prev_delay_ms: int | None = None
-        while True:
-            try:
-                response = await self._bounded_model_call(request, handler)
-                self._record_success()
-                return response
-            except GraphBubbleUp:
-                # Preserve LangGraph control-flow signals (interrupt/pause/resume).
-                self._release_half_open_probe()
-                raise
-            except Exception as exc:
-                retriable, reason = self._classify_error(exc)
-                max_attempts = self._max_attempts_for(exc, reason)
-                if retriable and attempt < max_attempts:
-                    wait_ms = self._build_retry_delay_ms(prev_delay_ms, exc, reason)
-                    prev_delay_ms = wait_ms
-                    logger.warning(
-                        "Transient LLM error on attempt %d/%d; retrying in %dms: %s",
-                        attempt,
-                        max_attempts,
-                        wait_ms,
-                        _extract_error_detail(exc),
-                    )
-                    await self._aemit_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts)
-                    await asyncio.sleep(wait_ms / 1000)
-                    attempt += 1
-                    continue
-                logger.warning(
-                    "LLM call failed after %d attempt(s): %s",
-                    attempt,
-                    _extract_error_detail(exc),
-                    exc_info=exc,
-                )
-                if retriable and reason != "burst_rate":
-                    self._record_failure()
-                else:
-                    # 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.
+        try:
+            while True:
+                try:
+                    response = await self._bounded_model_call(request, handler)
+                    self._record_success()
+                    return response
+                except GraphBubbleUp:
+                    # Preserve LangGraph control-flow signals (interrupt/pause/resume).
                     self._release_half_open_probe()
-                return self._build_user_fallback_message(exc, reason)
+                    raise
+                except Exception as exc:
+                    retriable, reason = self._classify_error(exc)
+                    max_attempts = self._max_attempts_for(exc, reason)
+                    if retriable and attempt < max_attempts:
+                        wait_ms = self._build_retry_delay_ms(prev_delay_ms, exc, reason)
+                        prev_delay_ms = wait_ms
+                        logger.warning(
+                            "Transient LLM error on attempt %d/%d; retrying in %dms: %s",
+                            attempt,
+                            max_attempts,
+                            wait_ms,
+                            _extract_error_detail(exc),
+                        )
+                        await self._aemit_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts)
+                        await asyncio.sleep(wait_ms / 1000)
+                        attempt += 1
+                        continue
+                    logger.warning(
+                        "LLM call failed after %d attempt(s): %s",
+                        attempt,
+                        _extract_error_detail(exc),
+                        exc_info=exc,
+                    )
+                    if retriable and reason != "burst_rate":
+                        self._record_failure()
+                    else:
+                        # 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)
+        except asyncio.CancelledError:
+            # Cancellation can arrive during admission, the provider call, retry
+            # event delivery, or backoff. It is not a provider failure.
+            self._release_half_open_probe(probe_token=probe_token)
+            raise
 
 
 def _matches_any(detail: str, patterns: tuple[str, ...]) -> bool:

backend/tests/test_llm_error_handling_middleware.py

@@ -288,6 +288,120 @@ def test_async_model_call_propagates_graph_bubble_up() -> None:
         asyncio.run(middleware.awrap_model_call(SimpleNamespace(), handler))
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize("cancel_during", ["provider", "backoff", "queue"])
+async def test_cancelled_recovery_probe_allows_next_model_call(cancel_during: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    middleware = _build_middleware(retry_max_attempts=1, circuit_failure_threshold=1, circuit_recovery_timeout_sec=0, max_concurrent_llm_calls=1)
+
+    async def unavailable(_request) -> AIMessage:
+        raise FakeError("Service unavailable", status_code=503)
+
+    await middleware.awrap_model_call(SimpleNamespace(), unavailable)
+    middleware.retry_max_attempts = 2
+    entered = asyncio.Event()
+
+    async def retry_sleep(_delay: float) -> None:
+        entered.set()
+        await asyncio.Event().wait()
+
+    if cancel_during == "backoff":
+        monkeypatch.setattr(asyncio, "sleep", retry_sleep)
+
+    async def recovering(_request) -> AIMessage:
+        assert cancel_during != "queue", "A queued request must not reach the provider"
+        if cancel_during == "backoff":
+            raise FakeError("Service unavailable", status_code=503)
+        entered.set()
+        await asyncio.Event().wait()
+        return AIMessage(content="unreachable")
+
+    occupied = asyncio.Event()
+    release_slot = asyncio.Event()
+
+    async def occupying_handler(_request) -> AIMessage:
+        occupied.set()
+        await release_slot.wait()
+        return AIMessage(content="other call finished")
+
+    async def run_probe():
+        if cancel_during == "queue":
+            entered.set()
+        return await middleware.awrap_model_call(SimpleNamespace(), recovering)
+
+    blocker = None
+    probe = None
+    try:
+        if cancel_during == "queue":
+            other_middleware = _build_middleware(max_concurrent_llm_calls=1)
+            blocker = asyncio.create_task(other_middleware.awrap_model_call(SimpleNamespace(), occupying_handler))
+            await asyncio.wait_for(occupied.wait(), timeout=1)
+        probe = asyncio.create_task(run_probe())
+        await asyncio.wait_for(entered.wait(), timeout=1)
+        probe.cancel("user stopped recovery")
+        with pytest.raises(asyncio.CancelledError, match="user stopped recovery"):
+            await probe
+    finally:
+        if probe is not None:
+            probe.cancel()
+            await asyncio.gather(probe, return_exceptions=True)
+        release_slot.set()
+        if blocker is not None:
+            await asyncio.wait_for(blocker, timeout=1)
+
+    async def healthy(_request) -> AIMessage:
+        return AIMessage(content="recovered")
+
+    result = await asyncio.wait_for(middleware.awrap_model_call(SimpleNamespace(), healthy), timeout=1)
+    assert result.content == "recovered"
+
+
+@pytest.mark.anyio
+async def test_cancelling_older_call_does_not_release_another_recovery_probe() -> None:
+    middleware = _build_middleware(retry_max_attempts=1, circuit_failure_threshold=1, circuit_recovery_timeout_sec=0)
+    old_entered = asyncio.Event()
+    probe_entered = asyncio.Event()
+    finish_probe = asyncio.Event()
+
+    async def old_handler(_request) -> AIMessage:
+        old_entered.set()
+        await asyncio.Event().wait()
+        return AIMessage(content="unreachable")
+
+    async def unavailable(_request) -> AIMessage:
+        raise FakeError("Service unavailable", status_code=503)
+
+    async def recovering(_request) -> AIMessage:
+        probe_entered.set()
+        await finish_probe.wait()
+        return AIMessage(content="recovered")
+
+    async def extra_handler(_request) -> AIMessage:
+        pytest.fail("Only the existing recovery probe may reach the provider")
+
+    old_call = asyncio.create_task(middleware.awrap_model_call(SimpleNamespace(), old_handler))
+    probe = None
+    try:
+        await asyncio.wait_for(old_entered.wait(), timeout=1)
+        await middleware.awrap_model_call(SimpleNamespace(), unavailable)
+        probe = asyncio.create_task(middleware.awrap_model_call(SimpleNamespace(), recovering))
+        await asyncio.wait_for(probe_entered.wait(), timeout=1)
+
+        old_call.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await old_call
+
+        blocked = await middleware.awrap_model_call(SimpleNamespace(), extra_handler)
+        assert blocked.additional_kwargs["error_type"] == "CircuitBreakerOpen"
+        finish_probe.set()
+        result = await probe
+        assert result.content == "recovered"
+    finally:
+        tasks = [old_call] if probe is None else [old_call, probe]
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+
 def test_circuit_half_open_graph_bubble_up_resets_probe() -> None:
     """Verify that GraphBubbleUp in half_open state resets probe_in_flight."""
     middleware = _build_middleware()