import logging from dataclasses import dataclass from datetime import datetime, timedelta, timezone from bson import ObjectId from app.services.message_ids import ( first_message_id, normalize_message_id, parse_message_id_list, ) from app.services.subject_normalize import normalize_subject logger = logging.getLogger(__name__) FALLBACK_WINDOW = timedelta(days=7) @dataclass(frozen=True) class ThreadResolution: thread_id: ObjectId | None parent_message_id: ObjectId | None thread_relation: str thread_resolution_method: str @dataclass(frozen=True) class KnownMessage: id: ObjectId thread_id: ObjectId internet_message_id: str account_id: str sent_at: datetime | None received_at: datetime | None created_at: datetime participants: frozenset[str] normalized_subject: str def _message_timestamp(message: KnownMessage) -> datetime: return message.sent_at or message.received_at or message.created_at def resolve_thread_from_known( *, account_id: str, in_reply_to: str | None, references: list[str], subject: str, participants: set[str], known_by_internet_id: dict[str, KnownMessage], known_threads: dict[ObjectId, list[KnownMessage]] | None = None, now: datetime | None = None, ) -> ThreadResolution: scoped = { key: value for key, value in known_by_internet_id.items() if value.account_id == account_id } in_reply = normalize_message_id(in_reply_to) if in_reply and in_reply in scoped: parent = scoped[in_reply] return ThreadResolution( thread_id=parent.thread_id, parent_message_id=parent.id, thread_relation="reply", thread_resolution_method="in_reply_to", ) reference_ids = [ normalize_message_id(value) for value in references if normalize_message_id(value) ] for reference in reversed(reference_ids): if reference in scoped: parent = scoped[reference] return ThreadResolution( thread_id=parent.thread_id, parent_message_id=parent.id, thread_relation="reply", thread_resolution_method="references", ) fallback = _conservative_fallback( account_id=account_id, subject=subject, participants=participants, known_by_internet_id=scoped, known_threads=known_threads or {}, now=now or datetime.now(timezone.utc), ) if fallback is not None: return fallback return ThreadResolution( thread_id=None, parent_message_id=None, thread_relation="new_thread", thread_resolution_method="new_thread", ) def _conservative_fallback( *, account_id: str, subject: str, participants: set[str], known_by_internet_id: dict[str, KnownMessage], known_threads: dict[ObjectId, list[KnownMessage]], now: datetime, ) -> ThreadResolution | None: normalized = normalize_subject(subject) if not normalized or not participants: return None candidates: dict[ObjectId, KnownMessage] = {} for message in known_by_internet_id.values(): if message.normalized_subject != normalized: continue if not participants.intersection(message.participants): continue timestamp = _message_timestamp(message) if timestamp < now - FALLBACK_WINDOW: continue candidates[message.thread_id] = message if len(candidates) != 1: return None parent = next(iter(candidates.values())) return ThreadResolution( thread_id=parent.thread_id, parent_message_id=parent.id, thread_relation="unknown", thread_resolution_method="fallback", ) def participants_from_parsed(parsed: dict) -> set[str]: addresses: set[str] = set() for key in ("sender",): sender = parsed.get(key) if isinstance(sender, dict) and sender.get("address"): addresses.add(sender["address"].lower()) for key in ("to", "cc", "reply_to"): for item in parsed.get(key, []) or []: if isinstance(item, dict) and item.get("address"): addresses.add(item["address"].lower()) return addresses async def lookup_known_messages( database, account_id: str, internet_ids: list[str], ) -> dict[str, KnownMessage]: normalized_ids = [ value for value in {normalize_message_id(item) for item in internet_ids} if value ] if not normalized_ids: return {} pipeline = [ { "$match": { "message_id": {"$in": normalized_ids}, "$or": [ {"message_scope": "EXTERNAL"}, {"message_scope": {"$exists": False}}, ], } }, { "$lookup": { "from": "mailbox_entries", "let": {"message_id": "$_id"}, "pipeline": [ { "$match": { "$expr": { "$and": [ {"$eq": ["$message_id", "$$message_id"]}, {"$eq": ["$account_id", account_id]}, ] } } }, {"$limit": 1}, ], "as": "entry", } }, {"$match": {"entry.0": {"$exists": True}}}, ] known: dict[str, KnownMessage] = {} cursor = await database.messages.aggregate(pipeline) async for document in cursor: internet_id = normalize_message_id(document.get("message_id")) if not internet_id: continue entry = document["entry"][0] thread_id = entry.get("thread_id") or document.get("thread_id") if thread_id is None: continue participants = participants_from_parsed(document) known[internet_id] = KnownMessage( id=document["_id"], thread_id=thread_id, internet_message_id=internet_id, account_id=account_id, sent_at=document.get("sent_at"), received_at=document.get("received_at"), created_at=document.get("created_at"), participants=frozenset(participants), normalized_subject=normalize_subject(document.get("subject")), ) return known async def resolve_thread( database, *, account_id: str, in_reply_to: str | None, references: list[str], subject: str, participants: set[str], ) -> ThreadResolution: search_ids = [ *([in_reply_to] if in_reply_to else []), *references, ] known = await lookup_known_messages(database, account_id, search_ids) resolution = resolve_thread_from_known( account_id=account_id, in_reply_to=in_reply_to, references=references, subject=subject, participants=participants, known_by_internet_id=known, ) logger.info( "thread_resolved account_id=%s method=%s relation=%s parent_found=%s", account_id, resolution.thread_resolution_method, resolution.thread_relation, resolution.parent_message_id is not None, ) return resolution