From fc8539b7b769a80fb6c5ec931d1c22caeffcdcb6 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Sun, 27 Sep 2026 13:11:55 +0530 Subject: [PATCH] fix(gateway): back off only runner-demoted events, not every None turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- gateway/platforms/base.py | 44 +++++++++++++-------- gateway/run_inbound.py | 15 +++++++ tests/gateway/test_busy_requeue_hot_loop.py | 25 +++++++++--- 3 files changed, 62 insertions(+), 22 deletions(-) diff --git a/gateway/platforms/base.py b/gateway/platforms/base.py index 76c85ce623..41cfb41c85 100644 --- a/gateway/platforms/base.py +++ b/gateway/platforms/base.py @@ -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: diff --git a/gateway/run_inbound.py b/gateway/run_inbound.py index f84b9e8796..da37a8007b 100644 --- a/gateway/run_inbound.py +++ b/gateway/run_inbound.py @@ -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 diff --git a/tests/gateway/test_busy_requeue_hot_loop.py b/tests/gateway/test_busy_requeue_hot_loop.py index d4803f502d..e9808a1320 100644 --- a/tests/gateway/test_busy_requeue_hot_loop.py +++ b/tests/gateway/test_busy_requeue_hot_loop.py @@ -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}"