"""Sync engine: bootstrap mailboxes, backfill, and the live IDLE/poll loop (SPEC §6.3).""" from __future__ import annotations import asyncio import contextlib import logging import random import uuid from dataclasses import dataclass, field from datetime import date, timedelta from typing import Any from sqlalchemy import delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from app.config import get_settings from app.events import get_event_bus from app.indexer.indexer import PersistResult, persist_message from app.models import Account, Draft, Mailbox, Message from app.providers.base import MailboxDelta, ProviderNotSupported from app.providers.imap_smtp import ImapSmtpProvider from app.search.factory import get_search_backend from app.sync.folder_names import infer_special_use from app.sync.provider_factory import build_provider logger = logging.getLogger(__name__) # Backfill priority: inbox, then sent/drafts, then everything else (SPEC §6.3). _PRIORITY = {"inbox": 0, "sent": 1, "drafts": 2} # Sync-window import strategy (SPEC §12 M2). "all" imports the most recent # _RECENT_DAYS first (so the account is usable fast), then the older remainder. _RECENT_DAYS = 30 _WINDOW_DAYS = {"3months": 90} async def _load_account(session: AsyncSession, account_id: uuid.UUID) -> Account: account = await session.get(Account, account_id) if account is None: raise ValueError(f"account {account_id} not found") return account async def _set_state(account_id: uuid.UUID, state: str, detail: str = "") -> None: from app.db import session_scope async with session_scope() as session: await session.execute( update(Account).where(Account.id == account_id).values(state=state, state_detail=detail) ) # Keep the sync:live registry in step with the account's state so the sync # coordinator claims it once live and drops it otherwise (refactor §2). from app.redis_client import get_redis from app.sync import leases redis = get_redis() if state == "live": await leases.mark_live(redis, account_id) else: await leases.mark_not_live(redis, account_id) bus = get_event_bus() account = None async with session_scope() as session: account = await session.get(Account, account_id) if account is not None: await bus.publish( "account.state_changed", {"user_id": str(account.user_id), "account_id": str(account_id), "state": state}, ) async def bootstrap_mailboxes( session: AsyncSession, provider: ImapSmtpProvider, account: Account ) -> list[Mailbox]: """List remote mailboxes and upsert local ``mailboxes`` rows (SPEC §6.3).""" remote = await provider.list_mailboxes() result: list[Mailbox] = [] for rm in remote: special = rm.special_use or infer_special_use(rm.path) existing = await session.scalar( select(Mailbox).where(Mailbox.account_id == account.id, Mailbox.path == rm.path) ) if existing is None: mailbox = Mailbox( account_id=account.id, user_id=account.user_id, path=rm.path, special_use=special, uidvalidity=rm.uidvalidity, ) session.add(mailbox) else: existing.special_use = special existing.uidvalidity = rm.uidvalidity mailbox = existing result.append(mailbox) await session.flush() return result def _mailbox_order(mailboxes: list[Mailbox]) -> list[Mailbox]: return sorted(mailboxes, key=lambda m: _PRIORITY.get(m.special_use or "", 99)) async def _index_and_notify(results: list[PersistResult], account: Account) -> None: """Index each persisted message and publish its ``message.new`` event.""" backend = get_search_backend() bus = get_event_bus() for r in results: await backend.index(r.search_doc) await bus.publish( "message.new", { "user_id": str(account.user_id), "account_id": str(account.id), "message_id": str(r.message.id), "thread_id": str(r.message.thread_id), }, ) async def _emit_batch( results: list[PersistResult], account: Account, mailbox_path: str, done: int ) -> None: await _index_and_notify(results, account) await get_event_bus().publish( "sync.progress", { "user_id": str(account.user_id), "account_id": str(account.id), "mailbox": mailbox_path, "done": done, }, ) @dataclass class _BackfillProgress: """Account-level backfill counter: persists ``backfill_done`` and emits ``sync.progress`` with the running total so the UI can show a progress bar.""" account: Account total: int done: int = 0 _mailboxes: set[str] = field(default_factory=set) async def record(self, results: list[PersistResult], mailbox_path: str) -> None: from app.db import session_scope await _index_and_notify(results, self.account) self.done += len(results) self._mailboxes.add(mailbox_path) async with session_scope() as session: await session.execute( update(Account) .where(Account.id == self.account.id) .values(backfill_done=self.done) ) await get_event_bus().publish( "sync.progress", { "user_id": str(self.account.user_id), "account_id": str(self.account.id), "mailbox": mailbox_path, "done": self.done, "total": self.total, }, ) async def backfill_mailbox( provider: ImapSmtpProvider, account: Account, mailbox: Mailbox, progress: _BackfillProgress | None = None, *, since: date | None = None, before: date | None = None, ) -> int: """Backfill one mailbox newest-first, indexing inline per batch (FR-30). ``since``/``before`` bound the import by date (sync window, SPEC §12 M2). When a ``progress`` tracker is given, batches count toward the account-level backfill total; otherwise each batch emits a standalone ``sync.progress``. """ from app.db import session_scope account_id, user_id, mailbox_id = account.id, account.user_id, mailbox.id special_use, path = mailbox.special_use, mailbox.path done = 0 async for batch in provider.backfill(path, batch_size=200, since=since, before=before): results: list[PersistResult] = [] max_uid = 0 async with session_scope() as session: for remote in batch: r = await persist_message( session, user_id=user_id, account_id=account_id, mailbox_id=mailbox_id, mailbox_special_use=special_use, remote=remote, ) if r is not None: results.append(r) max_uid = max(max_uid, remote.uid) await session.execute( update(Mailbox) .where(Mailbox.id == mailbox_id, Mailbox.last_seen_uid < max_uid) .values(last_seen_uid=max_uid) ) done += len(results) if progress is not None: await progress.record(results, path) else: await _emit_batch(results, account, path, done) return done async def _set_backfill_counts(account_id: uuid.UUID, *, total: int, done: int) -> None: from app.db import session_scope async with session_scope() as session: await session.execute( update(Account) .where(Account.id == account_id) .values(backfill_total=total, backfill_done=done) ) async def _sync_new_only(provider: ImapSmtpProvider, mailboxes: list[Mailbox]) -> None: """``new`` window: import no history — just move each mailbox watermark to the current tip so only mail arriving after setup is pulled (SPEC §12 M2).""" from app.db import session_scope for mailbox in mailboxes: tip = await provider.newest_uid(mailbox.path) async with session_scope() as session: await session.execute( update(Mailbox) .where(Mailbox.id == mailbox.id, Mailbox.last_seen_uid < tip) .values(last_seen_uid=tip) ) async def bootstrap_and_backfill(account_id: uuid.UUID) -> int: """Full initial sync for an account: mailboxes + windowed backfill; state=live.""" from app.db import session_scope async with session_scope() as session: account = await _load_account(session, account_id) provider = build_provider(account) mailboxes = _mailbox_order(await bootstrap_mailboxes(session, provider, account)) # Detach: reuse account/mailboxes values after the session closes. account_obj = account window = account.sync_window mailbox_list = list(mailboxes) try: if window == "new": # No backfill: set watermarks to the tip, go live immediately. await _sync_new_only(provider, mailbox_list) await _set_backfill_counts(account_id, total=0, done=0) await _set_state(account_id, "live") return 0 since = ( date.today() - timedelta(days=_WINDOW_DAYS[window]) if window in _WINDOW_DAYS else None ) # Progress denominator: count what we're about to import up front so the # bar has a stable total (SPEC §12 M2 backfill-progress UX). total = 0 for mailbox in mailbox_list: total += await provider.message_count(mailbox.path, since=since) await _set_backfill_counts(account_id, total=total, done=0) progress = _BackfillProgress(account_obj, total) recent_cutoff = date.today() - timedelta(days=_RECENT_DAYS) count = 0 for mailbox in mailbox_list: if since is not None: # 3months: single windowed pass count += await backfill_mailbox( provider, account_obj, mailbox, progress, since=since ) else: # all: recent window first (usable fast), then the remainder count += await backfill_mailbox( provider, account_obj, mailbox, progress, since=recent_cutoff ) count += await backfill_mailbox( provider, account_obj, mailbox, progress, before=recent_cutoff ) await _set_state(account_id, "live") except Exception as exc: # noqa: BLE001 - surface provider errors as account state logger.exception("backfill failed for account %s", account_id) await _set_state(account_id, "error", str(exc)[:500]) raise finally: await provider.close() return count async def apply_delta( provider: ImapSmtpProvider, account: Account, mailbox: Mailbox, delta: MailboxDelta ) -> None: """Apply a poll delta: new messages, flag changes, uidvalidity resets (SPEC §6.3).""" from app.db import session_scope bus = get_event_bus() if delta.new_uidvalidity is not None: # UIDVALIDITY changed: wipe this mailbox's rows and re-backfill (SPEC §6.3). async with session_scope() as session: await session.execute(delete(Message).where(Message.mailbox_id == mailbox.id)) await session.execute( update(Mailbox) .where(Mailbox.id == mailbox.id) .values(uidvalidity=delta.new_uidvalidity, last_seen_uid=0) ) await backfill_mailbox(provider, account, mailbox) return results: list[PersistResult] = [] max_uid = mailbox.last_seen_uid gone_ids: list[uuid.UUID] = [] async with session_scope() as session: for remote in delta.new: r = await persist_message( session, user_id=account.user_id, account_id=account.id, mailbox_id=mailbox.id, mailbox_special_use=mailbox.special_use, remote=remote, ) if r is not None: results.append(r) max_uid = max(max_uid, remote.uid) for change in delta.flag_changes: await session.execute( update(Message) .where(Message.mailbox_id == mailbox.id, Message.uid == change.uid) .values(flags=sorted(change.flags)) ) if delta.expunged_uids: # Server-side expunges (QRESYNC/VANISHED, SPEC §6.3). Capture the row ids # first so we can drop them from the search index after the delete. gone_ids = list( await session.scalars( select(Message.id).where( Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids) ) ) ) await session.execute( delete(Message).where( Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids) ) ) # A local draft whose mirrored copy vanished reverts to local-only. await session.execute( update(Draft) .where(Draft.account_id == account.id, Draft.imap_uid.in_(delta.expunged_uids)) .values(imap_uid=None) ) updates = {"last_seen_uid": max_uid} if delta.new_highestmodseq is not None: updates["highestmodseq"] = delta.new_highestmodseq await session.execute(update(Mailbox).where(Mailbox.id == mailbox.id).values(**updates)) await _emit_batch(results, account, mailbox.path, len(results)) if gone_ids: # Drop expunged messages from the search index (PG keeps its tsvector on the # row, so this is a no-op there; Meili needs the explicit removal) and tell # the UI to refresh, mirroring the drafts-reconcile notification. backend = get_search_backend() for mid in gone_ids: await backend.delete(mid) await bus.publish( "message.deleted", { "user_id": str(account.user_id), "account_id": str(account.id), "mailbox": mailbox.path, }, ) for change in delta.flag_changes: await bus.publish( "flags.changed", { "user_id": str(account.user_id), "account_id": str(account.id), "mailbox": mailbox.path, "uid": change.uid, "flags": sorted(change.flags), }, ) # Expunge reconciliation for Drafts. The mirror mechanism (and other clients) # expunge server copies, but a CONDSTORE poll doesn't report vanished UIDs, so # expunged drafts would linger in the DB and resurface as phantom "remote" # drafts. Drafts folders are small, so a full UID reconcile per poll is cheap. if mailbox.special_use == "drafts": await _reconcile_drafts(provider, account, mailbox) async def _reconcile_drafts( provider: ImapSmtpProvider, account: Account, mailbox: Mailbox ) -> None: """Drop DB rows for drafts that no longer exist on the server (were expunged).""" from app.db import session_scope try: server_uids = await provider.list_uids(mailbox.path) except Exception: logger.exception("drafts reconcile: list_uids failed for %s", mailbox.path) return async with session_scope() as session: db_uids = set( await session.scalars(select(Message.uid).where(Message.mailbox_id == mailbox.id)) ) gone = db_uids - server_uids if not gone: return await session.execute( delete(Message).where(Message.mailbox_id == mailbox.id, Message.uid.in_(gone)) ) # A local draft whose mirrored copy vanished becomes local-only again. await session.execute( update(Draft) .where(Draft.account_id == account.id, Draft.imap_uid.in_(gone)) .values(imap_uid=None) ) logger.info("drafts reconcile: removed %d expunged row(s) in %s", len(gone), mailbox.path) await get_event_bus().publish( "message.deleted", {"user_id": str(account.user_id), "account_id": str(account.id), "mailbox": mailbox.path}, ) async def live_loop( account_id: uuid.UUID, mailbox_path: str = "INBOX", stop: asyncio.Event | None = None ) -> None: """Maintain a live IDLE/poll loop for one mailbox (SPEC §6.3).""" from app.db import session_scope stop = stop or asyncio.Event() idle_fallback = get_settings().poll_warm async with session_scope() as session: account = await _load_account(session, account_id) account_obj = account provider = build_provider(account_obj) try: while not stop.is_set(): try: await provider.idle(mailbox_path) except ProviderNotSupported: await asyncio.sleep(idle_fallback) if stop.is_set(): break async with session_scope() as session: mailbox = await session.scalar( select(Mailbox).where( Mailbox.account_id == account_id, Mailbox.path == mailbox_path ) ) mailbox_obj = mailbox if mailbox_obj is None: break delta = await provider.poll( mailbox_path, last_seen_uid=mailbox_obj.last_seen_uid, highestmodseq=mailbox_obj.highestmodseq, ) await apply_delta(provider, account_obj, mailbox_obj, delta) finally: await provider.close() async def poll_folders_once(account_id: uuid.UUID, *, include_inbox: bool) -> None: """One CONDSTORE poll pass over an account's special-use folders. ``include_inbox`` also polls INBOX (used by the manual Refresh); the periodic secondary loop leaves INBOX to its dedicated IDLE loop. Custom (non special-use) folders are not polled. Drafts are polled so drafts made in other clients appear (persisted un-threaded). """ from app.db import session_scope async with session_scope() as session: account = await _load_account(session, account_id) account_obj = account stmt = select(Mailbox.id).where( Mailbox.account_id == account_id, Mailbox.special_use.isnot(None) ) if not include_inbox: stmt = stmt.where(Mailbox.special_use != "inbox") mailbox_ids = list(await session.scalars(stmt)) if not mailbox_ids: return provider = build_provider(account_obj) try: for mailbox_id in mailbox_ids: async with session_scope() as session: mailbox = await session.get(Mailbox, mailbox_id) mailbox_obj = mailbox if mailbox_obj is None: continue try: delta = await provider.poll( mailbox_obj.path, last_seen_uid=mailbox_obj.last_seen_uid, highestmodseq=mailbox_obj.highestmodseq, ) await apply_delta(provider, account_obj, mailbox_obj, delta) except Exception: logger.exception("poll failed: account=%s mailbox=%s", account_id, mailbox_obj.path) finally: await provider.close() async def _poll_secondary_once(account_id: uuid.UUID) -> None: await poll_folders_once(account_id, include_inbox=False) async def refresh_account_job(ctx: dict[str, Any], account_id: str) -> None: """Poll every special-use folder (incl. INBOX) once — the manual Refresh action.""" await poll_folders_once(uuid.UUID(account_id), include_inbox=True) async def secondary_poll_loop( account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None ) -> None: """Periodically CONDSTORE-poll non-INBOX folders (Sent/Drafts/Archive/Trash/Junk). INBOX has its own IDLE loop (real-time). This keeps the other folders in step with server-side changes (e.g. Gmail-filed sent mail) and changes made by other clients. CONDSTORE/QRESYNC makes each pass a cheap MODSEQ delta. Used by hot-tier runners alongside the INBOX IDLE loop. """ stop = stop or asyncio.Event() period = interval if interval is not None else get_settings().poll_warm first = True while not stop.is_set(): # Wait first (initial state was covered by backfill). Spread the fleet: the # first wait is a random point in the period, later waits jitter ±10%, so # accounts claimed together don't poll in lockstep (thundering herd). wait = random.uniform(0, period) if first else period * random.uniform(0.9, 1.1) first = False with contextlib.suppress(TimeoutError): await asyncio.wait_for(stop.wait(), timeout=wait) if stop.is_set(): break await _poll_secondary_once(account_id) async def poll_all_loop( account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None ) -> None: """Transient CONDSTORE poll of INBOX + special folders every ``interval`` seconds. The warm/cold tiers use this instead of a persistent IDLE connection: no socket is held between passes. Correctness matches the hot tier — only latency differs. """ stop = stop or asyncio.Event() first = True while not stop.is_set(): # Spread the fleet: first poll at a random point in the interval, then # jitter ±10% so a batch claimed together doesn't poll in lockstep. wait = random.uniform(0, interval) if first else interval * random.uniform(0.9, 1.1) first = False with contextlib.suppress(TimeoutError): await asyncio.wait_for(stop.wait(), timeout=wait) if stop.is_set(): break await poll_folders_once(account_id, include_inbox=True)