webmail-client/plugins/webmail_plugin_antispam/webmail_plugin_antispam/jobs.py

"""Provider-Junk bootstrap (SPEC §C.5).

A per-user Bayes classifier knows nothing on day one, which is exactly when the
filter is asked to judge the most mail. But the user already owns a labelled
corpus: whatever their provider filed as Junk is spam they agreed to lose, and
old Inbox mail they never moved is ham they agreed to keep. Feeding both gives
Bayes a warm start without asking the user to train anything.

Deliberately unhurried. This is a background nicety, not a deadline: it yields
between messages so it cannot monopolise the worker or the engine, and it gives
up quietly on anything that looks wrong rather than retrying into a wall.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from typing import Any

from app.plugins.sdk import sdk
from webmail_plugin_antispam import policy, rspamd

logger = logging.getLogger("plg.antispam")

_SDK = sdk("antispam")

# §C.5: up to 500 of each. Enough to shape a classifier, small enough that the
# whole pass is minutes rather than hours.
SPAM_SAMPLE = 500
HAM_SAMPLE = 500
# Ham comes from mail old enough that the user has plainly not triaged it — a
# message from this morning is not yet evidence of anything.
HAM_MIN_AGE_DAYS = 30
# A breath between messages: the engine is shared with live classification, and
# live mail must never queue behind a training backlog.
PAUSE_SECONDS = 0.05

# A Bayes classifier learns the ratio it is shown. Accounts do not sync their
# provider's Junk folder by default (accounts.sync_spam is off), so the spam side
# of this corpus is routinely tiny while the inbox side is thousands — teaching
# both as found would train "everything is ham" and make the filter worse than
# leaving it untrained. Two guards:
#
#   below MIN_SPAM_CORPUS   teach nothing at all, and leave the account unmarked
#                           so a later run can do it properly
#   above it                teach at most MAX_HAM_RATIO times as much ham as spam
MIN_SPAM_CORPUS = 10
MAX_HAM_RATIO = 3

# Accounts whose warm start is done. A new key on purpose: markers under the old
# one ("_bootstrapped") may have been written by the `-0` bug described in
# _teach — an account marked done after a total engine outage, with nothing
# taught, and never retried. Ignoring them costs every older install one repeat
# bootstrap, which is harmless: rspamd skips messages it has already learned.
_DONE_KEY = "bootstrap_done"


async def _teach(user_id: uuid.UUID, ids: list[uuid.UUID], *, spam: bool) -> int | None:
    """Number of messages learned, or None if the engine stopped accepting them.

    None rather than a negative count: the count was once negated to signal
    failure, which works for every case except the one that matters. If the
    FIRST message is refused, nothing has been taught yet, `-0` is `0`, and the
    caller's `< 0` test reads a total engine outage as a clean run of zero — then
    marks the account bootstrapped, so it is never retried.
    """
    taught = 0
    for message_id in ids:
        raw = await _SDK.messages.raw(user_id, message_id)
        if raw is None:
            continue
        if not await rspamd.learn(raw, user_key=str(user_id), spam=spam):
            # Engine gone. Stop rather than grind through hundreds of failures;
            # the account stays unmarked, so a later run picks it up again.
            return None
        taught += 1
        await asyncio.sleep(PAUSE_SECONDS)
    return taught


async def bootstrap_account(
    ctx: dict[str, Any], *, user_id: str, account_id: str, force: bool = False
) -> None:
    """Teach Bayes from one account's existing mail. Safe to call more than once.

    ``force`` runs it even for an account already marked done — the "learn from
    my Spam folder now" button. The provider-Junk setting still applies: the
    user turned that off for a reason.
    """
    uid, acc = uuid.UUID(user_id), uuid.UUID(account_id)

    settings = await policy.load_settings(uid)
    if not settings.learn_from_provider_junk or not settings.covers(account_id):
        return

    done = await _SDK.kv(uid).get(_DONE_KEY, [])
    if account_id in done and not force:
        return

    spam_ids = await _SDK.messages.sample(
        uid, account_id=acc, special_use="junk", limit=SPAM_SAMPLE
    )
    ham_ids = await _SDK.messages.sample(
        uid,
        account_id=acc,
        special_use="inbox",
        limit=HAM_SAMPLE,
        older_than_days=HAM_MIN_AGE_DAYS,
    )
    if len(spam_ids) < MIN_SPAM_CORPUS:
        # Nothing to be learned from an empty Junk folder, and plenty of harm in
        # learning only the ham beside it. Usually means the account is not
        # syncing its provider's spam folder.
        logger.info(
            "antispam bootstrap for %s skipped: only %d spam samples (need %d); "
            "enable spam-folder sync on the account to build a corpus",
            account_id,
            len(spam_ids),
            MIN_SPAM_CORPUS,
        )
        return

    ham_ids = ham_ids[: len(spam_ids) * MAX_HAM_RATIO]

    spam_taught = await _teach(uid, spam_ids, spam=True)
    if spam_taught is None:
        logger.warning(
            "antispam bootstrap for %s stopped: engine unavailable", account_id
        )
        return
    ham_taught = await _teach(uid, ham_ids, spam=False)
    if ham_taught is None:
        logger.warning(
            "antispam bootstrap for %s stopped: engine unavailable", account_id
        )
        return
    if spam_taught == 0:
        # Every sample was unreadable. Nothing was learned, so do not record a
        # warm start that did not happen.
        logger.warning(
            "antispam bootstrap for %s: no samples could be read", account_id
        )
        return

    if account_id not in done:
        await _SDK.kv(uid).set(_DONE_KEY, [*done, account_id])
    logger.info(
        "antispam bootstrap for %s: learned %d spam, %d ham",
        account_id,
        spam_taught,
        ham_taught,
    )
    await _SDK.publish(
        "plugin.antispam.bootstrapped",
        {
            "user_id": user_id,
            "account_id": account_id,
            "spam": spam_taught,
            "ham": ham_taught,
        },
    )


async def bootstrap_sweep(ctx: dict[str, Any]) -> None:
    """Find accounts that still need a warm start, and give them one.

    The `account.state_changed` hook catches an account the moment it finishes
    importing, which covers accounts added from now on. This covers everything
    else, and it is the reason bootstrapping is not a thing anyone has to
    remember to do:

    * accounts that existed before this plugin did — their transition to `live`
      happened long ago and will not happen again;
    * accounts whose Junk folder was empty at the time and has since filled, or
      which only started syncing the provider's spam folder later;
    * a transition missed because the worker was restarting.

    Cheap to repeat: `bootstrap_account` returns immediately for an account
    already marked done, and declines (without marking) one whose corpus is still
    too small to learn from — so this simply retries until there is something
    worth learning.
    """
    for user_id, account_id in await _SDK.accounts.live():
        try:
            await bootstrap_account(
                ctx, user_id=str(user_id), account_id=str(account_id)
            )
        except Exception as exc:  # noqa: BLE001 - one bad account must not stop the sweep
            logger.warning(
                "antispam bootstrap sweep failed for %s: %s", account_id, exc
            )