perf(email): fetch headers before body, use UID SEARCH to skip re-fetch

The IMAP poll loop previously downloaded the entire message body for every UNSEEN message before running any filter (self-sent, SPF/DKIM, allow-list), and only learned the UID by parsing it back out of that fetch response. Rejected messages stay unseen and get re-fetched in full on every subsequent poll. Switch to UID SEARCH (UIDs come back directly, no per-message fetch needed to learn them) so already-processed UIDs are skipped before any network fetch, then fetch headers only to evaluate every filter — the full body is downloaded only for messages that pass every check and are actually delivered. No behavior change: filter outcomes and \Seen semantics are unchanged.
2026-08-31 08:13:11 +03:00 · 2026-08-30 17:38:44 +08:00
parent 1ac1b35c84
commit f573ecfe56
2 changed files with 113 additions and 198 deletions
@@ -430,7 +430,13 @@ class EmailChannel(BaseChannel):
        skipped_uids: set[str],
        cycle_uids: set[str],
    ) -> list[dict[str, Any]] | None:
-        """Fetch messages by arbitrary IMAP search criteria."""
+        """Fetch messages by arbitrary IMAP search criteria.
+
+        Uses UID SEARCH so already-processed UIDs are recognized before any
+        FETCH at all, then fetches headers only to evaluate every filter — the
+        full body (and any attachments) is downloaded only for messages that
+        pass every check and are actually going to be delivered.
+        """
        mailbox = self.config.imap_mailbox or "INBOX"

        client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True)
@@ -438,29 +444,30 @@ class EmailChannel(BaseChannel):
            return messages

        try:
-            status, data = client.search(None, *search_criteria)
-            if status != "OK" or not data:
+            status, data = client.uid("SEARCH", None, *search_criteria)
+            if status != "OK" or not data or not data[0]:
                return messages

-            ids = data[0].split()
-            if limit > 0 and len(ids) > limit:
-                ids = ids[-limit:]
-            for imap_id in ids:
-                status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)")
+            uids = [raw.decode("ascii", errors="ignore") for raw in data[0].split()]
+            if limit > 0 and len(uids) > limit:
+                uids = uids[-limit:]
+
+            features: _ServerFeatures | None = None
+
+            for uid in uids:
+                if not uid or uid in cycle_uids:
+                    continue
+                if dedupe and uid in self._processed_uids:
+                    continue
+
+                status, fetched = client.uid("FETCH", uid, "(BODY.PEEK[HEADER])")
                if status != "OK" or not fetched:
                    continue
-
-                raw_bytes = self._extract_message_bytes(fetched)
-                if raw_bytes is None:
+                header_bytes = self._extract_message_bytes(fetched)
+                if header_bytes is None:
                    continue

-                uid = self._extract_uid(fetched)
-                if uid and uid in cycle_uids:
-                    continue
-                if dedupe and uid and uid in self._processed_uids:
-                    continue
-
-                parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
+                parsed = BytesParser(policy=policy.default).parsebytes(header_bytes)
                sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
                if not sender:
                    continue
@@ -468,8 +475,7 @@ class EmailChannel(BaseChannel):
                    self.logger.info("From {} ignored: matches bot-owned address", sender)
                    self._remember_processed_uid(uid, dedupe, cycle_uids)
                    if mark_seen:
-                        client.store(imap_id, "+FLAGS", "\\Seen")
-                    if uid:
+                        features = self._mark_seen_uid(client, uid, features)
                    skipped_uids.add(uid)
                    continue

@@ -482,7 +488,6 @@ class EmailChannel(BaseChannel):
                        sender,
                    )
                    self._remember_processed_uid(uid, dedupe, cycle_uids)
-                    if uid:
                    skipped_uids.add(uid)
                    continue
                if self.config.verify_dkim and not dkim_pass:
@@ -492,18 +497,26 @@ class EmailChannel(BaseChannel):
                        sender,
                    )
                    self._remember_processed_uid(uid, dedupe, cycle_uids)
-                    if uid:
                    skipped_uids.add(uid)
                    continue

                if not self.is_allowed(sender):
                    self._remember_processed_uid(uid, dedupe, cycle_uids)
                    if mark_seen:
-                        client.store(imap_id, "+FLAGS", "\\Seen")
-                    if uid:
+                        features = self._mark_seen_uid(client, uid, features)
                    skipped_uids.add(uid)
                    continue

+                # Passed every filter — only now fetch the full message body
+                # (and any attachments) for the message we're actually delivering.
+                status, full_fetched = client.uid("FETCH", uid, "(BODY.PEEK[])")
+                if status != "OK" or not full_fetched:
+                    continue
+                raw_bytes = self._extract_message_bytes(full_fetched)
+                if raw_bytes is None:
+                    continue
+                parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
+
                subject = self._decode_header_value(parsed.get("Subject", ""))
                date_value = parsed.get("Date", "")
                message_id = parsed.get("Message-ID", "").strip()
@@ -556,10 +569,19 @@ class EmailChannel(BaseChannel):
                self._remember_processed_uid(uid, dedupe, cycle_uids)

                if mark_seen:
-                    client.store(imap_id, "+FLAGS", "\\Seen")
+                    features = self._mark_seen_uid(client, uid, features)
        finally:
            self._close_imap_client(client)

+    def _mark_seen_uid(
+        self, client: Any, uid: str, features: _ServerFeatures | None
+    ) -> _ServerFeatures:
+        """Mark a single UID \\Seen, reusing session-learned STORE support."""
+        if features is None:
+            features = self._server_features(client)
+        self._uid_store_flag(client, uid, "\\Seen", features)
+        return features
+
    def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None:
        if self.config.imap_use_ssl:
            client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
@@ -714,11 +736,14 @@ class EmailChannel(BaseChannel):
        return data[0].split()[0]

    def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool:
+        return self._uid_store_flag(client, uid, "\\Deleted", features)
+
+    def _uid_store_flag(self, client: Any, uid: str, flag: str, features: _ServerFeatures) -> bool:
        # Optimistic path: try UID STORE first because UID is stable and avoids
        # sequence-number lookup. If this fails once for the session, remember it
        # and use the sequence STORE fallback directly for remaining UIDs.
        if features.uid_store is not False:
-            status, _ = client.uid("STORE", uid, "+FLAGS", "(\\Deleted)")
+            status, _ = client.uid("STORE", uid, "+FLAGS", f"({flag})")
            if status == "OK":
                features.uid_store = True
                return True
@@ -728,12 +753,12 @@ class EmailChannel(BaseChannel):
        # unreliable: resolve the current sequence number from UID and use STORE.
        imap_id = self._lookup_imap_id_by_uid(client, uid)
        if not imap_id:
-            self.logger.warning("Post-action skipped: UID {} not found", uid)
+            self.logger.warning("Could not locate UID {} to set flag {}", uid, flag)
            return False

-        status, _ = client.store(imap_id, "+FLAGS", "\\Deleted")
+        status, _ = client.store(imap_id, "+FLAGS", flag)
        if status != "OK":
-            self.logger.warning("Post-action failed: could not mark UID {} as deleted", uid)
+            self.logger.warning("Failed to set flag {} on UID {}", flag, uid)
            return False
        return True

@@ -773,16 +798,6 @@ class EmailChannel(BaseChannel):
                    return bytes(fetched_item[1])
        return None

-    @staticmethod
-    def _extract_uid(fetched: list[Any]) -> str:
-        for item in fetched:
-            if isinstance(item, tuple) and item and isinstance(item[0], (bytes, bytearray)):
-                head = bytes(item[0]).decode("utf-8", errors="ignore")
-                m = re.search(r"UID\s+(\d+)", head)
-                if m:
-                    return m.group(1)
-        return ""
-
    @staticmethod
    def _decode_header_value(value: str) -> str:
        if not value:
@@ -53,30 +53,7 @@ def _make_raw_email(
 def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
    raw = _make_raw_email(subject="Invoice", body="Please pay")

-    class FakeIMAP:
-        def __init__(self) -> None:
-            self.store_calls: list[tuple[bytes, str, str]] = []
-
-        def login(self, _user: str, _pw: str):
-            return "OK", [b"logged in"]
-
-        def select(self, _mailbox: str):
-            return "OK", [b"1"]
-
-        def search(self, *_args):
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
-
-        def store(self, imap_id: bytes, op: str, flags: str):
-            self.store_calls.append((imap_id, op, flags))
-            return "OK", [b""]
-
-        def logout(self):
-            return "BYE", [b""]
-
-    fake = FakeIMAP()
+    fake = _make_fake_imap(raw, uid=b"123")
    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)

    channel = EmailChannel(_make_config(), MessageBus())
@@ -86,7 +63,7 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
    assert items[0]["sender"] == "alice@example.com"
    assert items[0]["subject"] == "Invoice"
    assert "Please pay" in items[0]["content"]
-    assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
+    assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
    assert skipped_uids == set()

    # Same UID should be deduped in-process.
@@ -98,26 +75,8 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
 def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> None:
    raw = _make_raw_email(subject="Invoice", body="Please pay")

-    class FakeIMAP:
-        def login(self, _user: str, _pw: str):
-            return "OK", [b"logged in"]
-
-        def select(self, _mailbox: str):
-            return "OK", [b"1"]
-
-        def search(self, *_args):
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
-
-        def store(self, _imap_id: bytes, _op: str, _flags: str):
-            return "OK", [b""]
-
-        def logout(self):
-            return "BYE", [b""]
-
-    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
+    fake = _make_fake_imap(raw, uid=b"123")
+    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)

    channel = EmailChannel(_make_config(post_action="delete"), MessageBus())
    items, skipped_uids = channel._fetch_new_messages()
@@ -130,26 +89,10 @@ def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> No
 def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None:
    raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")

-    class FakeIMAP:
-        def login(self, _user: str, _pw: str):
-            return "OK", [b"logged in"]
-
-        def select(self, _mailbox: str):
-            return "OK", [b"1"]
-
-        def search(self, *_args):
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
-
-        def store(self, _imap_id: bytes, _op: str, _flags: str):
-            return "OK", [b""]
-
-        def logout(self):
-            return "BYE", [b""]
-
-    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
+    monkeypatch.setattr(
+        "nanobot.channels.email.runtime.imaplib.IMAP4_SSL",
+        lambda _h, _p: _make_fake_imap(raw, uid=b"123"),
+    )

    channel_skip = EmailChannel(
        _make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True),
@@ -545,30 +488,7 @@ async def test_start_keeps_post_actions_for_successful_emails_when_later_deliver
 def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) -> None:
    raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")

-    class FakeIMAP:
-        def __init__(self) -> None:
-            self.store_calls: list[tuple[bytes, str, str]] = []
-
-        def login(self, _user: str, _pw: str):
-            return "OK", [b"logged in"]
-
-        def select(self, _mailbox: str):
-            return "OK", [b"1"]
-
-        def search(self, *_args):
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
-
-        def store(self, imap_id: bytes, op: str, flags: str):
-            self.store_calls.append((imap_id, op, flags))
-            return "OK", [b""]
-
-        def logout(self):
-            return "BYE", [b""]
-
-    fake = FakeIMAP()
+    fake = _make_fake_imap(raw, uid=b"123")
    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)

    channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus())
@@ -576,7 +496,7 @@ def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) ->

    assert items == []
    assert skipped_uids == {"123"}
-    assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
+    assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls

    # Same UID should still be deduped after being ignored.
    items_again, skipped_again = channel._fetch_new_messages()
@@ -614,37 +534,14 @@ def test_fetch_new_messages_skips_self_sent_across_identity_sources(
    imap_username matches, and must be case-insensitive."""
    raw = _make_raw_email(from_addr=from_header, subject="Loop test")

-    class FakeIMAP:
-        def __init__(self) -> None:
-            self.store_calls: list[tuple[bytes, str, str]] = []
-
-        def login(self, _user: str, _pw: str):
-            return "OK", [b"logged in"]
-
-        def select(self, _mailbox: str):
-            return "OK", [b"1"]
-
-        def search(self, *_args):
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
-
-        def store(self, imap_id: bytes, op: str, flags: str):
-            self.store_calls.append((imap_id, op, flags))
-            return "OK", [b""]
-
-        def logout(self):
-            return "BYE", [b""]
-
-    fake = FakeIMAP()
+    fake = _make_fake_imap(raw, uid=b"123")
    monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)

    channel = EmailChannel(_make_config(**config_override), MessageBus())
    items, _ = channel._fetch_new_messages()

    assert items == []
-    assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
+    assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls


 def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None:
@@ -662,15 +559,16 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
        def select(self, _mailbox: str):
            return "OK", [b"1"]

-        def search(self, *_args):
+        def uid(self, command: str, *args):
+            if command == "SEARCH":
                self.search_calls += 1
                if fail_once["pending"]:
                    fail_once["pending"] = False
                    raise imaplib.IMAP4.abort("socket error")
-            return "OK", [b"1"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
+                return "OK", [b"123"]
+            if command == "FETCH":
                return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
+            return "OK", [b""]

        def store(self, imap_id: bytes, op: str, flags: str):
            self.store_calls.append((imap_id, op, flags))
@@ -700,10 +598,7 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
 def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypatch) -> None:
    raw_first = _make_raw_email(subject="First", body="First body")
    raw_second = _make_raw_email(subject="Second", body="Second body")
-    mailbox_state = {
-        b"1": {"uid": b"123", "raw": raw_first, "seen": False},
-        b"2": {"uid": b"124", "raw": raw_second, "seen": False},
-    }
+    mailbox_state = {"123": raw_first, "124": raw_second}
    fail_once = {"pending": True}

    class FlakyIMAP:
@@ -713,20 +608,18 @@ def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypa
        def select(self, _mailbox: str):
            return "OK", [b"2"]

-        def search(self, *_args):
-            unseen_ids = [imap_id for imap_id, item in mailbox_state.items() if not item["seen"]]
-            return "OK", [b" ".join(unseen_ids)]
-
-        def fetch(self, imap_id: bytes, _parts: str):
-            if imap_id == b"2" and fail_once["pending"]:
+        def uid(self, command: str, *args):
+            if command == "SEARCH":
+                keys = " ".join(sorted(mailbox_state.keys(), key=int))
+                return "OK", [keys.encode()]
+            if command == "FETCH":
+                uid = args[0]
+                if uid == "124" and fail_once["pending"]:
                    fail_once["pending"] = False
                    raise imaplib.IMAP4.abort("socket error")
-            item = mailbox_state[imap_id]
-            header = b"%s (UID %s BODY[] {200})" % (imap_id, item["uid"])
-            return "OK", [(header, item["raw"]), b")"]
-
-        def store(self, imap_id: bytes, _op: str, _flags: str):
-            mailbox_state[imap_id]["seen"] = True
+                raw = mailbox_state[uid]
+                header = f"{uid} (UID {uid} BODY[] {{200}})".encode()
+                return "OK", [(header, raw), b")"]
            return "OK", [b""]

        def logout(self):
@@ -1044,12 +937,13 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
        def select(self, _mailbox: str):
            return "OK", [b"1"]

-        def search(self, *_args):
-            self.search_args = _args
-            return "OK", [b"5"]
-
-        def fetch(self, _imap_id: bytes, _parts: str):
+        def uid(self, command: str, *args):
+            if command == "SEARCH":
+                self.search_args = args
+                return "OK", [b"999"]
+            if command == "FETCH":
                return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"]
+            return "OK", [b""]

        def store(self, imap_id: bytes, op: str, flags: str):
            self.store_calls.append((imap_id, op, flags))
@@ -1070,7 +964,7 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m

    assert len(items) == 1
    assert items[0]["subject"] == "Status"
-    # search(None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
+    # uid("SEARCH", None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
    assert fake.search_args is not None
    assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
    assert fake.store_calls == []
@@ -1080,11 +974,12 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
 # Security: Anti-spoofing tests for Authentication-Results verification
 # ---------------------------------------------------------------------------

-def _make_fake_imap(raw: bytes):
+def _make_fake_imap(raw: bytes, uid: bytes = b"500"):
    """Return a FakeIMAP class pre-loaded with the given raw email."""
    class FakeIMAP:
        def __init__(self) -> None:
            self.store_calls: list[tuple[bytes, str, str]] = []
+            self.uid_calls: list[tuple] = []

        def login(self, _user: str, _pw: str):
            return "OK", [b"logged in"]
@@ -1092,11 +987,16 @@ def _make_fake_imap(raw: bytes):
        def select(self, _mailbox: str):
            return "OK", [b"1"]

-        def search(self, *_args):
-            return "OK", [b"1"]
+        def capability(self):
+            return "OK", [b"IMAP4rev1"]

-        def fetch(self, _imap_id: bytes, _parts: str):
-            return "OK", [(b"1 (UID 500 BODY[] {200})", raw), b")"]
+        def uid(self, command: str, *args):
+            self.uid_calls.append((command, *args))
+            if command == "SEARCH":
+                return "OK", [uid]
+            if command == "FETCH":
+                return "OK", [(b"1 (UID " + uid + b" BODY[] {200})", raw), b")"]
+            return "OK", [b""]

        def store(self, imap_id: bytes, op: str, flags: str):
            self.store_calls.append((imap_id, op, flags))
@@ -1292,7 +1192,7 @@ def test_fetch_new_messages_ignores_unauthorized_sender_before_attachments(monke

    assert channel._fetch_new_messages() == ([], {"500"})
    assert called["attachments"] is False
-    assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
+    assert ("STORE", "500", "+FLAGS", "(\\Seen)") in fake.uid_calls


 def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None: