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

"""Outbound user actions: flag changes and moves (SPEC §6.3).

The API writes the change to the DB first, then enqueues a job here to push it to
the mail server. On provider failure the job retries (arq, ×5 exponential) and, on
final failure, reverts the DB change and emits ``flags.changed`` with
``reverted=true`` so the UI can roll back.
"""

from __future__ import annotations

import logging
import uuid
from collections.abc import Iterable
from typing import Any

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.db import session_scope
from app.events import get_event_bus
from app.models import Account, Draft, Mailbox, Message, Thread
from app.sync.provider_factory import build_provider
from app.threading_.threader import denormalize_thread

logger = logging.getLogger(__name__)

# User-settable IMAP flags (SPEC §8.2 is:read/starred/answered). \Deleted/\Draft
# are managed by move/expunge and drafts, not this surface.
ALLOWED_FLAGS = frozenset({"\\Seen", "\\Flagged", "\\Answered"})
_MAX_TRIES = 5

# Folder names to create when a target special-use folder is missing on the server.
_DEFAULT_FOLDER_NAME = {"archive": "Archive", "trash": "Trash", "junk": "Spam"}

# Moves into these folders record where the message came from (restore_mailbox_id)
# so the Restore action can return it to its original folder.
_REMOVAL_FOLDERS = frozenset({"archive", "trash", "junk"})


def sanitize_flags(flags: Iterable[str]) -> frozenset[str]:
    return frozenset(f for f in flags if f in ALLOWED_FLAGS)


async def apply_flags_to_messages(
    session: AsyncSession, messages: list[Message], add: frozenset[str], remove: frozenset[str]
) -> set[uuid.UUID]:
    """Apply a flag delta to messages in the DB and re-denormalize their threads."""
    thread_ids: set[uuid.UUID] = set()
    for message in messages:
        message.flags = sorted((set(message.flags) | add) - remove)
        if message.thread_id is not None:
            thread_ids.add(message.thread_id)
    await session.flush()
    for tid in thread_ids:
        await denormalize_thread(session, tid)
    return thread_ids


# --- worker jobs -------------------------------------------------------------


async def push_flags_job(
    ctx: dict[str, Any], message_id: str, add: list[str], remove: list[str]
) -> None:
    """Push a flag change to the mail server; revert the DB on final failure."""
    mid = uuid.UUID(message_id)
    async with session_scope() as session:
        message = await session.get(Message, mid)
        if message is None:
            return
        account = await session.get(Account, message.account_id)
        mailbox = await session.get(Mailbox, message.mailbox_id)
        if account is None or mailbox is None:
            return
        user_id, uid, path = message.user_id, message.uid, mailbox.path

    provider = build_provider(account)
    try:
        await provider.push_flags(path, uid, frozenset(add), frozenset(remove))
    except Exception:
        if ctx.get("job_try", 1) >= _MAX_TRIES:
            await _revert_flags(mid, user_id, add, remove)
            return
        raise  # arq retries with backoff
    finally:
        await provider.close()


async def _revert_flags(
    message_id: uuid.UUID, user_id: uuid.UUID, add: list[str], remove: list[str]
) -> None:
    async with session_scope() as session:
        message = await session.get(Message, message_id)
        if message is None:
            return
        # Undo the change: remove what we added, add back what we removed.
        message.flags = sorted((set(message.flags) - set(add)) | set(remove))
        await session.flush()
        if message.thread_id is not None:
            await denormalize_thread(session, message.thread_id)
        flags = list(message.flags)
    await get_event_bus().publish(
        "flags.changed",
        {"user_id": str(user_id), "message_id": str(message_id), "flags": flags, "reverted": True},
    )
    logger.warning("reverted flag change for message %s after push failures", message_id)


async def move_job(
    ctx: dict[str, Any], message_id: str, dst_mailbox_id: str, src_mailbox_id: str
) -> None:
    """Move a message on the server; update its uid, or revert the DB on failure."""
    mid = uuid.UUID(message_id)
    dst_id = uuid.UUID(dst_mailbox_id)
    async with session_scope() as session:
        message = await session.get(Message, mid)
        if message is None:
            return
        account = await session.get(Account, message.account_id)
        dst = await session.get(Mailbox, dst_id)
        src = await session.get(Mailbox, uuid.UUID(src_mailbox_id))
        if account is None or dst is None or src is None:
            return
        src_path = src.path
        uid, dst_path, user_id = message.uid, dst.path, message.user_id

    provider = build_provider(account)
    try:
        new_uid = await provider.move(uid, src_path, dst_path)
    except Exception:
        if ctx.get("job_try", 1) >= _MAX_TRIES:
            logger.warning("move failed for message %s; leaving optimistic state", mid)
            return
        raise
    finally:
        await provider.close()

    async with session_scope() as session:
        message = await session.get(Message, mid)
        if message is not None:
            message.uid = new_uid
            message.mailbox_id = dst_id
    await get_event_bus().publish(
        "message.changed",
        {"user_id": str(user_id), "message_id": str(mid), "mailbox_id": str(dst_id)},
    )


async def move_thread_to_folder_job(ctx: dict[str, Any], thread_id: str, special_use: str) -> None:
    """Move every message in a thread to each account's special-use folder.

    Creates the target folder (e.g. Archive) if the account doesn't have one.
    Used by the Archive and Delete (-> Trash) thread actions.
    """
    tid = uuid.UUID(thread_id)
    async with session_scope() as session:
        messages = list(await session.scalars(select(Message).where(Message.thread_id == tid)))
        if not messages:
            return
        user_id = messages[0].user_id
        account_ids = list({m.account_id for m in messages})
        accounts = {
            a.id: a
            for a in await session.scalars(select(Account).where(Account.id.in_(account_ids)))
        }
        mailboxes = list(
            await session.scalars(select(Mailbox).where(Mailbox.account_id.in_(account_ids)))
        )
        path_by_id = {mb.id: mb.path for mb in mailboxes}
        target_by_account = {mb.account_id: mb for mb in mailboxes if mb.special_use == special_use}
        target_id_path = {mb.account_id: (mb.id, mb.path) for mb in target_by_account.values()}
        snapshot = [(m.id, m.account_id, m.mailbox_id, m.uid) for m in messages]

    # (message_id, account_id, target_path, new_uid) applied after the network moves.
    moves: list[tuple[uuid.UUID, uuid.UUID, str, int]] = []
    to_create: list[tuple[uuid.UUID, str, int]] = []  # account_id, path, uidvalidity

    for account_id in account_ids:
        provider = build_provider(accounts[account_id])
        try:
            if account_id in target_id_path:
                target_id, target_path = target_id_path[account_id]
            else:
                target_id, target_path = None, _DEFAULT_FOLDER_NAME.get(special_use, special_use)
                await provider.create_mailbox(target_path)
                remote = await provider.list_mailboxes()
                uidv = next((r.uidvalidity for r in remote if r.path == target_path), 0)
                to_create.append((account_id, target_path, uidv))
            for mid, m_account, m_mailbox, m_uid in snapshot:
                if m_account != account_id or m_mailbox == target_id:
                    continue
                src_path = path_by_id.get(m_mailbox, "INBOX")
                try:
                    new_uid = await provider.move(m_uid, src_path, target_path)
                except Exception:
                    logger.exception("move failed for message %s -> %s", mid, target_path)
                    continue
                moves.append((mid, account_id, target_path, new_uid))
        finally:
            await provider.close()

    async with session_scope() as session:
        created: dict[tuple[uuid.UUID, str], uuid.UUID] = {}
        for account_id, path, uidv in to_create:
            mailbox = Mailbox(
                account_id=account_id,
                user_id=user_id,
                path=path,
                special_use=special_use,
                uidvalidity=uidv,
            )
            session.add(mailbox)
            await session.flush()
            created[(account_id, path)] = mailbox.id
        for mid, account_id, target_path, new_uid in moves:
            resolved = target_id_path.get(account_id, (None, ""))[0] or created.get(
                (account_id, target_path)
            )
            message = await session.get(Message, mid)
            if message is not None and resolved is not None:
                # Remember where it came from so Restore returns it there, not the Inbox.
                if special_use in _REMOVAL_FOLDERS and message.mailbox_id != resolved:
                    message.restore_mailbox_id = message.mailbox_id
                message.mailbox_id = resolved
                message.uid = new_uid

    await get_event_bus().publish(
        "thread.updated", {"user_id": str(user_id), "thread_id": str(tid)}
    )


async def restore_thread_job(ctx: dict[str, Any], thread_id: str) -> None:
    """Move every message in a thread back to its original folder (Restore action).

    Each message returns to ``restore_mailbox_id`` (the folder it was in before it
    was archived/trashed) — e.g. a deleted Sent message goes back to Sent. Messages
    with no recorded origin fall back to the account's Inbox.
    """
    tid = uuid.UUID(thread_id)
    async with session_scope() as session:
        messages = list(await session.scalars(select(Message).where(Message.thread_id == tid)))
        if not messages:
            return
        user_id = messages[0].user_id
        account_ids = list({m.account_id for m in messages})
        accounts = {
            a.id: a
            for a in await session.scalars(select(Account).where(Account.id.in_(account_ids)))
        }
        mailboxes = list(
            await session.scalars(select(Mailbox).where(Mailbox.account_id.in_(account_ids)))
        )
        path_by_id = {mb.id: mb.path for mb in mailboxes}
        inbox_by_account = {
            mb.account_id: mb.id for mb in mailboxes if mb.special_use == "inbox"
        }
        # (message_id, account_id, source_path, target_id) — resolve the target now.
        plan: list[tuple[uuid.UUID, uuid.UUID, str, uuid.UUID, int]] = []
        for m in messages:
            target_id = m.restore_mailbox_id
            if target_id is None or target_id not in path_by_id:
                target_id = inbox_by_account.get(m.account_id)
            if target_id is None or target_id == m.mailbox_id:
                continue  # nowhere to restore to, or already there
            plan.append(
                (m.id, m.account_id, path_by_id.get(m.mailbox_id, "INBOX"), target_id, m.uid)
            )

    # (message_id, target_id, new_uid) applied after the network moves.
    moves: list[tuple[uuid.UUID, uuid.UUID, int]] = []
    for account_id in account_ids:
        provider = build_provider(accounts[account_id])
        try:
            for mid, m_account, src_path, target_id, uid in plan:
                if m_account != account_id:
                    continue
                target_path = path_by_id[target_id]
                try:
                    new_uid = await provider.move(uid, src_path, target_path)
                except Exception:
                    logger.exception("restore move failed for message %s -> %s", mid, target_path)
                    continue
                moves.append((mid, target_id, new_uid))
        finally:
            await provider.close()

    async with session_scope() as session:
        for mid, target_id, new_uid in moves:
            message = await session.get(Message, mid)
            if message is not None:
                message.mailbox_id = target_id
                message.uid = new_uid
                message.restore_mailbox_id = None

    await get_event_bus().publish(
        "thread.updated", {"user_id": str(user_id), "thread_id": str(tid)}
    )


async def delete_thread_job(ctx: dict[str, Any], thread_id: str) -> None:
    """Permanently delete every message in a thread (server EXPUNGE + DB rows).

    Used by the Trash folder's Delete action. Best-effort per message on the
    server; the DB rows (and cascaded attachments/body parts) are always removed.
    """
    tid = uuid.UUID(thread_id)
    async with session_scope() as session:
        messages = list(await session.scalars(select(Message).where(Message.thread_id == tid)))
        if not messages:
            return
        user_id = messages[0].user_id
        account_ids = list({m.account_id for m in messages})
        accounts = {
            a.id: a
            for a in await session.scalars(select(Account).where(Account.id.in_(account_ids)))
        }
        mailboxes = list(
            await session.scalars(select(Mailbox).where(Mailbox.account_id.in_(account_ids)))
        )
        path_by_id = {mb.id: mb.path for mb in mailboxes}
        # (account_id, mailbox_path) -> [uids] to expunge on the server.
        by_box: dict[tuple[uuid.UUID, str], list[int]] = {}
        for m in messages:
            by_box.setdefault((m.account_id, path_by_id.get(m.mailbox_id, "INBOX")), []).append(
                m.uid
            )

    for account_id in account_ids:
        provider = build_provider(accounts[account_id])
        try:
            for (acct, path), uids in by_box.items():
                if acct != account_id:
                    continue
                try:
                    await provider.delete_messages(path, uids)
                except Exception:
                    logger.exception("expunge failed for %s in %s", uids, path)
        finally:
            await provider.close()

    async with session_scope() as session:
        msgs = list(await session.scalars(select(Message).where(Message.thread_id == tid)))
        msg_ids = [m.id for m in msgs]
        # Drop draft references to the thread/messages (FKs lack ON DELETE CASCADE).
        if msg_ids:
            await session.execute(
                update(Draft)
                .where(Draft.reply_to_message_id.in_(msg_ids))
                .values(reply_to_message_id=None)
            )
        await session.execute(update(Draft).where(Draft.thread_id == tid).values(thread_id=None))
        for m in msgs:
            await session.delete(m)
        thread = await session.get(Thread, tid)
        if thread is not None:
            await session.delete(thread)

    await get_event_bus().publish(
        "thread.updated", {"user_id": str(user_id), "thread_id": str(tid)}
    )