fix(gateway): back off only runner-demoted events, not every None turn

The drain back-off keyed on `response is None`, but None is the normal
return for every streamed turn and for queue/steer busy modes, so
ordinary chained follow-ups were delayed 0.25s -> 5s forever. The busy
fast-path in run_inbound now tags the adapter's pending head
(`_busy_requeued`) where it demotes the same inbound event (or its
rewrite-hook copy) back into the queue (interrupt demotion, queue mode,
steer fallback); the drain backs off only for a tagged event and
otherwise resets the counter and dispatches at once — single reset owner.

Also: clear _requeue_counts on cancel_session_processing, stale-lock
heal, session end and shutdown; restore the pending event if cancelled
during the back-off sleep; reuse agent.retry_utils.jittered_backoff.
The regression test now chains 3 genuine follow-ups after a streamed
(None) turn and requires each to dispatch in <0.1s (red on pre-fold).
This commit is contained in:
kshitijk4poor
2026-09-27 13:11:55 +05:30
committed by kshitij
parent 506cc14def
commit fc8539b7b7
3 changed files with 62 additions and 22 deletions

View File

@@ -1925,8 +1925,8 @@ class BasePlatformAdapter(ABC):
# could drop a newer guard.
self._active_sessions: Dict[str, asyncio.Event] = {}
self._pending_messages: Dict[str, MessageEvent] = {}
# Consecutive in-band drains per session whose handler returned nothing while an event sat
# pending (the runner's busy-demotion re-queue); drives the drain back-off (#123229).
# Consecutive in-band drains per session of an event the runner DEMOTED back into the queue
# (tagged ``_busy_requeued``); drives the drain back-off (#123229).
self._requeue_counts: Dict[str, int] = {}
self._pending_text_batches: Dict[str, MessageEvent] = {}
self._pending_text_batch_tasks: Dict[str, asyncio.Task] = {}
@@ -3862,6 +3862,7 @@ class BasePlatformAdapter(ABC):
self.name, session_key)
self._active_sessions.pop(session_key, None)
self._pending_messages.pop(session_key, None)
self._requeue_counts.pop(session_key, None)
self._session_tasks.pop(session_key, None)
self._discard_text_debounce(session_key)
return True
@@ -3899,6 +3900,7 @@ class BasePlatformAdapter(ABC):
"""Cancel in-flight processing for one session. ``release_guard=False`` keeps the guard so
reset-like commands finish atomically; the await is bounded (5s) so a wedged finally can't
stall."""
self._requeue_counts.pop(session_key, None)
task = self._session_tasks.pop(session_key, None)
if task is not None and not task.done():
logger.debug("[%s] Cancelling active processing for session %s", self.name, session_key)
@@ -4476,8 +4478,6 @@ class BasePlatformAdapter(ABC):
await self._run_processing_hook("on_processing_start", event)
event._turn_marker_handoff = self.gateway_runner is not None # it can release the marker
response = await self._message_handler(event)
if response is not None:
self._requeue_counts.pop(session_key, None) # the handler consumed its event
# A muted diagnostic wake ran for the session; its reply is not presented. The
# policy read binds the routed profile; delivery itself stays in the launch scope.
with self._media_delivery_scope(event.source):
@@ -4554,7 +4554,7 @@ class BasePlatformAdapter(ABC):
self._clear_session_guard(session_key)
await self._stop_typing_refresh(event.source.chat_id, typing_task, metadata=_thread_metadata)
self._spawn_drain_task(pending_event, session_key,
delay=self._requeue_backoff_delay(session_key, response))
delay=self._requeue_backoff_delay(session_key, pending_event))
return # Drain task owns the session now.
except asyncio.CancelledError:
expected = asyncio.current_task() in self._expected_cancelled_tasks
@@ -4586,21 +4586,26 @@ class BasePlatformAdapter(ABC):
_REQUEUE_BACKOFF_INITIAL_SECONDS = 0.25
_REQUEUE_BACKOFF_MAX_SECONDS = 5.0
def _requeue_backoff_delay(self, session_key: str, response: Any) -> float:
"""Delay before re-dispatching a follow-up the handler left pending without answering.
def _requeue_backoff_delay(self, session_key: str, pending_event: MessageEvent) -> float:
"""Delay before re-dispatching the popped follow-up.
The runner's busy-demotion puts the event back (possibly a rewritten copy, so this counts
per session, not by event identity) and returns None; re-dispatching at once hot-loops for
the whole busy window (#123229). The first re-queue stays immediate — restart auto-resume
relies on one self-bounce — then back off exponentially to a cap. Defers, never drops."""
if response is not None:
Only an event the runner DEMOTED back into this queue (tagged ``_busy_requeued`` at the
re-queue site — rewrite-hook copies included) backs off; re-dispatching it at once hot-loops
for the whole busy window (#123229). Anything else — a genuine follow-up after a normal
(e.g. streamed, None-returning) turn — resets the counter and runs immediately. The first
demotion stays immediate (restart auto-resume relies on one self-bounce), then back off
exponentially to a cap. Defers, never drops."""
if not getattr(pending_event, "_busy_requeued", False):
self._requeue_counts.pop(session_key, None)
return 0.0
from agent.retry_utils import jittered_backoff
pending_event._busy_requeued = False # the runner re-tags it if it demotes it again
attempts = self._requeue_counts.get(session_key, 0)
self._requeue_counts[session_key] = attempts + 1
if attempts == 0:
return 0.0
delay = min(self._REQUEUE_BACKOFF_MAX_SECONDS,
self._REQUEUE_BACKOFF_INITIAL_SECONDS * 2 ** min(attempts - 1, 16))
delay = jittered_backoff(attempts, base_delay=self._REQUEUE_BACKOFF_INITIAL_SECONDS,
max_delay=self._REQUEUE_BACKOFF_MAX_SECONDS, jitter_ratio=0.0)
(logger.info if attempts == 1 else logger.debug)(
"[%s] Handler re-queued a pending event for %s again (session busy elsewhere); "
"backing off %.2fs", self.name, session_key, delay)
@@ -4619,7 +4624,12 @@ class BasePlatformAdapter(ABC):
async def _drain_after(self, pending_event: MessageEvent, session_key: str, delay: float) -> None:
if delay > 0:
await asyncio.sleep(delay)
try:
await asyncio.sleep(delay)
except asyncio.CancelledError:
# Not yet processing: hand the event back unless something newer took the slot.
self._pending_messages.setdefault(session_key, pending_event)
raise
await self._process_message_background(pending_event, session_key)
def _clear_session_guard(self, session_key: str) -> None:
@@ -4645,6 +4655,7 @@ class BasePlatformAdapter(ABC):
self._release_session_guard(session_key, guard=interrupt_event)
if session_key not in self._active_sessions:
self._session_tasks.pop(session_key, None)
self._requeue_counts.pop(session_key, None)
async def cancel_background_tasks(self) -> None:
"""Cancel in-flight background tasks (shutdown/replacement); 5s bound each,
@@ -4673,7 +4684,8 @@ class BasePlatformAdapter(ABC):
for state in self._text_debounce_store().values():
state.cancel_timer()
for bucket in (self._background_tasks, self._expected_cancelled_tasks, self._session_tasks,
self._pending_messages, self._active_sessions, self._text_debounce_store()):
self._pending_messages, self._active_sessions, self._requeue_counts,
self._text_debounce_store()):
bucket.clear()
def has_pending_interrupt(self, session_key: str) -> bool:

View File

@@ -708,9 +708,11 @@ class GatewayInboundMixin:
if effective_busy_input_mode == "queue":
logger.debug("PRIORITY queue follow-up for session %s", _quick_key)
self._queue_or_replace_pending_event(_quick_key, event)
self._hm_tag_busy_requeue(source, _quick_key)
return None
if effective_busy_input_mode == "steer":
self._hm_busy_steer(event, running_agent, _quick_key)
self._hm_tag_busy_requeue(source, _quick_key)
return None
# Subagent protection: an interrupt cascades through ``_active_children`` and aborts
# in-flight delegate_task work (/stop reached its handler above — still an escape hatch).
@@ -725,8 +727,21 @@ class GatewayInboundMixin:
return None
logger.info("PRIORITY interrupt demoted to queue for session %s %s", _quick_key, _demote)
self._queue_or_replace_pending_event(_quick_key, event)
self._hm_tag_busy_requeue(source, _quick_key)
return None
def _hm_tag_busy_requeue(self, source: SessionSource, _quick_key: str) -> None:
"""Mark the adapter's pending head as DEMOTED by this busy fast-path (the same inbound event,
or its rewrite-hook copy / merge, went back into the queue while an agent runs). Only tagged
events make the adapter drain back off (#123229); genuine follow-ups after a normal turn
stay immediate. A steer that was absorbed leaves nothing new pending; the head, if any, is
still blocked on the same running agent, so tagging it only defers it."""
adapter = self._delivery_adapter_for(source)
pending = getattr(adapter, "_pending_messages", None) if adapter else None
head = pending.get(_quick_key) if isinstance(pending, dict) else None
if head is not None:
head._busy_requeued = True
def _hm_quick_commands(self) -> dict:
"""User-defined ``quick_commands`` mapping from config (empty dict when unset/malformed)."""
cfg = self.config

View File

@@ -143,25 +143,38 @@ async def test_requeued_busy_event_does_not_hot_loop(rewrite_hook):
@pytest.mark.asyncio
async def test_requeued_event_runs_once_the_agent_finishes():
"""The back-off must defer, not drop: once the running agent is gone the event is processed."""
"""The back-off must defer, not drop: once the running agent is gone the event is processed.
And it must key on the runner's demotion, not on ``None``: every streamed turn returns None,
so chained genuine follow-ups after it must each dispatch immediately."""
adapter = _Adapter()
runner, agent, sk = _runner_with_running_agent(adapter, compression_in_flight=True)
handled = []
handled, starts, ends = [], [], []
real_handle = runner._handle_message
async def handler(event):
if sk not in runner._running_agents:
handled.append(event.text)
return "done"
starts.append(time.monotonic())
await asyncio.sleep(0.05)
ends.append(time.monotonic())
return None # streamed turn: the body was already delivered
return await real_handle(event)
adapter.set_message_handler(handler)
await adapter.handle_message(MessageEvent(text="queued msg", source=_source(), message_id="m2"))
await asyncio.sleep(0.6)
runner._running_agents.pop(sk) # the long turn finishes
for _ in range(100):
for _ in range(1000):
if handled:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.01)
for i in range(1, 4): # a genuine follow-up queued during each (None-returning) turn
await adapter.handle_message(MessageEvent(text=f"f{i}", source=_source(), message_id=f"f{i}"))
for _ in range(100):
if len(handled) > i:
break
await asyncio.sleep(0.01)
await adapter.cancel_background_tasks()
assert handled == ["queued msg"]
assert handled == ["queued msg", "f1", "f2", "f3"]
gaps = [starts[i] - ends[i - 1] for i in range(1, len(starts))]
assert all(g < 0.1 for g in gaps), f"genuine follow-ups backed off: {gaps}"