webmail-client/backend/app/threading_/threader.py

"""JWZ threading — incremental per-message algorithm (SPEC §7.2, §7.3).

Cross-account by design: candidate lookup is scoped by ``user_id`` only (I-3).
Called after the new :class:`Message` row is flushed (so it has an ``id`` and its
normalized id columns); assigns ``thread_id`` and recomputes the thread's
denormalized fields.
"""

from __future__ import annotations

import uuid
from datetime import timedelta

from sqlalchemy import ColumnElement, delete, exists, func, literal, select, update
from sqlalchemy import any_ as sa_any
from sqlalchemy import or_ as sa_or
from sqlalchemy.ext.asyncio import AsyncSession

from app.models import Draft, Message, Thread
from app.threading_.normalize import is_synthetic, normalize_subject

_SUBJECT_FALLBACK_WINDOW = timedelta(days=7)


def _recipient_emails(message: Message) -> set[str]:
    """Recipient (To/Cc) emails of a message, lowercased.

    Subject-fallback joining keys on recipients rather than all participants so a
    bulk sender that reuses one subject to many recipients (e.g. a newsletter or
    a static welcome mail) does not collapse unrelated deliveries into one thread.
    A genuine reply is always addressed to someone already in the thread.
    """
    emails: set[str] = set()
    for field in (message.to_addr, message.cc_addr):
        for addr in field:
            email = addr.get("email")
            if email:
                emails.add(email.lower())
    return emails


def _parent_refs(message: Message) -> list[str]:
    """Non-synthetic parent ids this message explicitly declares (References +
    In-Reply-To). A non-empty result means the message states its parentage, so it
    must not be subject-merged into an unrelated conversation (§7.3)."""
    refs: list[str] = list(message.references_hdr)
    if message.in_reply_to:
        refs.append(message.in_reply_to)
    out: list[str] = []
    for r in refs:
        if r and not is_synthetic(r) and r not in out:
            out.append(r)
    return out


def _candidate_ids(message: Message) -> list[str]:
    candidates: list[str] = list(message.references_hdr)
    if message.in_reply_to:
        candidates.append(message.in_reply_to)
    # Own id may let an earlier reply (that referenced us) join when we arrive late.
    if message.message_id_hdr and not is_synthetic(message.message_id_hdr):
        candidates.append(message.message_id_hdr)
    # Dedupe, drop synthetic ids (never used to join, §7.1).
    seen: list[str] = []
    for c in candidates:
        if c and not is_synthetic(c) and c not in seen:
            seen.append(c)
    return seen


async def _find_by_references(
    session: AsyncSession,
    user_id: uuid.UUID,
    message_id: uuid.UUID,
    candidates: list[str],
    parent_refs: list[str],
    own_id: str | None,
) -> list[uuid.UUID]:
    """Threads joined by reference (SPEC §7.2.2).

    Matches bidirectionally: existing messages whose id is one of my candidate ids
    (I reference them / duplicate id), existing messages that reference *my* id
    (late-arriving parent), OR existing messages that share one of my declared
    parents — a sibling reply to the same original, even when neither the parent
    nor a mutual reference is present. Synthetic ids are never used to join (§7.1).
    """
    conditions: list[ColumnElement[bool]] = []
    if candidates:
        conditions.append(Message.message_id_hdr.in_(candidates))
    if parent_refs:
        # Siblings: an existing message that references, or is a reply to, one of
        # my parents. This groups replies to a common (possibly absent) original
        # without letting the subject fallback merge replies to *different* ones.
        conditions.append(Message.references_hdr.overlap(parent_refs))
        conditions.append(Message.in_reply_to.in_(parent_refs))
    if own_id is not None and not is_synthetic(own_id):
        # own_id = ANY(references_hdr): an existing message that references me.
        conditions.append(literal(own_id) == sa_any(Message.references_hdr))
        conditions.append(Message.in_reply_to == own_id)
    if not conditions:
        return []
    stmt = (
        select(Message.thread_id)
        .where(
            Message.user_id == user_id,
            Message.thread_id.is_not(None),
            Message.id != message_id,
            sa_or(*conditions),
        )
        .distinct()
    )
    rows = await session.scalars(stmt)
    return [tid for tid in rows if tid is not None]


async def _find_by_subject(
    session: AsyncSession, user_id: uuid.UUID, message: Message
) -> uuid.UUID | None:
    subject_norm = normalize_subject(message.subject)
    # Never join when this message both lacks references and has an empty subject (§7.3).
    if not subject_norm:
        return None
    anchor = message.internal_date
    lo = anchor - _SUBJECT_FALLBACK_WINDOW
    hi = anchor + _SUBJECT_FALLBACK_WINDOW
    stmt = (
        select(Thread)
        .where(
            Thread.user_id == user_id,
            Thread.subject_norm == subject_norm,
            Thread.last_message_at >= lo,
            Thread.last_message_at <= hi,
        )
        .order_by(Thread.last_message_at.desc())
    )
    my_recipients = _recipient_emails(message)
    if not my_recipients:
        return None
    for thread in await session.scalars(stmt):
        participant_emails = {
            (p.get("email") or "").lower() for p in thread.participants if p.get("email")
        }
        if my_recipients & participant_emails:
            return thread.id
    return None


async def _merge_threads(
    session: AsyncSession, thread_ids: list[uuid.UUID]
) -> tuple[uuid.UUID, list[uuid.UUID]]:
    """Merge ``thread_ids`` keeping the oldest ``last_message_at`` as canonical (§7.2.5)."""
    threads = list(await session.scalars(select(Thread).where(Thread.id.in_(thread_ids))))
    canonical = min(threads, key=lambda t: t.last_message_at)
    others = [t.id for t in threads if t.id != canonical.id]
    if others:
        await session.execute(
            update(Message).where(Message.thread_id.in_(others)).values(thread_id=canonical.id)
        )
        await session.execute(delete(Thread).where(Thread.id.in_(others)))
    return canonical.id, others


async def _create_thread(session: AsyncSession, message: Message) -> uuid.UUID:
    thread = Thread(
        user_id=message.user_id,
        subject_norm=normalize_subject(message.subject),
        last_message_at=message.internal_date,
    )
    session.add(thread)
    await session.flush()
    return thread.id


async def denormalize_thread(session: AsyncSession, thread_id: uuid.UUID) -> None:
    """Recompute a thread's denormalized fields from its messages (SPEC §7.2.6).

    The thread row is locked *before* its messages are read: two transactions
    recounting one thread at once (read in the webmail while a sync brings the
    same flag from a phone) then take turns, and the second reads what the
    first committed — instead of both reading first and the slower one writing
    an old count (a message shown read, its thread still "unread"). The
    messages are reloaded from the database, not taken from the session's
    copies, which a bulk UPDATE may have left behind."""
    thread = await session.scalar(
        select(Thread)
        .where(Thread.id == thread_id)
        .with_for_update()
        .execution_options(populate_existing=True)
    )
    if thread is None:
        return
    messages = list(
        await session.scalars(
            select(Message)
            .where(Message.thread_id == thread_id)
            .order_by(Message.internal_date.desc())
            .execution_options(populate_existing=True)
        )
    )
    if not messages:
        # A reply still being written keeps its (now empty) thread: the draft
        # points at it.
        if await session.scalar(select(exists().where(Draft.thread_id == thread_id))):
            thread.message_count = 0
            thread.unread_count = 0
            return
        await session.execute(delete(Thread).where(Thread.id == thread_id))
        return

    newest = messages[0]
    participants: list[dict[str, str]] = []
    seen_emails: set[str] = set()
    for msg in messages:  # newest first => recency order
        for addr in list(msg.from_addr) + list(msg.to_addr) + list(msg.cc_addr):
            email = (addr.get("email") or "").lower()
            if email and email not in seen_emails:
                seen_emails.add(email)
                participants.append({"name": addr.get("name", ""), "email": addr.get("email", "")})

    thread.message_count = len(messages)
    thread.unread_count = sum(1 for m in messages if "\\Seen" not in m.flags)
    thread.has_attachments = any(m.has_attachments for m in messages)
    thread.pinned = any("\\Flagged" in m.flags for m in messages)
    thread.account_ids = sorted({m.account_id for m in messages})
    thread.last_message_at = newest.internal_date
    thread.snippet = newest.snippet
    thread.subject_norm = normalize_subject(newest.subject) or thread.subject_norm
    thread.participants = participants


async def thread_message(session: AsyncSession, message: Message) -> uuid.UUID:
    """Assign ``message`` to a thread (creating/attaching/merging) and denormalize.

    Returns the resulting thread id and sets ``message.thread_id``.
    """
    candidates = _candidate_ids(message)
    parent_refs = _parent_refs(message)
    own_id = message.message_id_hdr
    thread_ids = await _find_by_references(
        session, message.user_id, message.id, candidates, parent_refs, own_id
    )

    # Subject fallback only for messages that declare no parentage (a broken/absent
    # References chain). A message that explicitly replies to a specific parent is
    # never subject-merged into an unrelated thread — e.g. two auto-replies to
    # different originals that happen to share the same subject (§7.3).
    if not thread_ids and not parent_refs:
        subject_thread = await _find_by_subject(session, message.user_id, message)
        if subject_thread is not None:
            thread_ids = [subject_thread]

    if not thread_ids:
        thread_id = await _create_thread(session, message)
    elif len(thread_ids) == 1:
        thread_id = thread_ids[0]
    else:
        thread_id, _merged_away = await _merge_threads(session, thread_ids)

    message.thread_id = thread_id
    await session.flush()
    await denormalize_thread(session, thread_id)
    await session.flush()
    return thread_id


async def repair_threads(session: AsyncSession, user_id: uuid.UUID) -> int:
    """Recount the user's threads whose stored counters disagree with their
    messages (and drop threads left without any). A safety net run nightly
    (app/indexer/gc.py): every write path recounts its threads, but a missed
    one would otherwise show a read message as unread until it changed again.
    Returns how many were repaired."""
    seen = Message.flags.any("\\Seen")  # type: ignore[arg-type]
    actual = (
        select(
            Message.thread_id.label("tid"),
            func.count().label("n"),
            func.count().filter(~seen).label("unread"),
        )
        .where(Message.user_id == user_id, Message.thread_id.is_not(None))
        .group_by(Message.thread_id)
        .subquery()
    )
    stale = list(
        await session.scalars(
            select(Thread.id)
            .outerjoin(actual, actual.c.tid == Thread.id)
            .where(
                Thread.user_id == user_id,
                (actual.c.tid.is_(None) & (Thread.message_count != 0))
                | (actual.c.tid.is_(None) & ~exists().where(Draft.thread_id == Thread.id))
                | (actual.c.n != Thread.message_count)
                | (actual.c.unread != Thread.unread_count),
            )
        )
    )
    for tid in sorted(stale):
        await denormalize_thread(session, tid)
    return len(stale)