Ver Fonte

fix: reliably clean up check-in interactions

AI-Co-Authored-By: Codex
chendeben há 2 meses atrás
pai
commit
34227d0488

+ 26 - 0
tests/test_message_cleanup.py

@@ -40,9 +40,11 @@ class FakeClient:
 def cleanup(app_modules):
     module = app_modules.load("wbb.services.message_cleanup")
     module._scheduled_messages.clear()
+    module._cleanup_tasks.clear()
     module._active_interactions.clear()
     yield module
     module._scheduled_messages.clear()
+    module._cleanup_tasks.clear()
     module._active_interactions.clear()
 
 
@@ -73,6 +75,30 @@ async def test_private_messages_are_not_scheduled(cleanup):
     assert message.delete_calls == []
 
 
+async def test_cleanup_task_is_referenced_until_deletion_finishes(
+    cleanup, monkeypatch
+):
+    real_sleep = asyncio.sleep
+    release = asyncio.Event()
+
+    async def blocked_sleep(_delay: int) -> None:
+        await release.wait()
+
+    monkeypatch.setattr(cleanup.asyncio, "sleep", blocked_sleep)
+    message = FakeMessage(-1001, 17, text="签到")
+
+    cleanup.schedule_group_cleanup(message)
+    await real_sleep(0)
+
+    assert len(cleanup._cleanup_tasks) == 1
+    release.set()
+    await real_sleep(0)
+    await real_sleep(0)
+
+    assert message.delete_calls == [True]
+    assert cleanup._cleanup_tasks == set()
+
+
 async def test_command_context_tracks_only_responses_in_the_same_group(
     cleanup, monkeypatch
 ):

+ 48 - 3
tests/test_point_settings.py

@@ -23,9 +23,11 @@ async def test_enabling_checkin_button_sends_hideable_keyboard(
     )
 
     assert rules["checkin_button_enabled"] is True
-    send_message.assert_awaited_once()
-    chat_id, text = send_message.await_args.args
-    reply_markup = send_message.await_args.kwargs["reply_markup"]
+    assert send_message.await_count == 2
+    remove_call, keyboard_call = send_message.await_args_list
+    assert isinstance(remove_call.kwargs["reply_markup"], ReplyKeyboardRemove)
+    chat_id, text = keyboard_call.args
+    reply_markup = keyboard_call.kwargs["reply_markup"]
     assert chat_id == -100
     assert "签到快捷按钮已开启" in text
     assert isinstance(reply_markup, ReplyKeyboardMarkup)
@@ -33,6 +35,8 @@ async def test_enabling_checkin_button_sends_hideable_keyboard(
     assert reply_markup.resize_keyboard is True
     assert reply_markup.one_time_keyboard is True
     assert reply_markup.keyboard[0][0].text == "签到"
+    stored = await points.rulesdb.find_one({"chat_id": -100})
+    assert stored["checkin_keyboard_version"] == settings.CHECKIN_KEYBOARD_VERSION
 
 
 async def test_checkin_button_sync_is_idempotent_and_follows_points_state(
@@ -80,3 +84,44 @@ async def test_checkin_button_send_failure_rolls_back_rule(app_modules, monkeypa
 
     rules = await points.get_point_rules(-100)
     assert rules["checkin_button_enabled"] is False
+
+
+async def test_active_checkin_keyboards_are_migrated_only_once(
+    app_modules, monkeypatch
+):
+    points = app_modules.load("wbb.utils.dbpoints")
+    settings = app_modules.load("wbb.services.point_settings")
+    await points.set_point_rules(
+        -100,
+        {
+            "enabled": True,
+            "checkin_enabled": True,
+            "checkin_button_enabled": True,
+        },
+    )
+    send_message = AsyncMock(
+        side_effect=[
+            SimpleNamespace(id=1),
+            SimpleNamespace(id=2),
+        ]
+    )
+    monkeypatch.setattr(settings.app, "send_message", send_message)
+
+    refreshed = await settings.refresh_active_checkin_keyboards()
+
+    assert refreshed == 1
+    assert send_message.await_count == 2
+    assert isinstance(
+        send_message.await_args_list[0].kwargs["reply_markup"],
+        ReplyKeyboardRemove,
+    )
+    assert isinstance(
+        send_message.await_args_list[1].kwargs["reply_markup"],
+        ReplyKeyboardMarkup,
+    )
+    stored = await points.rulesdb.find_one({"chat_id": -100})
+    assert stored["checkin_keyboard_version"] == settings.CHECKIN_KEYBOARD_VERSION
+
+    send_message.reset_mock()
+    assert await settings.refresh_active_checkin_keyboards() == 0
+    send_message.assert_not_awaited()

+ 7 - 0
wbb/modules/points.py

@@ -16,6 +16,7 @@ from wbb.services.point_settings import (
     CHECKIN_BUTTON_TEXT,
     apply_point_rules,
     checkin_button_active,
+    refresh_active_checkin_keyboards,
 )
 from wbb.utils.dbpoints import (
     SOURCE_ADMIN,
@@ -338,6 +339,12 @@ async def points_activity_listener(_, message: Message):
 
 
 async def _reconcile_points_loop() -> None:
+    try:
+        refreshed = await refresh_active_checkin_keyboards()
+        if refreshed:
+            log.info(f"已刷新 {refreshed} 个群的签到快捷按钮")
+    except Exception as exc:
+        log.error(f"刷新签到快捷按钮失败:{exc}")
     await asyncio.sleep(60)
     while True:
         try:

+ 4 - 1
wbb/services/message_cleanup.py

@@ -10,6 +10,7 @@ GROUP_INTERACTION_TTL_SECONDS = 60
 
 _LOGGER = logging.getLogger(__name__)
 _scheduled_messages: set[tuple[int, int]] = set()
+_cleanup_tasks: set[asyncio.Task[Any]] = set()
 _active_interactions: WeakKeyDictionary[asyncio.Task[Any], int] = WeakKeyDictionary()
 _TRACKED_SEND_METHODS = (
     "send_message",
@@ -98,10 +99,12 @@ def schedule_group_cleanup(
         if key in _scheduled_messages:
             continue
         _scheduled_messages.add(key)
-        asyncio.create_task(
+        task = asyncio.create_task(
             _delete_later(message, delay_seconds),
             name=f"group-interaction-cleanup:{key[0]}:{key[1]}",
         )
+        _cleanup_tasks.add(task)
+        task.add_done_callback(_cleanup_tasks.discard)
 
 
 def is_group_cleanup_scheduled(message: Any) -> bool:

+ 51 - 8
wbb/services/point_settings.py

@@ -4,15 +4,19 @@ from typing import Any
 
 from pyrogram.types import KeyboardButton, ReplyKeyboardMarkup, ReplyKeyboardRemove
 
-from wbb import app
+from wbb import app, log
+from wbb.services.message_cleanup import schedule_group_cleanup
 from wbb.utils.dbpoints import (
     PointsError,
     get_point_rules,
+    list_checkin_keyboard_migrations,
+    mark_checkin_keyboard_version,
     normalize_point_rules,
     set_point_rules,
 )
 
 CHECKIN_BUTTON_TEXT = "签到"
+CHECKIN_KEYBOARD_VERSION = 1
 
 
 def checkin_button_active(rules: dict[str, Any]) -> bool:
@@ -35,18 +39,52 @@ def checkin_keyboard() -> ReplyKeyboardMarkup:
 
 
 async def _sync_checkin_keyboard(chat_id: int, enabled: bool) -> None:
-    if enabled:
-        text = "签到快捷按钮已开启。点击下方“签到”即可领取每日积分。"
-        reply_markup = checkin_keyboard()
-    else:
-        text = "签到快捷按钮已关闭。"
-        reply_markup = ReplyKeyboardRemove(selective=False)
+    messages = []
     try:
-        await app.send_message(chat_id, text, reply_markup=reply_markup)
+        if enabled:
+            messages.append(
+                await app.send_message(
+                    chat_id,
+                    "正在更新签到快捷按钮。",
+                    reply_markup=ReplyKeyboardRemove(selective=False),
+                )
+            )
+            messages.append(
+                await app.send_message(
+                    chat_id,
+                    "签到快捷按钮已开启。点击下方“签到”即可领取每日积分。",
+                    reply_markup=checkin_keyboard(),
+                )
+            )
+        else:
+            messages.append(
+                await app.send_message(
+                    chat_id,
+                    "签到快捷按钮已关闭。",
+                    reply_markup=ReplyKeyboardRemove(selective=False),
+                )
+            )
     except Exception as exc:
         raise PointsError(
             "无法更新群底部签到按钮,请确认机器人可以在群内发言。"
         ) from exc
+    schedule_group_cleanup(*messages)
+
+
+async def refresh_active_checkin_keyboards() -> int:
+    chat_ids = await list_checkin_keyboard_migrations(CHECKIN_KEYBOARD_VERSION)
+    refreshed = 0
+    for chat_id in chat_ids:
+        try:
+            await _sync_checkin_keyboard(chat_id, True)
+            await mark_checkin_keyboard_version(
+                chat_id,
+                CHECKIN_KEYBOARD_VERSION,
+            )
+            refreshed += 1
+        except Exception as exc:
+            log.error(f"刷新群 {chat_id} 的签到快捷按钮失败:{exc}")
+    return refreshed
 
 
 async def apply_point_rules(chat_id: int, values: dict[str, Any]) -> dict[str, Any]:
@@ -60,6 +98,11 @@ async def apply_point_rules(chat_id: int, values: dict[str, Any]) -> dict[str, A
         return saved
     try:
         await _sync_checkin_keyboard(chat_id, target_active)
+        if target_active:
+            await mark_checkin_keyboard_version(
+                chat_id,
+                CHECKIN_KEYBOARD_VERSION,
+            )
     except Exception:
         await set_point_rules(chat_id, current)
         raise

+ 36 - 0
wbb/utils/dbpoints.py

@@ -210,6 +210,42 @@ async def set_point_rules(chat_id: int, values: dict[str, Any]) -> dict[str, Any
     return normalized
 
 
+async def list_checkin_keyboard_migrations(
+    keyboard_version: int,
+) -> list[int]:
+    await ensure_point_indexes()
+    cursor = rulesdb.find(
+        {
+            "enabled": True,
+            "checkin_enabled": True,
+            "checkin_button_enabled": True,
+            "$or": [
+                {"checkin_keyboard_version": {"$exists": False}},
+                {"checkin_keyboard_version": {"$lt": int(keyboard_version)}},
+            ],
+        },
+        {"chat_id": 1},
+    )
+    documents = await cursor.to_list(length=10000)
+    return [int(document["chat_id"]) for document in documents]
+
+
+async def mark_checkin_keyboard_version(
+    chat_id: int,
+    keyboard_version: int,
+) -> None:
+    await ensure_point_indexes()
+    await rulesdb.update_one(
+        {"chat_id": int(chat_id)},
+        {
+            "$set": {
+                "checkin_keyboard_version": int(keyboard_version),
+                "updated_at": utc_now(),
+            }
+        },
+    )
+
+
 async def _reconcile_account_locked(chat_id: int, user_id: int) -> dict[str, Any]:
     cached = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id})
     pipeline = [