lukasandClaude Opus 5
8eb699a05f
Warming Bayes hung on a single event — `account.state_changed` reaching "live". That covers an account added from now on and nothing else. An account that went live before this plugin existed never gets one; neither does one whose Junk folder was empty at the time and filled up later, or one whose owner enabled spam-folder sync afterwards. A worker restarting at the wrong moment loses it too. The result was a feature that quietly required someone to run it by hand per account, which is not a feature. An hourly sweep now walks every live account and offers each a bootstrap. That works because the per-account job was already cheap to repeat: it returns at once for an account already done, and declines — without marking it done — when the corpus is still too small to learn from. So the sweep simply keeps asking until there is something worth learning, then does it once. `s.accounts.live()` is the new SDK surface it needs: cross-user by design, because a sweep runs for the instance rather than a request. Ids and state only, so reading anything still goes through the user-scoped accessors. One failing account cannot stop the sweep — the mailbox that raises is logged and the rest still get their turn. docs/plugins.md now says this in general terms next to `JobSpec`: a state-change hook is one-shot, so pair it with a sweep when the work must eventually happen for everything. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018SqJYt7jD7oUEkoY68sBZo
331 lines
11 KiB
Python
331 lines
11 KiB
Python
"""C2 learning loop: what teaches, what doesn't, and what only teaches once.
|
|
|
|
Passive learning (SPEC §C.5) reads the user's own filing as a verdict. The value
|
|
is in the restraint — most folder moves say nothing about spam, and teaching the
|
|
same message twice skews the corpus.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from webmail_plugin_antispam import hook, jobs, policy
|
|
|
|
SETTINGS = policy.Settings(
|
|
enabled_accounts=[],
|
|
move_threshold=12.0,
|
|
banner_threshold=6.0,
|
|
network_checks=False,
|
|
learn_from_provider_junk=True,
|
|
)
|
|
|
|
|
|
class _FakeSdk:
|
|
def __init__(self, learned: str | None = None) -> None:
|
|
self.learned_before = learned
|
|
self.published: list[tuple[str, dict[str, Any]]] = []
|
|
self.recorded: list[dict[str, Any]] = []
|
|
self.enqueued: list[dict[str, Any]] = []
|
|
|
|
outer = self
|
|
|
|
class _M:
|
|
async def envelope(self, *_a: object, **_k: object) -> dict[str, Any]:
|
|
return {"account_id": "acc-1"}
|
|
|
|
async def raw(self, *_a: object, **_k: object) -> bytes:
|
|
return b"From: x@y.test\r\n\r\nbody\r\n"
|
|
|
|
self.messages = _M()
|
|
|
|
async def _enqueue(_f: object, **kw: Any) -> None:
|
|
outer.enqueued.append(kw)
|
|
|
|
self.enqueue = _enqueue
|
|
|
|
async def publish(self, event: str, payload: dict[str, Any]) -> None:
|
|
self.published.append((event, payload))
|
|
|
|
|
|
@pytest.fixture
|
|
def wired(monkeypatch: pytest.MonkeyPatch): # type: ignore[no-untyped-def]
|
|
def _install(*, learned: str | None = None, engine_up: bool = True) -> _FakeSdk:
|
|
fake = _FakeSdk(learned)
|
|
taught: list[bool] = []
|
|
|
|
async def _learn(_raw: bytes, *, user_key: str, spam: bool) -> bool:
|
|
taught.append(spam)
|
|
return engine_up
|
|
|
|
async def _already(_u: object, _m: object) -> str | None:
|
|
return fake.learned_before
|
|
|
|
async def _remember(_u: object, _m: object, **kw: Any) -> None:
|
|
fake.recorded.append(kw)
|
|
|
|
async def _settings(_u: object) -> policy.Settings:
|
|
return SETTINGS
|
|
|
|
monkeypatch.setattr(hook, "_SDK", fake)
|
|
monkeypatch.setattr(hook.rspamd, "learn", _learn)
|
|
monkeypatch.setattr(hook, "_already_learned", _already)
|
|
monkeypatch.setattr(hook, "_remember_learned", _remember)
|
|
monkeypatch.setattr(policy, "load_settings", _settings)
|
|
fake.taught = taught # type: ignore[attr-defined]
|
|
return fake
|
|
|
|
return _install
|
|
|
|
|
|
async def _move(before: str | None, now: str | None) -> None:
|
|
await hook.on_message_changed(
|
|
{
|
|
"user_id": str(uuid.uuid4()),
|
|
"message_id": str(uuid.uuid4()),
|
|
"previous_special_use": before,
|
|
"special_use": now,
|
|
}
|
|
)
|
|
|
|
|
|
async def test_filing_into_spam_teaches_spam(wired) -> None: # type: ignore[no-untyped-def]
|
|
fake = wired()
|
|
await _move("inbox", "junk")
|
|
assert fake.taught == [True]
|
|
assert fake.recorded[0]["verdict"] == "spam"
|
|
|
|
|
|
async def test_rescuing_from_spam_teaches_ham(wired) -> None: # type: ignore[no-untyped-def]
|
|
fake = wired()
|
|
await _move("junk", "inbox")
|
|
assert fake.taught == [False]
|
|
assert fake.recorded[0]["verdict"] == "ham"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("before", "now"),
|
|
[
|
|
("inbox", "archive"), # tidying up
|
|
("junk", "trash"), # deleting spam is not a verdict about ham
|
|
("archive", "inbox"), # unarchiving says nothing about spam
|
|
("sent", "archive"),
|
|
("inbox", "inbox"), # not a move at all
|
|
],
|
|
)
|
|
async def test_other_moves_teach_nothing(wired, before, now) -> None: # type: ignore[no-untyped-def]
|
|
fake = wired()
|
|
await _move(before, now)
|
|
assert fake.taught == []
|
|
assert fake.recorded == []
|
|
|
|
|
|
async def test_the_same_lesson_is_not_taught_twice(wired) -> None: # type: ignore[no-untyped-def]
|
|
"""Users move things back and forth; the corpus should not count that twice."""
|
|
fake = wired(learned="spam")
|
|
await _move("inbox", "junk")
|
|
assert fake.taught == []
|
|
|
|
|
|
async def test_reversing_a_lesson_does_teach(wired) -> None: # type: ignore[no-untyped-def]
|
|
"""Having taught spam, moving it back out must be able to teach ham."""
|
|
fake = wired(learned="spam")
|
|
await _move("junk", "inbox")
|
|
assert fake.taught == [False]
|
|
|
|
|
|
async def test_nothing_is_recorded_when_the_engine_is_down(wired) -> None: # type: ignore[no-untyped-def]
|
|
"""`learned` stays unset so a later move can teach it again."""
|
|
fake = wired(engine_up=False)
|
|
await _move("inbox", "junk")
|
|
assert fake.taught == [True]
|
|
assert fake.recorded == []
|
|
|
|
|
|
# --- bootstrap ---------------------------------------------------------------
|
|
|
|
|
|
async def test_going_live_enqueues_the_bootstrap(wired) -> None: # type: ignore[no-untyped-def]
|
|
fake = wired()
|
|
await hook.on_account_state_changed(
|
|
{"user_id": str(uuid.uuid4()), "account_id": "acc-1", "state": "live"}
|
|
)
|
|
assert fake.enqueued == [{"user_id": fake.enqueued[0]["user_id"], "account_id": "acc-1"}]
|
|
|
|
|
|
async def test_other_states_do_not_enqueue(wired) -> None: # type: ignore[no-untyped-def]
|
|
fake = wired()
|
|
for state in ("backfilling", "error", "paused"):
|
|
await hook.on_account_state_changed(
|
|
{"user_id": str(uuid.uuid4()), "account_id": "acc-1", "state": state}
|
|
)
|
|
assert fake.enqueued == []
|
|
|
|
|
|
async def test_bootstrap_respects_the_setting(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Opting out of provider-Junk learning means no training pass at all."""
|
|
sampled: list[str] = []
|
|
|
|
async def _settings(_u: object) -> policy.Settings:
|
|
return policy.Settings(
|
|
enabled_accounts=[],
|
|
move_threshold=12.0,
|
|
banner_threshold=6.0,
|
|
network_checks=False,
|
|
learn_from_provider_junk=False,
|
|
)
|
|
|
|
class _M:
|
|
async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
|
|
sampled.append(str(kw.get("special_use")))
|
|
return []
|
|
|
|
class _Fake:
|
|
messages = _M()
|
|
|
|
monkeypatch.setattr(policy, "load_settings", _settings)
|
|
monkeypatch.setattr(jobs, "_SDK", _Fake())
|
|
await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
|
|
assert sampled == []
|
|
|
|
|
|
async def test_bootstrap_refuses_a_lopsided_corpus(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Three spam beside a thousand ham would teach "everything is ham".
|
|
|
|
Accounts do not sync the provider's spam folder by default, so this is the
|
|
common case, not an edge case.
|
|
"""
|
|
taught: list[bool] = []
|
|
|
|
async def _settings(_u: object) -> policy.Settings:
|
|
return SETTINGS
|
|
|
|
class _Kv:
|
|
async def get(self, _k: str, _d: object = None) -> list[str]:
|
|
return []
|
|
|
|
async def set(self, _k: str, _v: object) -> None:
|
|
raise AssertionError("must not mark the account done")
|
|
|
|
class _M:
|
|
async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
|
|
n = 3 if kw.get("special_use") == "junk" else 1000
|
|
return [uuid.uuid4() for _ in range(n)]
|
|
|
|
async def raw(self, *_a: object, **_k: object) -> bytes:
|
|
return b"x"
|
|
|
|
class _Fake:
|
|
messages = _M()
|
|
|
|
def kv(self, _u: object) -> _Kv:
|
|
return _Kv()
|
|
|
|
async def _learn(_raw: bytes, **kw: object) -> bool:
|
|
taught.append(bool(kw.get("spam")))
|
|
return True
|
|
|
|
monkeypatch.setattr(policy, "load_settings", _settings)
|
|
monkeypatch.setattr(jobs, "_SDK", _Fake())
|
|
monkeypatch.setattr(jobs.rspamd, "learn", _learn)
|
|
|
|
await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
|
|
assert taught == [], "taught from a corpus too lopsided to learn from"
|
|
|
|
|
|
async def test_bootstrap_caps_ham_at_a_multiple_of_spam(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
marked: list[object] = []
|
|
|
|
async def _settings(_u: object) -> policy.Settings:
|
|
return SETTINGS
|
|
|
|
class _Kv:
|
|
async def get(self, _k: str, _d: object = None) -> list[str]:
|
|
return []
|
|
|
|
async def set(self, _k: str, v: object) -> None:
|
|
marked.append(v)
|
|
|
|
class _M:
|
|
async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
|
|
n = 20 if kw.get("special_use") == "junk" else 1000
|
|
return [uuid.uuid4() for _ in range(n)]
|
|
|
|
async def raw(self, *_a: object, **_k: object) -> bytes:
|
|
return b"x"
|
|
|
|
class _Fake:
|
|
messages = _M()
|
|
published: list[object] = []
|
|
|
|
def kv(self, _u: object) -> _Kv:
|
|
return _Kv()
|
|
|
|
async def publish(self, _e: str, _p: object) -> None:
|
|
return None
|
|
|
|
taught: list[bool] = []
|
|
|
|
async def _learn(_raw: bytes, **kw: object) -> bool:
|
|
taught.append(bool(kw.get("spam")))
|
|
return True
|
|
|
|
monkeypatch.setattr(policy, "load_settings", _settings)
|
|
monkeypatch.setattr(jobs, "_SDK", _Fake())
|
|
monkeypatch.setattr(jobs.rspamd, "learn", _learn)
|
|
monkeypatch.setattr(jobs, "PAUSE_SECONDS", 0)
|
|
|
|
await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
|
|
assert taught.count(True) == 20
|
|
assert taught.count(False) == 60, "ham should be capped at 3x the spam count"
|
|
assert marked, "a completed bootstrap must mark the account done"
|
|
|
|
|
|
# --- the sweep: bootstrapping must not depend on catching an event -----------
|
|
|
|
|
|
async def test_sweep_visits_every_live_account(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""An account that went live before the plugin existed still gets a warm start."""
|
|
visited: list[tuple[str, str]] = []
|
|
accounts = [(uuid.uuid4(), uuid.uuid4()) for _ in range(3)]
|
|
|
|
class _Accounts:
|
|
async def live(self) -> list[tuple[uuid.UUID, uuid.UUID]]:
|
|
return accounts
|
|
|
|
class _Fake:
|
|
accounts = _Accounts()
|
|
|
|
async def _one(_ctx: object, *, user_id: str, account_id: str) -> None:
|
|
visited.append((user_id, account_id))
|
|
|
|
monkeypatch.setattr(jobs, "_SDK", _Fake())
|
|
monkeypatch.setattr(jobs, "bootstrap_account", _one)
|
|
await jobs.bootstrap_sweep({})
|
|
assert visited == [(str(u), str(a)) for u, a in accounts]
|
|
|
|
|
|
async def test_one_failing_account_does_not_stop_the_sweep(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
done: list[str] = []
|
|
accounts = [(uuid.uuid4(), uuid.uuid4()) for _ in range(3)]
|
|
|
|
class _Accounts:
|
|
async def live(self) -> list[tuple[uuid.UUID, uuid.UUID]]:
|
|
return accounts
|
|
|
|
class _Fake:
|
|
accounts = _Accounts()
|
|
|
|
async def _one(_ctx: object, *, user_id: str, account_id: str) -> None:
|
|
if account_id == str(accounts[1][1]):
|
|
raise RuntimeError("that mailbox is on fire")
|
|
done.append(account_id)
|
|
|
|
monkeypatch.setattr(jobs, "_SDK", _Fake())
|
|
monkeypatch.setattr(jobs, "bootstrap_account", _one)
|
|
await jobs.bootstrap_sweep({})
|
|
assert done == [str(accounts[0][1]), str(accounts[2][1])]
|