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, 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 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)."""
    messages = list(
        await session.scalars(
            select(Message)
            .where(Message.thread_id == thread_id)
            .order_by(Message.internal_date.desc())
        )
    )
    thread = await session.get(Thread, thread_id)
    if thread is None:
        return
    if not messages:
        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