"""Regression test for #17758 — chained pending-message drains must not grow the call stack. Before the fix, ``_process_message_background`` finished a turn, found a pending follow-up, and drained it via ``await self._process_message_background(pending_event, session_key)``. Each queued follow-up added a frame to the call stack instead of starting fresh, so under sustained pending-queue activity the C stack would exhaust at ~2000 nested frames and the process would crash with SIGSEGV. After the fix, the in-band drain spawns a fresh task (mirroring the late-arrival drain pattern), so the stack stays bounded regardless of chain length. We assert the invariant directly: count nested ``_process_message_background`` frames at handler entry across a chain of N follow-ups. Recursion makes depth grow linearly (1, 2, 3, …, N); task spawning keeps it constant (1 every time). """ import asyncio import sys from unittest.mock import AsyncMock import pytest from gateway.config import Platform, PlatformConfig from gateway.platforms.base import ( BasePlatformAdapter, MessageEvent, MessageType, ) from gateway.session import SessionSource, build_session_key class _StubAdapter(BasePlatformAdapter): async def connect(self, *, is_reconnect: bool = False): pass async def disconnect(self): pass async def send(self, chat_id, text, **kwargs): return None async def get_chat_info(self, chat_id): return {} def _make_adapter(): adapter = _StubAdapter(PlatformConfig(enabled=True, token="t"), Platform.TELEGRAM) adapter._send_with_retry = AsyncMock(return_value=None) return adapter def _make_event(text="hi", chat_id="42"): return MessageEvent( text=text, message_type=MessageType.TEXT, source=SessionSource(platform=Platform.TELEGRAM, chat_id=chat_id, chat_type="dm"), ) def _sk(chat_id="42"): return build_session_key( SessionSource(platform=Platform.TELEGRAM, chat_id=chat_id, chat_type="dm") ) def _count_pmb_frames() -> int: """Walk the current call stack and count nested ``_process_message_background`` frames. Used to detect recursive in-band drains.""" f = sys._getframe() n = 0 while f is not None: if f.f_code.co_name == "_process_message_background": n += 1 f = f.f_back return n @pytest.mark.asyncio async def test_in_band_drain_does_not_grow_stack(): """Issue #17758: chained pending-message drains must not recurse. Queue a fresh pending message inside each handler invocation so the in-band drain block fires for every turn in the chain. After N turns, the recorded stack depth at handler entry must stay bounded. Pre-fix, depths would be 1, 2, 3, …, N; post-fix, depths are 1 every time because each drain runs in its own task. """ N = 12 adapter = _make_adapter() sk = _sk() depths: list[int] = [] next_index = [1] async def handler(event): depths.append(_count_pmb_frames()) if next_index[0] < N: adapter._pending_messages[sk] = _make_event(text=f"M{next_index[0]}") next_index[0] += 1 return "ok" adapter._message_handler = handler await adapter.handle_message(_make_event(text="M0")) # Drain the chain. Each turn schedules the next via the in-band # drain block, so we wait until N handler runs have completed and # the session has been released. # 2000 * 0.01s = 20s budget: the old 4s budget flaked on loaded CI # runners (11/12 turns completed; main run 33455779041). for _ in range(2000): if len(depths) >= N and sk not in adapter._active_sessions: break await asyncio.sleep(0.01) await adapter.cancel_background_tasks() assert len(depths) == N, ( f"expected {N} handler runs in the chain, got {len(depths)}: depths={depths!r}" ) max_depth = max(depths) assert max_depth <= 2, ( f"in-band drain is recursing instead of spawning a fresh task — " f"stack depth grew with chain length: {depths!r}" ) # --------------------------------------------------------------------------- # Follow-up guardrails (belt-and-suspenders on top of the #17758 fix). # # The in-band drain hand-off changed cleanup semantics in three subtle ways # that the original fix reasoned about but didn't test directly. These # tests pin each invariant so future refactors can't silently regress them. # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_normal_path_releases_session_guard(): """The common path — one message, nothing queued — must still fully release ``_active_sessions[sk]`` and ``_session_tasks[sk]`` through the end-of-finally block. The #17758 fix moved ``_release_session_guard(...)`` under an ``if current_task is self._session_tasks.get(session_key)`` conditional. For the 99%-common case (no pending message, no handoff) ``current_task`` IS the stored task, so the guard must still fire. This test would fail if the conditional were ever tightened in a way that dropped the normal path.""" adapter = _make_adapter() sk = _sk() async def handler(event): return "ok" adapter._message_handler = handler await adapter.handle_message(_make_event(text="solo")) # Wait for the single-shot handler to fully unwind. for _ in range(200): if sk not in adapter._active_sessions and sk not in adapter._session_tasks: break await asyncio.sleep(0.01) await adapter.cancel_background_tasks() assert sk not in adapter._active_sessions, ( "normal-path unwind left _active_sessions[sk] populated — future " "messages would take the busy-handler path forever" ) assert sk not in adapter._session_tasks, ( "normal-path unwind left _session_tasks[sk] populated — " "stale-lock detection will treat a dead task as alive" ) @pytest.mark.asyncio async def test_drain_task_cancellation_releases_session(): """If the in-band drain task is cancelled (e.g. user sent ``/stop`` mid-drain), the session guard and task registry must still get cleaned up — the cancelled drain task's own ``finally`` runs and fires ``_release_session_guard``. The #17758 fix transfers ownership of ``_session_tasks[sk]`` to the drain task; the drain task's ``except asyncio.CancelledError`` branch must then own the cleanup. Without this test a future refactor could move cancellation handling in a way that leaves the session permanently pinned as busy after a cancel.""" adapter = _make_adapter() sk = _sk() turn_started = asyncio.Event() drain_hit_handler = asyncio.Event() async def handler(event): if event.text == "M0": # Queue a pending follow-up so an in-band drain task gets spawned. adapter._pending_messages[sk] = _make_event(text="M1") turn_started.set() return "ok" # M1 is the drained follow-up — hang so we can cancel the drain task. drain_hit_handler.set() try: await asyncio.sleep(0.2) except asyncio.CancelledError: raise adapter._message_handler = handler await adapter.handle_message(_make_event(text="M0")) # Wait for the drain task to actually start running M1. await asyncio.wait_for(drain_hit_handler.wait(), timeout=2) # Cancel the drain task mid-handler. drain_task = adapter._session_tasks.get(sk) assert drain_task is not None, "in-band drain did not install a drain task" assert not drain_task.done(), "drain task finished before we could cancel" drain_task.cancel() # Drain task's finally must release both registries. for _ in range(200): if sk not in adapter._active_sessions and sk not in adapter._session_tasks: break await asyncio.sleep(0.01) await adapter.cancel_background_tasks() assert sk not in adapter._active_sessions, ( "cancelled drain task did not release _active_sessions[sk] — " "the session stays permanently pinned as busy after a /stop mid-drain" ) assert sk not in adapter._session_tasks, ( "cancelled drain task did not release _session_tasks[sk] — " "stale-lock detection will treat the dead task as alive" ) @pytest.mark.asyncio async def test_late_arrival_drain_still_fires_when_no_in_band_drain(): """The late-arrival drain in ``finally`` must still spawn a fresh task when no in-band drain preceded it. Pre-#17758 this path already existed; the #17758 follow-up guard only re-queues when ``_session_tasks[sk] is not current_task``. For a late-arrival with no in-band drain, ``_session_tasks[sk]`` IS the current task, so the ``else`` branch must fire and spawn a drain task for the queued message. Queue a pending message *after* M0's handler returns (so the in-band drain block sees nothing) but *before* ``finally`` runs the late-arrival check — we do this by hooking ``_stop_typing``, which runs in finally before the late-arrival check.""" adapter = _make_adapter() sk = _sk() results: list[str] = [] original_stop_typing = getattr(adapter, "stop_typing", None) async def injecting_stop_typing(chat_id): # Simulate a message landing during the cleanup awaits. adapter._pending_messages[sk] = _make_event(text="late") if original_stop_typing: await original_stop_typing(chat_id) adapter.stop_typing = injecting_stop_typing async def handler(event): results.append(event.text) return "ok" adapter._message_handler = handler await adapter.handle_message(_make_event(text="first")) # Wait for the late-arrival drain task to finish the second event. # 2000 * 0.01s = 20s budget: the old 4s budget flaked on loaded CI # runners (11/12 turns completed; main run 33455779041). for _ in range(2000): if "late" in results and sk not in adapter._active_sessions: break await asyncio.sleep(0.01) await adapter.cancel_background_tasks() assert "first" in results, "original message handler did not run" assert "late" in results, ( "late-arrival drain did not spawn a drain task — a message that " "landed during cleanup awaits was silently dropped" )