Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions gateway/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please port this helper across the current watcher surface, not only the four historical callsites. Current start() has additional untracked Kanban, handoff, async-delegation, scale-to-zero, and drain-control watchers at gateway/run.py:7302-7383, plus a post-turn process-watcher drain at gateway/run.py:11727-11740.

self._background_tasks.add(task)
task.add_done_callback(self._background_tasks.discard)
return task




Expand Down Expand Up @@ -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:
Expand All @@ -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")

Expand Down Expand Up @@ -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)

Expand Down
64 changes: 64 additions & 0 deletions tests/gateway/test_gateway_shutdown.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Loading