From 3b0bb3b8bb6f0125b1bf000fcade334af01694a1 Mon Sep 17 00:00:00 2001 From: Dannyzen <659908+Dannyzen@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:09:53 +0530 Subject: [PATCH] fix(gateway): keep event loop alive during /compress and Relay drain Offload manual /compress temporary-agent cleanup through the existing bounded off-loop helper so a slow agent.close() cannot freeze the gateway event loop, heartbeat, or platform polling. Guarantee Relay transport teardown even when the runner cancels adapter.disconnect() during go_idle: shielded finally, 2s drain-path idle ACK budget under the 5s outer disconnect budget, and bounded supervisor/reader/ws.close awaits. Original commits: - fix(gateway): offload manual /compress cleanup from the event loop - fix(gateway): tear down Relay transport even if go_idle is cancelled - fix(gateway): keep Relay disconnect budgets inside the runner window By @Dannyzen (PR #78027), salvaged onto current main. --- gateway/relay/adapter.py | 48 +++++++++--- gateway/relay/ws_transport.py | 21 +++-- gateway/slash_commands.py | 9 ++- tests/gateway/relay/test_relay_adapter.py | 46 +++++++++++ tests/gateway/test_compress_command.py | 96 +++++++++++++++++++++++ 5 files changed, 202 insertions(+), 18 deletions(-) diff --git a/gateway/relay/adapter.py b/gateway/relay/adapter.py index d64e2fb00b8a7..73f7cf2941028 100644 --- a/gateway/relay/adapter.py +++ b/gateway/relay/adapter.py @@ -31,6 +31,14 @@ from gateway.session import SessionSource logger = logging.getLogger(__name__) +# Keep the drain-path going-idle ACK budget strictly under the runner's default +# adapter disconnect timeout (5s). If go_idle consumes the whole outer budget, +# cancellation can fire before transport.disconnect() and leave the websocket +# open. Paired with transport teardown budgets of 1s each for supervisor, +# reader, and ws.close (~3s), the full drain path stays inside 5s. +_RELAY_GO_IDLE_ON_DISCONNECT_TIMEOUT_S = 2.0 +_RELAY_REVOCATION_MONITOR_TEARDOWN_TIMEOUT_S = 1.0 + def _utf16_len(text: str) -> int: """Count UTF-16 code units (Telegram's length unit).""" @@ -816,8 +824,11 @@ class RelayAdapter(BasePlatformAdapter): if self._revocation_monitor is not None: self._revocation_monitor.cancel() try: - await self._revocation_monitor - except (asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown + await asyncio.wait_for( + self._revocation_monitor, + timeout=_RELAY_REVOCATION_MONITOR_TEARDOWN_TIMEOUT_S, + ) + except (asyncio.TimeoutError, asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown pass self._revocation_monitor = None if self._transport is not None: @@ -831,15 +842,32 @@ class RelayAdapter(BasePlatformAdapter): # the ack (Q-5.3c). Best-effort + guarded: a transport without go_idle # (the stub) or a failed/timed-out ack must not block shutdown — we # proceed to disconnect exactly as before, no regression. - go_idle = getattr(self._transport, "go_idle", None) - if callable(go_idle): + # + # transport.disconnect() runs in finally so an outer cancellation + # during go_idle (runner default adapter budget is 5s) still closes + # the socket/supervisor instead of leaking them. shield() keeps the + # teardown await itself from being cancelled mid-flight. + try: + go_idle = getattr(self._transport, "go_idle", None) + if callable(go_idle): + try: + result: Any = go_idle( + timeout_s=_RELAY_GO_IDLE_ON_DISCONNECT_TIMEOUT_S + ) + if asyncio.iscoroutine(result): + await result + except Exception: # noqa: BLE001 - going-idle is an optimization, never blocks drain + logger.debug( + "relay going_idle failed during drain", exc_info=True + ) + finally: try: - result: Any = go_idle() - if asyncio.iscoroutine(result): - await result - except Exception: # noqa: BLE001 - going-idle is an optimization, never blocks drain - logger.debug("relay going_idle failed during drain", exc_info=True) - await self._transport.disconnect() + await asyncio.shield(self._transport.disconnect()) + except Exception: # noqa: BLE001 - teardown must not block outer cancel propagation + logger.debug( + "relay transport disconnect failed during drain", + exc_info=True, + ) async def go_dormant(self) -> bool: """Quiesce the relay for a scale-to-zero suspend (D12 / Phase 0). diff --git a/gateway/relay/ws_transport.py b/gateway/relay/ws_transport.py index a4bfc4f6b2607..442453ab74b7f 100644 --- a/gateway/relay/ws_transport.py +++ b/gateway/relay/ws_transport.py @@ -53,6 +53,10 @@ WEBSOCKETS_AVAILABLE = websockets is not None # How long to wait for the handshake descriptor and for each outbound result. _HANDSHAKE_TIMEOUT_S = 30.0 _OUTBOUND_TIMEOUT_S = 30.0 +# Bound supervisor/reader/ws.close awaits so a wedged peer cannot stall +# adapter.disconnect. Three sequential awaits at 1.0s stay under the runner's +# default 5s adapter disconnect budget (plus the 2s go_idle ACK budget). +_TEARDOWN_AWAIT_TIMEOUT_S = 1.0 # Phase 7 Unit 7d-B: the application close code the connector sends when it # rejects/revokes a gateway's WS upgrade auth (mirrors the connector's @@ -501,23 +505,26 @@ class WebSocketRelayTransport: if self._supervisor is not None: self._supervisor.cancel() try: - await self._supervisor - except (asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown + await asyncio.wait_for( + self._supervisor, timeout=_TEARDOWN_AWAIT_TIMEOUT_S + ) + except (asyncio.TimeoutError, asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown pass self._supervisor = None if self._reader is not None: self._reader.cancel() try: - await self._reader - except (asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown + await asyncio.wait_for(self._reader, timeout=_TEARDOWN_AWAIT_TIMEOUT_S) + except (asyncio.TimeoutError, asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown pass self._reader = None if self._ws is not None: try: - await self._ws.close() - except Exception: # noqa: BLE001 + await asyncio.wait_for(self._ws.close(), timeout=_TEARDOWN_AWAIT_TIMEOUT_S) + except (asyncio.TimeoutError, asyncio.CancelledError, Exception): # noqa: BLE001 pass - self._ws = None + finally: + self._ws = None # Fail any in-flight outbound waiters so callers don't hang. for fut in self._pending.values(): if not fut.done(): diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index e719eb2ac88a0..7b87e435055cb 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -4147,7 +4147,14 @@ class GatewaySlashCommandsMixin: # Evict cached agent so next turn rebuilds system prompt # from current files (SOUL.md, memory, etc.). self._evict_cached_agent(session_key) - self._cleanup_agent_resources(tmp_agent) + # Off-loop + bounded: temporary-agent teardown can block on + # subprocess/network/SQLite work. Running it inline freezes the + # gateway loop and stalls platform polling / heartbeat, the same + # wedge class fixed for /new (#35994) and hygiene/shutdown + # (#53175). + await self._cleanup_agent_resources_off_loop( + tmp_agent, context="manual compression" + ) lines = [f"🗜️ {summary['headline']}"] if focus_topic: lines.append(t("gateway.compress.focus_line", topic=focus_topic)) diff --git a/tests/gateway/relay/test_relay_adapter.py b/tests/gateway/relay/test_relay_adapter.py index ba93036067137..137c0b15cd870 100644 --- a/tests/gateway/relay/test_relay_adapter.py +++ b/tests/gateway/relay/test_relay_adapter.py @@ -1,5 +1,7 @@ """RelayAdapter capability-advertisement tests (relay Phase 1, Task 1.1).""" +import asyncio + import pytest from gateway.config import Platform, PlatformConfig @@ -274,3 +276,47 @@ async def test_get_chat_info_local_fallback_when_not_advertised(): info = await a.get_chat_info("chan-1") assert info == {"name": "chan-1", "type": "dm"} assert t.calls == [] + + +class _HangOnIdleTransport: + """Transport that hangs in go_idle so outer disconnect cancellation can race it.""" + + def __init__(self): + self.go_idle_started = asyncio.Event() + self.go_idle_timeouts: list[float] = [] + self.disconnect_calls = 0 + + def set_inbound_handler(self, h): # noqa: D401 + self._h = h + + async def go_idle(self, timeout_s: float = 10.0): + self.go_idle_timeouts.append(timeout_s) + self.go_idle_started.set() + await asyncio.sleep(3600) + return False + + async def disconnect(self): + self.disconnect_calls += 1 + + +@pytest.mark.asyncio +async def test_disconnect_tears_down_transport_when_go_idle_is_cancelled(): + """Runner disconnect budgets can cancel adapter.disconnect mid go_idle. + + The gateway runner's default adapter disconnect budget is 5s, while + transport.go_idle defaults to 10s. If cancellation lands during the idle + handshake, transport.disconnect must still run so the websocket/supervisor + cannot outlive the adapter. + """ + transport = _HangOnIdleTransport() + adapter = RelayAdapter(PlatformConfig(), make_desc(platform="discord"), transport=transport) + + task = asyncio.create_task(adapter.disconnect()) + await asyncio.wait_for(transport.go_idle_started.wait(), timeout=1.0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert transport.disconnect_calls == 1 + assert transport.go_idle_timeouts + assert transport.go_idle_timeouts[0] < 5.0 diff --git a/tests/gateway/test_compress_command.py b/tests/gateway/test_compress_command.py index 599ad2b161559..4e502447b75dd 100644 --- a/tests/gateway/test_compress_command.py +++ b/tests/gateway/test_compress_command.py @@ -1,5 +1,7 @@ """Tests for gateway /compress user-facing messaging.""" +import asyncio +import threading from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -420,3 +422,97 @@ async def test_compress_command_single_profile_skips_profile_resolution(): runner._resolve_profile_home_for_source.assert_not_called() runner._shutdown_executor() + + +@pytest.mark.asyncio +async def test_compress_command_cleanup_does_not_block_event_loop(): + """Manual /compress must not run agent teardown on the gateway event loop. + + #53175 offloaded session-expiry, hygiene, and shutdown cleanup, but the + manual /compress finally still called ``_cleanup_agent_resources`` inline. + A slow ``agent.close()`` there freezes the whole loop and stops the + runtime-status heartbeat from advancing — the same wedge class as the + original incident. + + Observation must happen from a side thread: if cleanup blocks the event + loop, an ``await``-based waiter cannot sample ticks until close returns, + which falsely looks healthy after the block ends. + """ + import time + + history = _make_history() + compressed = [ + history[0], + {"role": "assistant", "content": "compressed summary"}, + history[-1], + ] + runner = _make_runner(history) + + close_started = threading.Event() + release_close = threading.Event() + + def slow_close(): + close_started.set() + release_close.wait(timeout=5) + + agent_instance = MagicMock() + agent_instance.shutdown_memory_provider = MagicMock() + agent_instance.close = slow_close + agent_instance._cached_system_prompt = "" + agent_instance.tools = None + agent_instance.context_compressor.has_content_to_compress.return_value = True + agent_instance.context_compressor._last_compress_aborted = False + agent_instance.context_compressor._last_summary_fallback_used = False + agent_instance.context_compressor._last_summary_dropped_count = 0 + agent_instance.context_compressor._last_summary_error = None + agent_instance.context_compressor._last_aux_model_failure_model = None + agent_instance.context_compressor._last_aux_model_failure_error = None + agent_instance.session_id = "sess-1" + agent_instance._compress_context.return_value = (compressed, "") + agent_instance._compression_skipped_due_to_lock = False + agent_instance._session_messages = None + + ticks = {"n": 0} + stop = threading.Event() + observed = {} + + async def _heartbeat(): + while not stop.is_set(): + ticks["n"] += 1 + await asyncio.sleep(0.005) + + def _observer(): + # threading.Event wait does not need the event loop. Sample ticks + # while close() is still held so an on-loop teardown is visible. + if not close_started.wait(timeout=5): + observed["error"] = "close() never started" + release_close.set() + return + baseline = ticks["n"] + time.sleep(0.12) + observed["ticks_during_block"] = ticks["n"] - baseline + release_close.set() + + hb = asyncio.create_task(_heartbeat()) + observer = threading.Thread(target=_observer, name="compress-cleanup-observer", daemon=True) + observer.start() + + with ( + patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "***"}), + patch("gateway.run._resolve_gateway_model", return_value="test-model"), + patch("run_agent.AIAgent", return_value=agent_instance), + patch("agent.model_metadata.estimate_request_tokens_rough", return_value=100), + ): + result = await runner._handle_compress_command(_make_event()) + + observer.join(timeout=5) + stop.set() + await hb + runner._shutdown_executor() + + assert "Compressed:" in result + assert "error" not in observed, observed.get("error") + assert observed.get("ticks_during_block", 0) >= 5, ( + "event loop was blocked during manual /compress cleanup: only " + f"{observed.get('ticks_during_block')} ticks while agent.close() was running" + )