"""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, ProviderAuthError, 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.health import AccountHealth, is_temporary from app.sync.provider_factory import build_provider from app.threading_.threader import denormalize_thread logger = logging.getLogger(__name__) # Backfill priority: inbox, then sent/drafts, then everything else (SPEC §6.3). # archive (Gmail All Mail) sorts last so its dedupe pass runs after the system folders. _PRIORITY = {"inbox": 0, "sent": 1, "drafts": 2} # Gmail sync (SPEC: Gmail sync): the only folders backfilled — All Mail (archive) is # the complete store, the rest are the system folders needed for their views. Every # custom-label folder is skipped; labels come from X-GM-LABELS instead. _GMAIL_SYNC_FOLDERS = {"inbox", "sent", "drafts", "trash", "junk", "archive"} # Progress denominator for Gmail: All Mail already contains INBOX/Sent/Drafts/labels, # so only these three are counted (Trash/Spam sit outside All Mail). _GMAIL_DENOM_FOLDERS = {"archive", "trash", "junk"} # 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) ) if state == "live": # (re)connected and synced: no sync problem any more await session.execute( update(Account) .where(Account.id == account_id) .values(sync_error=None, sync_error_since=None) ) # 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) payload: dict[str, object] = { "user_id": str(account.user_id), "account_id": str(account.id), "message_id": str(r.message.id), "thread_id": str(r.message.thread_id), } # Gmail custom labels ride along so the labels plugin can map them to chips # (SPEC: Gmail sync). Only present on Gmail accounts with labelled messages. if r.gmail_labels: payload["gmail_labels"] = list(r.gmail_labels) await bus.publish("message.new", payload) 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, dedupe_account_scope: bool = False, ) -> 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``. ``dedupe_account_scope`` drops messages already synced under another mailbox of the account (Gmail All Mail pass — SPEC: Gmail sync / no-duplication). """ 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, dedupe_account_scope=dedupe_account_scope, ) 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) # Gmail: sync only the system folders + All Mail and read each message's # labels (X-GM-LABELS) instead of visiting every label folder. This avoids # the INBOX/All-Mail/label duplication and lets labels map to chips lazily, # tied to the sync window (SPEC: Gmail sync). gmail = provider.gmail_labels_supported if gmail: mailbox_list = [m for m in mailbox_list if m.special_use in _GMAIL_SYNC_FOLDERS] # Spam is opt-in: skip the Junk folder unless the account enabled it (§6.3). if not account.sync_spam: mailbox_list = [m for m in mailbox_list if m.special_use != "junk"] 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). For Gmail only # All Mail + Trash + Spam are counted — INBOX/Sent/Drafts are subsets of All # Mail, so counting them too would inflate the total past the deduped `done`. total = 0 for mailbox in mailbox_list: if gmail and mailbox.special_use not in _GMAIL_DENOM_FOLDERS: continue 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: # Dedupe the Gmail All Mail pass against messages already stored under a # system folder (INBOX/Sent/…), which the priority order backfills first. dedupe = gmail and mailbox.special_use == "archive" if since is not None: # 3months: single windowed pass count += await backfill_mailbox( provider, account_obj, mailbox, progress, since=since, dedupe_account_scope=dedupe, ) else: # all: recent window first (usable fast), then the remainder count += await backfill_mailbox( provider, account_obj, mailbox, progress, since=recent_cutoff, dedupe_account_scope=dedupe, ) count += await backfill_mailbox( provider, account_obj, mailbox, progress, before=recent_cutoff, dedupe_account_scope=dedupe, ) 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: # Their threads are recounted: a message not brought back by the # re-import (gone meanwhile) must not stay counted. wiped_threads: set[uuid.UUID] = { tid for tid in await session.scalars( select(Message.thread_id).where(Message.mailbox_id == mailbox.id) ) if tid is not None } await session.execute(delete(Message).where(Message.mailbox_id == mailbox.id)) for tid in sorted(wiped_threads): await denormalize_thread(session, tid) 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) # Threads whose denormalized counters this batch invalidates. A flag set # on the server (read on a phone, filed by a server-side rule) and an # expunge both change what `threads.unread_count` should say, and neither # goes through apply_flags_to_messages — so without this the row keeps its # old count and the list shows a message as unread that the reading pane # shows as read. touched_threads: set[uuid.UUID] = set() 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.flag_changes: touched_threads.update( await session.scalars( select(Message.thread_id).where( Message.mailbox_id == mailbox.id, Message.uid.in_([c.uid for c in delta.flag_changes]), Message.thread_id.is_not(None), ) ) ) 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 = ( await session.execute( select(Message.id, Message.thread_id).where( Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids) ) ) ).all() gone_ids = [row[0] for row in gone] touched_threads.update(row[1] for row in gone if row[1] is not None) 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) ) # After the deletes, so a thread left with no messages is removed rather # than recounted (denormalize_thread handles that case). for tid in sorted(touched_threads): # one lock order (denormalize_thread) await denormalize_thread(session, tid) 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 results and mailbox.special_use == "inbox": # New mail while live (backfills never come through here): tell the # user's devices, a few seconds later (app/push/service.py). from app.push.service import enqueue_new_mail await enqueue_new_mail(account.user_id, [r.message.id for r in 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, health: AccountHealth | 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: await provider.verify() # signed in: the account is reachable if health is not None: await health.ok() first = True while not stop.is_set(): # Catch up before each IDLE, not only after it: IDLE reports what # arrives from now on, so mail that came while this loop was not # running (a new account's first minutes, a reconnect, a worker # restart) would otherwise wait for the next new message or the # IDLE timeout (10 min) to show. 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 not None: 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) elif not first: break # the folder is gone # else: the account's first sync has not recorded the folder yet. first = False if stop.is_set(): break try: await provider.idle(mailbox_path) except ProviderNotSupported: await asyncio.sleep(idle_fallback) finally: await provider.close() async def poll_folders_once( account_id: uuid.UUID, *, include_inbox: bool, health: AccountHealth | None = None ) -> 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). ``archive`` (Gmail All Mail) is never live-polled — it is huge and cold, and backfill already covers it. """ 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), Mailbox.special_use != "archive", ) if not include_inbox: stmt = stmt.where(Mailbox.special_use != "inbox") if not account.sync_spam: # Spam is opt-in (§6.3) stmt = stmt.where(Mailbox.special_use != "junk") 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 as exc: # Server gone, or the sign-in refused: the other folders would only # fail the same way (20 s each). End the pass; the caller decides. if is_temporary(exc) or isinstance(exc, ProviderAuthError): raise logger.exception("poll failed: account=%s mailbox=%s", account_id, mailbox_obj.path) if health is not None: await health.ok() finally: await provider.close() async def _poll_secondary_once(account_id: uuid.UUID, health: AccountHealth | None = None) -> None: await poll_folders_once(account_id, include_inbox=False, health=health) 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, health: AccountHealth | 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, health) async def poll_all_loop( account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None, health: AccountHealth | 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, health=health)