webmail-client/backend/app/sync/supervisor.py

"""Live-sync supervisor (SPEC §6.3).

Runs inside the arq worker: on startup it launches one IDLE/poll loop per live
account so new mail and flag changes flow to the UI over WebSocket. Loops
self-heal (reconnect with backoff) and are cancelled on worker shutdown.

NOTE(spec): accounts that become live *after* worker startup are picked up on the
next worker restart; reacting to ``account.state_changed`` live is a later refinement.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import uuid
from typing import Any

from sqlalchemy import select

from app.db import session_scope
from app.models import Account
from app.sync.engine import live_loop, secondary_poll_loop

logger = logging.getLogger(__name__)

_RECONNECT_BACKOFF = 30.0


async def _run_account(account_id: uuid.UUID, stop: asyncio.Event) -> None:
    while not stop.is_set():
        try:
            await live_loop(account_id, "INBOX", stop)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("live loop for account %s crashed; reconnecting", account_id)
            with contextlib.suppress(TimeoutError):
                await asyncio.wait_for(stop.wait(), timeout=_RECONNECT_BACKOFF)


async def _run_secondary(account_id: uuid.UUID, stop: asyncio.Event) -> None:
    while not stop.is_set():
        try:
            await secondary_poll_loop(account_id, stop)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("secondary poll loop for %s crashed; retrying", account_id)
            with contextlib.suppress(TimeoutError):
                await asyncio.wait_for(stop.wait(), timeout=_RECONNECT_BACKOFF)


async def start_supervisor(ctx: dict[str, Any]) -> None:
    """Launch a live loop per live account (called from the worker on_startup)."""
    async with session_scope() as session:
        account_ids = list(await session.scalars(select(Account.id).where(Account.state == "live")))
    stop = asyncio.Event()
    ctx["supervisor_stop"] = stop
    tasks: list[asyncio.Task[None]] = []
    for aid in account_ids:
        tasks.append(asyncio.create_task(_run_account(aid, stop)))  # INBOX IDLE
        tasks.append(asyncio.create_task(_run_secondary(aid, stop)))  # other folders poll
    ctx["supervisor_tasks"] = tasks
    logger.info("live supervisor started for %d account(s)", len(account_ids))


async def stop_supervisor(ctx: dict[str, Any]) -> None:
    """Signal and cancel all live loops (called from the worker on_shutdown)."""
    stop = ctx.get("supervisor_stop")
    if stop is not None:
        stop.set()
    tasks = ctx.get("supervisor_tasks", [])
    for task in tasks:
        task.cancel()
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)