From acc87e8f56a4e181cf277c67922dbf9657035e10 Mon Sep 17 00:00:00 2001
From: RaresKeY <158580472+RaresKeY@users.noreply.github.com>
Date: Mon, 20 Jul 2026 21:35:06 +0000
Subject: [PATCH] fix(email): retire stale urgency accounts

---
 src/builtin_actions.py                 | 143 +++++++++++++--
 tests/test_email_urgency_checkpoint.py | 244 +++++++++++++++++++++++++
 2 files changed, 369 insertions(+), 18 deletions(-)

diff --git a/src/builtin_actions.py b/src/builtin_actions.py
index d4c1d0c59..5418e31ae 100644
--- a/src/builtin_actions.py
+++ b/src/builtin_actions.py
@@ -221,6 +221,21 @@ def _email_urgency_account_key(message_key):
     return str(message_key).split(":", 1)[0]
 
 
+def _email_urgency_payload_account_ids(state):
+    """Return account IDs that still own user-visible urgency payload."""
+    if not isinstance(state, dict):
+        return set()
+
+    per_uid = state.get("per_uid", {})
+    per_uid_keys = per_uid if isinstance(per_uid, dict) else {}
+    return {
+        _email_urgency_account_key(key) for key in per_uid_keys
+    } | {
+        _email_urgency_account_key(key)
+        for key in _email_urgency_string_set(state.get("notified_uids", []))
+    }
+
+
 def _email_urgency_stale_accounts(
     prior,
     base_account_generations,
@@ -248,6 +263,8 @@ def _merge_email_urgency_state(
     fully_scanned_account_ids,
     base_account_generations,
     timestamp,
+    retired_account_ids=(),
+    base_payload_account_ids=(),
 ):
     """Merge a scan without letting an older snapshot erase newer facts."""
     prior_per_uid = prior.get("per_uid", {})
@@ -255,18 +272,39 @@ def _merge_email_urgency_state(
         prior_per_uid = {}
     complete = {str(account_id) for account_id in fully_scanned_account_ids}
     prior_generations = _email_urgency_account_generations(prior)
+    retire_requested = {str(account_id) for account_id in retired_account_ids}
     observed_accounts = {
         _email_urgency_account_key(key) for key in per_uid_scores
-    } | complete
+    } | complete | retire_requested
     stale_accounts = _email_urgency_stale_accounts(
         prior,
         base_account_generations,
         observed_accounts,
     )
-    fresh_complete = complete - stale_accounts
+    prior_payload_accounts = _email_urgency_payload_account_ids(prior)
+    base_payload_accounts = {
+        str(account_id) for account_id in base_payload_account_ids
+    }
+    # A selected account can be absent from the base snapshot. If another
+    # worker creates its first payload before this transaction wins the lock,
+    # membership itself is a fence even when both snapshots normalize to the
+    # legacy generation zero.
+    retired_accounts = {
+        account_id
+        for account_id in retire_requested - stale_accounts
+        if not (
+            account_id in prior_payload_accounts
+            and account_id not in base_payload_accounts
+        )
+    }
+    fresh_complete = complete - stale_accounts - retired_accounts
     changed_accounts = set(fresh_complete)
 
-    merged_per_uid = dict(prior_per_uid)
+    merged_per_uid = {
+        key: value
+        for key, value in prior_per_uid.items()
+        if _email_urgency_account_key(key) not in retired_accounts
+    }
     for key in list(merged_per_uid):
         account_id = _email_urgency_account_key(key)
         if account_id in fresh_complete:
@@ -280,17 +318,21 @@ def _merge_email_urgency_state(
     # additive without another fresh scan.
     for key, value in per_uid_scores.items():
         account_id = _email_urgency_account_key(key)
-        if account_id in stale_accounts:
+        if account_id in stale_accounts or account_id in retired_accounts:
             continue
         if merged_per_uid.get(key) != value:
             changed_accounts.add(account_id)
         merged_per_uid[key] = value
 
     prior_notified = _email_urgency_string_set(prior.get("notified_uids", []))
-    merged_notified = set(prior_notified)
+    merged_notified = {
+        key
+        for key in prior_notified
+        if _email_urgency_account_key(key) not in retired_accounts
+    }
     for key in _email_urgency_string_set(notified_uids) - prior_notified:
         account_id = _email_urgency_account_key(key)
-        if account_id in stale_accounts:
+        if account_id in stale_accounts or account_id in retired_accounts:
             continue
         merged_notified.add(key)
         changed_accounts.add(account_id)
@@ -314,6 +356,19 @@ def _merge_email_urgency_state(
         generation["checkpoint"] += 1
         if account_id in fresh_complete:
             generation["complete"] += 1
+    for account_id in retired_accounts:
+        # Keep a generation-only tombstone. It prevents a scan that started
+        # before deletion/disable from resurrecting the retired account, while
+        # a later re-enabled account can capture this checkpoint and replace
+        # the tombstone normally. Do not advance an already-empty tombstone on
+        # every cleanup pass.
+        had_generation = account_id in prior_generations
+        if account_id in prior_payload_accounts or not had_generation:
+            generation = next_generations.setdefault(
+                account_id,
+                {"checkpoint": 0, "complete": 0},
+            )
+            generation["checkpoint"] += 1
 
     total_unread = 0
     total_urgent = 0
@@ -2216,16 +2271,14 @@ async def action_check_email_urgency(owner: str, **kwargs) -> Tuple[str, bool]:
             "shopping", "social", "work", "personal", "legal", "support", "promo",
         }
 
-        # ── 1. Resolve LLM candidates (utility primary + utility fallbacks; fall
-        # through to default chat as a last resort).
+        # Resolve with the task owner as before, but defer the availability
+        # gate until after authoritative account cleanup. State retirement must
+        # still run when no model is configured.
         from src.task_endpoint import resolve_task_candidates
         candidates = resolve_task_candidates(owner=owner)
-        if not candidates:
-            return "No LLM endpoint available", False
-
         target_account_id = _email_task_account_id(kwargs)
 
-        # ── 2. Enumerate enabled accounts. Match this task's owner AND fall
+        # ── 1. Enumerate enabled accounts. Match this task's owner AND fall
         # back to the legacy "unowned account whose imap_user / from_address
         # == this owner" pattern — same rule `_get_email_config` uses, so a
         # pre-multi-user account row still gets picked up for the seeded task.
@@ -2242,15 +2295,69 @@ async def action_check_email_urgency(owner: str, **kwargs) -> Tuple[str, bool]:
             accounts = q.all()
         finally:
             db.close()
+
+        # Capture the checkpoint basis before any cleanup or IMAP work. A full
+        # owner-wide enumeration authoritatively retires payload belonging to
+        # accounts that are no longer enabled/visible. A scoped task may retire
+        # only its selected missing/disabled account. Existing accounts remain
+        # present here even if their later network scan fails, so transient
+        # IMAP failure never erases their last known state.
+        base_state = _read_email_urgency_state(STATE_PATH)
+        base_account_generations = _email_urgency_account_generations(
+            base_state
+        )
+        base_payload_account_ids = _email_urgency_payload_account_ids(base_state)
+        enabled_account_ids = {str(account.id) for account in accounts}
+        if target_account_id:
+            retired_account_ids = (
+                {str(target_account_id)}
+                if str(target_account_id) not in enabled_account_ids
+                else set()
+            )
+        else:
+            retired_account_ids = (
+                base_payload_account_ids - enabled_account_ids
+            )
+
+        if retired_account_ids:
+            async def _retire_accounts(prior):
+                next_state = _merge_email_urgency_state(
+                    prior,
+                    owner=owner,
+                    per_uid_scores={},
+                    notified_uids=prior.get("notified_uids", []),
+                    all_unread_keys=set(),
+                    fully_scanned_account_ids=set(),
+                    base_account_generations=base_account_generations,
+                    timestamp=_time.time(),
+                    retired_account_ids=retired_account_ids,
+                    base_payload_account_ids=base_payload_account_ids,
+                )
+                return None, next_state
+
+            await _run_email_urgency_state_transaction(
+                STATE_PATH,
+                STATE_LOCK_DB,
+                _retire_accounts,
+            )
+            # Cleanup may have advanced tombstone generations. Capture the
+            # exact committed basis that the subsequent scan must compare.
+            base_state = _read_email_urgency_state(STATE_PATH)
+            base_account_generations = _email_urgency_account_generations(
+                base_state
+            )
+            base_payload_account_ids = _email_urgency_payload_account_ids(
+                base_state
+            )
+
         if not accounts:
             raise TaskNoop("no email accounts configured")
 
-        # Capture each account's checkpoint generation before any IMAP work.
-        # The locked merge compares this basis with the latest committed state,
-        # so an older scan cannot destructively replace a newer worker's facts.
-        base_account_generations = _email_urgency_account_generations(
-            _read_email_urgency_state(STATE_PATH)
-        )
+        # ── 2. Account retirement above is state maintenance and does not
+        # depend on model availability. Scanning still requires the utility
+        # primary/fallback candidates resolved for this task owner.
+        if not candidates:
+            return "No LLM endpoint available", False
 
         urgency_prompt = settings.get("urgent_email_prompt", "")
         per_uid_scores = {}   # key = "<acc_id>:<uid>" → {"score": 0-3, "reason": "..."}
diff --git a/tests/test_email_urgency_checkpoint.py b/tests/test_email_urgency_checkpoint.py
index 5392776ae..2053a58f7 100644
--- a/tests/test_email_urgency_checkpoint.py
+++ b/tests/test_email_urgency_checkpoint.py
@@ -289,6 +289,104 @@ def test_newer_partial_checkpoint_fences_older_complete_snapshot():
     }
 
 
+def test_authoritative_retirement_prunes_payload_and_recomputes_api_totals():
+    from src.builtin_actions import _merge_email_urgency_state
+
+    prior = {
+        "owner": "alice",
+        "per_uid": {
+            "acct-a:1": {"score": 1, "unread": True},
+            "acct-b:1": {"score": 3, "unread": True},
+        },
+        "notified_uids": ["acct-b:1"],
+        "account_generations": {
+            "acct-a": {"checkpoint": 1, "complete": 1},
+            "acct-b": {"checkpoint": 1, "complete": 1},
+        },
+    }
+
+    merged = _merge_email_urgency_state(
+        prior,
+        owner="alice",
+        per_uid_scores={},
+        notified_uids=prior["notified_uids"],
+        all_unread_keys=set(),
+        fully_scanned_account_ids=set(),
+        base_account_generations=prior["account_generations"],
+        timestamp=300.0,
+        retired_account_ids={"acct-b"},
+        base_payload_account_ids={"acct-a", "acct-b"},
+    )
+
+    assert merged["per_uid"] == {
+        "acct-a:1": {"score": 1, "unread": True},
+    }
+    assert merged["notified_uids"] == []
+    assert merged["total_unread"] == 1
+    assert merged["total_urgent"] == 0
+    assert merged["max_score"] == 1
+    assert merged["account_generations"]["acct-b"] == {
+        "checkpoint": 2,
+        "complete": 1,
+    }
+
+
+def test_authoritative_retirement_respects_generation_and_membership_fences():
+    from src.builtin_actions import _merge_email_urgency_state
+
+    newer = {
+        "owner": "alice",
+        "per_uid": {
+            "acct-b:2": {"score": 3, "unread": True, "reason": "newer"},
+        },
+        "notified_uids": ["acct-b:2"],
+        "account_generations": {
+            "acct-b": {"checkpoint": 2, "complete": 2},
+        },
+    }
+    generation_fenced = _merge_email_urgency_state(
+        newer,
+        owner="alice",
+        per_uid_scores={},
+        notified_uids=newer["notified_uids"],
+        all_unread_keys=set(),
+        fully_scanned_account_ids=set(),
+        base_account_generations={
+            "acct-b": {"checkpoint": 1, "complete": 1},
+        },
+        timestamp=300.0,
+        retired_account_ids={"acct-b"},
+        base_payload_account_ids={"acct-b"},
+    )
+    assert generation_fenced["per_uid"] == newer["per_uid"]
+    assert generation_fenced["notified_uids"] == ["acct-b:2"]
+    assert generation_fenced["account_generations"]["acct-b"] == {
+        "checkpoint": 2,
+        "complete": 2,
+    }
+
+    legacy_first_write = {
+        "owner": "alice",
+        "per_uid": {
+            "acct-b:3": {"score": 2, "unread": True, "reason": "concurrent"},
+        },
+        "notified_uids": [],
+    }
+    membership_fenced = _merge_email_urgency_state(
+        legacy_first_write,
+        owner="alice",
+        per_uid_scores={},
+        notified_uids=[],
+        all_unread_keys=set(),
+        fully_scanned_account_ids=set(),
+        base_account_generations={},
+        timestamp=300.0,
+        retired_account_ids={"acct-b"},
+        base_payload_account_ids=set(),
+    )
+    assert membership_fenced["per_uid"] == legacy_first_write["per_uid"]
+
+
 @pytest.mark.asyncio
 async def test_waiting_transaction_cancellation_does_not_leak_lock(tmp_path):
     from src.builtin_actions import _run_email_urgency_state_transaction
@@ -664,6 +762,146 @@ async def test_account_scoped_actions_merge_disjoint_checkpoints(
     assert state["total_urgent"] == 2
 
 
+@pytest.mark.asyncio
+async def test_full_enumeration_retires_deleted_or_disabled_account(
+    monkeypatch,
+    tmp_path,
+):
+    import routes.note_routes as note_routes
+
+    builtin_actions, runtime = _configure_action(
+        monkeypatch, tmp_path, ["acct-a", "acct-b"]
+    )
+
+    async def delivered(**_kwargs):
+        return {
+            "browser_sent": True,
+            "email_sent": False,
+            "ntfy_sent": False,
+            "webhook_sent": False,
+        }
+
+    monkeypatch.setattr(note_routes, "dispatch_reminder", delivered)
+    await builtin_actions.action_check_email_urgency("alice")
+
+    # The production query returns only enabled, owner-visible accounts. A
+    # deleted row and a disabled row are therefore the same authoritative
+    # absence at this boundary.
+    runtime["accounts"] = ["acct-a"]
+    await builtin_actions.action_check_email_urgency("alice")
+
+    state = json.loads(
+        (tmp_path / "email_urgency_state_alice.json").read_text(encoding="utf-8")
+    )
+    assert set(state["per_uid"]) == {"acct-a:1"}
+    assert state["notified_uids"] == ["acct-a:1"]
+    assert state["total_unread"] == 1
+    assert state["total_urgent"] == 1
+    assert state["max_score"] == 3
+    assert state["account_generations"]["acct-b"] == {
+        "checkpoint": 2,
+        "complete": 1,
+    }
+
+
+@pytest.mark.asyncio
+async def test_zero_enabled_accounts_cleanup_precedes_model_resolution(
+    monkeypatch,
+    tmp_path,
+):
+    import routes.note_routes as note_routes
+    from src import task_endpoint
+
+    builtin_actions, runtime = _configure_action(
+        monkeypatch, tmp_path, ["acct-a"]
+    )
+
+    async def delivered(**_kwargs):
+        return {
+            "browser_sent": True,
+            "email_sent": False,
+            "ntfy_sent": False,
+            "webhook_sent": False,
+        }
+
+    monkeypatch.setattr(note_routes, "dispatch_reminder", delivered)
+    await builtin_actions.action_check_email_urgency("alice")
+    runtime["accounts"] = []
+
+    model_resolution_owners = []
+
+    def no_model_available(*_args, **kwargs):
+        model_resolution_owners.append(kwargs.get("owner"))
+        return []
+
+    monkeypatch.setattr(
+        task_endpoint,
+        "resolve_task_candidates",
+        no_model_available,
+    )
+    with pytest.raises(builtin_actions.TaskNoop):
+        await builtin_actions.action_check_email_urgency("alice")
+    assert model_resolution_owners == ["alice"]
+
+    state = json.loads(
+        (tmp_path / "email_urgency_state_alice.json").read_text(encoding="utf-8")
+    )
+    assert state["per_uid"] == {}
+    assert state["notified_uids"] == []
+    assert state["total_unread"] == 0
+    assert state["total_urgent"] == 0
+    assert state["max_score"] == 0
+    assert state["account_generations"]["acct-a"] == {
+        "checkpoint": 2,
+        "complete": 1,
+    }
+
+
+@pytest.mark.asyncio
+async def test_scoped_missing_account_retires_only_selected_payload(
+    monkeypatch,
+    tmp_path,
+):
+    import routes.note_routes as note_routes
+
+    builtin_actions, runtime = _configure_action(
+        monkeypatch, tmp_path, ["acct-a", "acct-b"]
+    )
+
+    async def delivered(**_kwargs):
+        return {
+            "browser_sent": True,
+            "email_sent": False,
+            "ntfy_sent": False,
+            "webhook_sent": False,
+        }
+
+    monkeypatch.setattr(note_routes, "dispatch_reminder", delivered)
+    await builtin_actions.action_check_email_urgency("alice")
+
+    # A production account-id filter returns no row when the selected account
+    # was deleted or disabled. Other accounts are outside this scoped query and
+    # must remain untouched.
+    runtime["accounts"] = []
+    with pytest.raises(builtin_actions.TaskNoop):
+        await builtin_actions.action_check_email_urgency(
+            "alice",
+            prompt='{"account_id":"acct-b"}',
+        )
+
+    state = json.loads(
+        (tmp_path / "email_urgency_state_alice.json").read_text(encoding="utf-8")
+    )
+    assert set(state["per_uid"]) == {"acct-a:1"}
+    assert state["notified_uids"] == ["acct-a:1"]
+    assert state["total_unread"] == 1
+    assert state["total_urgent"] == 1
+    assert state["account_generations"]["acct-b"] == {
+        "checkpoint": 2,
+        "complete": 1,
+    }
+
+
 @pytest.mark.asyncio
 async def test_failed_account_scan_preserves_checkpoint_until_recovery(
     monkeypatch,
@@ -696,6 +934,12 @@ async def test_failed_account_scan_preserves_checkpoint_until_recovery(
     )
     assert failed_state["notified_uids"] == ["acct-a:1", "acct-b:1"]
     assert set(failed_state["per_uid"]) == {"acct-a:1", "acct-b:1"}
+    assert failed_state["total_unread"] == 2
+    assert failed_state["total_urgent"] == 2
+    assert failed_state["account_generations"]["acct-b"] == {
+        "checkpoint": 1,
+        "complete": 1,
+    }
 
     runtime["failures"] = set()
     runtime["accounts"] = ["acct-b"]