From 5cbad5fa3739caca496ae60d054c926d95ebbc4e Mon Sep 17 00:00:00 2001
From: Dusk1e <yusufalweshdemir@gmail.com>
Date: Thu, 9 Apr 2026 22:54:40 +0300
Subject: [PATCH] fix(gateway): track unowned watcher tasks in GatewayRunner

---
 gateway/run.py                         | 22 +++++++--
 tests/gateway/test_gateway_shutdown.py | 64 ++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/gateway/run.py b/gateway/run.py
index 339954f5bef40..482a854857b12 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -570,6 +570,12 @@ def __init__(self, config: Optional[GatewayConfig] = None):
         # Track background tasks to prevent garbage collection mid-execution
         self._background_tasks: set = set()
 
+    def _track_background_task(self, task: asyncio.Task) -> asyncio.Task:
+        """Retain a strong reference to a fire-and-forget task until completion."""
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.discard)
+        return task
+
 
 
 
@@ -1246,13 +1252,17 @@ async def start(self) -> bool:
             from tools.process_registry import process_registry
             while process_registry.pending_watchers:
                 watcher = process_registry.pending_watchers.pop(0)
-                asyncio.create_task(self._run_process_watcher(watcher))
+                self._track_background_task(
+                    asyncio.create_task(self._run_process_watcher(watcher))
+                )
                 logger.info("Resumed watcher for recovered process %s", watcher.get("session_id"))
         except Exception as e:
             logger.error("Recovered watcher setup error: %s", e)
 
         # Start background session expiry watcher for proactive memory flushing
-        asyncio.create_task(self._session_expiry_watcher())
+        self._track_background_task(
+            asyncio.create_task(self._session_expiry_watcher())
+        )
 
         # Start background reconnection watcher for platforms that failed at startup
         if self._failed_platforms:
@@ -1261,7 +1271,9 @@ async def start(self) -> bool:
                 len(self._failed_platforms),
                 ", ".join(p.value for p in self._failed_platforms),
             )
-        asyncio.create_task(self._platform_reconnect_watcher())
+        self._track_background_task(
+            asyncio.create_task(self._platform_reconnect_watcher())
+        )
 
         logger.info("Press Ctrl+C to stop")
         
@@ -3003,7 +3015,9 @@ async def _handle_message_with_agent(self, event, source, _quick_key: str):
                 from tools.process_registry import process_registry
                 while process_registry.pending_watchers:
                     watcher = process_registry.pending_watchers.pop(0)
-                    asyncio.create_task(self._run_process_watcher(watcher))
+                    self._track_background_task(
+                        asyncio.create_task(self._run_process_watcher(watcher))
+                    )
             except Exception as e:
                 logger.error("Process watcher setup error: %s", e)
 
diff --git a/tests/gateway/test_gateway_shutdown.py b/tests/gateway/test_gateway_shutdown.py
index 439fbfdb0519b..a0ddaf37810d5 100644
--- a/tests/gateway/test_gateway_shutdown.py
+++ b/tests/gateway/test_gateway_shutdown.py
@@ -37,6 +37,32 @@ def _source(chat_id="123456", chat_type="dm"):
     )
 
 
+class _FakeProcessRegistry:
+    def __init__(self, watchers=None):
+        self.pending_watchers = list(watchers or [])
+
+    def recover_from_checkpoint(self):
+        return 0
+
+
+@pytest.mark.asyncio
+async def test_track_background_task_releases_completed_tasks():
+    runner = object.__new__(GatewayRunner)
+    runner._background_tasks = set()
+
+    async def complete_soon():
+        await asyncio.sleep(0)
+
+    task = runner._track_background_task(asyncio.create_task(complete_soon()))
+
+    assert task in runner._background_tasks
+
+    await task
+    await asyncio.sleep(0)
+
+    assert runner._background_tasks == set()
+
+
 @pytest.mark.asyncio
 async def test_cancel_background_tasks_cancels_inflight_message_processing():
     adapter = StubAdapter()
@@ -105,3 +131,41 @@ async def block_forever(_event):
     assert runner._pending_messages == {}
     assert runner._pending_approvals == {}
     assert runner._shutdown_event.is_set() is True
+
+
+@pytest.mark.asyncio
+async def test_runner_start_tracks_startup_watchers(monkeypatch, tmp_path):
+    config = GatewayConfig(
+        platforms={Platform.TELEGRAM: PlatformConfig(enabled=False, token="***")},
+        sessions_dir=tmp_path / "sessions",
+    )
+    runner = GatewayRunner(config)
+
+    blocker = asyncio.Event()
+
+    async def hold(*_args, **_kwargs):
+        await blocker.wait()
+
+    runner._send_update_notification = AsyncMock(return_value=False)
+    runner._run_process_watcher = AsyncMock(side_effect=hold)
+    runner._session_expiry_watcher = AsyncMock(side_effect=hold)
+    runner._platform_reconnect_watcher = AsyncMock(side_effect=hold)
+
+    import tools.process_registry as pr_module
+
+    monkeypatch.setattr(
+        pr_module,
+        "process_registry",
+        _FakeProcessRegistry([{"session_id": "proc_1"}]),
+    )
+
+    ok = await runner.start()
+
+    assert ok is True
+    assert len(runner._background_tasks) == 3
+
+    with patch("gateway.status.remove_pid_file"), patch("gateway.status.write_runtime_status"):
+        await runner.stop()
+
+    await asyncio.sleep(0)
+    assert runner._background_tasks == set()