Skip to content

fix(gateway): track all background tasks in _background_tasks set - #11630

Closed
teknium1 wants to merge 1 commit into
mainfrom
fix/gateway-untracked-bg-tasks
Closed

fix(gateway): track all background tasks in _background_tasks set#11630
teknium1 wants to merge 1 commit into
mainfrom
fix/gateway-untracked-bg-tasks

Conversation

@teknium1

Copy link
Copy Markdown
Contributor

What this PR does (zoomed out)

Five long-running asyncio.create_task() call sites in gateway/run.py were fire-and-forget — no reference stored anywhere, not added to self._background_tasks, invisible to _stop_impl's cancel loop. They kept running (and pinning self, every cached agent, every adapter) past gateway shutdown.

The leak

Each call site looked like:

asyncio.create_task(self._session_expiry_watcher())

The bound-method coroutine holds an implicit reference to self (the GatewayRunner). That reference keeps the whole runner and everything it owns (adapters, _agent_cache, _session_db, etc.) alive until the process itself exits — even after stop() 'completed'. Gateway restarts looked clean but actually leaked the previous runner until Python exited.

Python GC could also reap the task object early under memory pressure, at which point the coroutine silently disappears (task was only held by the task variable we threw away).

Fixed sites

All in gateway/run.py:

Location Task
start() startup drain loop _run_process_watcher (per recovered watcher from crash-recovery checkpoint)
start() _session_expiry_watcher (forever loop, per runner)
start() _platform_reconnect_watcher (forever loop, per runner)
_handle_message() drain _run_process_watcher (mid-run per-process check_interval)
_schedule_update_notification_watch() _watch_update_progress

Every site now uses the pattern already established elsewhere:

_task = asyncio.create_task(...)
self._background_tasks.add(_task)
_task.add_done_callback(self._background_tasks.discard)

_update_notification_task keeps its self attribute so the idempotency check (if existing and not existing.done(): return) still works. It's now stored in both places.

_stop_impl is unchanged — its existing for _task in list(self._background_tasks): _task.cancel() loop already handles everything it sees.

Tests

tests/gateway/test_background_task_tracking.py — 7 cases:

  • Update notification task tracked in _background_tasks
  • Done callback removes the update task on natural completion
  • Idempotency: calling _schedule_update_notification_watch twice while one is running doesn't spawn a duplicate or leak into _background_tasks
  • Per-process watcher tracking (mid-run spawn pattern)
  • Session expiry watcher tracking
  • Platform reconnect watcher tracking
  • Full spawn → cancel → done-callback cleanup cycle
pytest tests/gateway/test_background_task_tracking.py         7 passed
pytest tests/gateway/test_gateway_shutdown.py etc.             66 passed

Context

Part of the memory-leak audit series. PR 2 of 5, following #11565 (bounded agent cache). Next up: SessionStore._entries pruning (PR 3).

Five long-running asyncio.create_task() call sites in gateway/run.py
were fire-and-forget — no reference stored anywhere, not added to
self._background_tasks, invisible to _stop_impl's cancel loop.

The impact:
  * Python could GC the task object prematurely under memory pressure.
  * Tasks survived stop()/restart because cancel() never reached them.
  * Each bound-method coroutine held a reference to self (GatewayRunner),
    pinning every adapter, cached AIAgent, and session object until
    the process actually exited — even though stop() had 'completed'.

Fixed sites (gateway/run.py):
  * start(): _session_expiry_watcher      — forever loop
  * start(): _platform_reconnect_watcher  — forever loop
  * start(): _run_process_watcher         — per recovered watcher
  * _handle_message(): _run_process_watcher — per mid-run watcher
  * _schedule_update_notification_watch(): _watch_update_progress

All now use the established pattern:
    _task = asyncio.create_task(...)
    self._background_tasks.add(_task)
    _task.add_done_callback(self._background_tasks.discard)

_update_notification_task keeps its self attribute too (the 'is a
watcher already running?' idempotency check in
_schedule_update_notification_watch depends on it).

Tests: 7 cases in tests/gateway/test_background_task_tracking.py
covering each tracking site, done-callback deregistration on normal
completion, idempotency of the update-notification scheduler, and
the full spawn+cancel+cleanup cycle.

No changes to _stop_impl — its existing cancel loop already handles
everything in _background_tasks correctly.
@teknium1

Copy link
Copy Markdown
Contributor Author

Skipping this one. Low ROI relative to the rest of the audit — Python GC risk is documented but rarely observed, and the forever-loop watchers already poll self._running so they exit cleanly on shutdown. Leaving the tracking pattern as-is until we see an actual leak attributable to this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant