feat(connect): delete QR code message after connect flow completes

- Track QR photo message_id via unified sensitive_msg_ids list in _qr_login_tasks
- Delete QR message in _poll_qr_login finally block (success, timeout, error, cancel)
- Add sensitive_msg_ids param to _register_qr_login_task and _poll_qr_login
- Simplify on_connect_cancel_callback: let background task finally handle cleanup
- Add architectural comments explaining three connect-flow state dicts
- Add tests: QR message deleted on success, timeout, and task cancellation
- Update existing _poll_qr_login tests for new sensitive_msg_ids param
- Update test_cancels_qr_task for new _qr_login_tasks dict structure

Author: oponfil

handlers/connect_handler.py

@@ -35,15 +35,22 @@
 from utils.telegram_user import upsert_effective_user
 
 
-# ====== /connect ======
-
-_qr_login_tasks: dict[int, asyncio.Task] = {}
-
-# Ожидающие 2FA-пароля: {user_id: {client, language_code, bot, chat_id}}
+# --- Состояние connect-flow ---
+# Три отдельных dict'а: QR-flow (фоновый polling + cleanup) и phone-flow (многошаговая
+# state-machine с таймаутами) имеют разную логику и жизненный цикл.
+# _pending_2fa — промежуточная фаза QR-flow: после сканирования QR потребовался 2FA-пароль,
+# polling-задача уже завершилась, но клиент остаётся живым для check_password.
+
+# QR-flow: фоновая задача polling + ID сообщений для cleanup
+# {user_id: {"task": Task, "sensitive_msg_ids": list[int], "chat_id": int}}
+_qr_login_tasks: dict[int, dict] = {}
+
+# QR-flow, фаза 2FA: клиент ожидает ввода пароля после успешного сканирования QR
+# {user_id: {"client": Client, "language_code": str, "bot": Bot, "chat_id": int}}
 _pending_2fa: dict[int, dict] = {}
 
-# Ожидающие phone-логина: {user_id: {state, client, phone_number, phone_code_hash, expires_at, ...}}
-# state: "awaiting_phone" | "awaiting_code" | "awaiting_2fa"
+# Phone-flow: многошаговая state-machine (ввод номера → подтверждение → код → 2FA)
+# {user_id: {"state": str, "client": Client, ..., "sensitive_msg_ids": list[int], "expires_at": float}}
 _pending_phone: dict[int, dict] = {}
 
 
@@ -124,8 +131,11 @@ async def cancel_pending_phone(
 
 def _get_qr_login_task(user_id: int) -> asyncio.Task | None:
     """Возвращает активную QR-задачу пользователя."""
-    task = _qr_login_tasks.get(user_id)
-    if task and task.done():
+    entry = _qr_login_tasks.get(user_id)
+    if entry is None:
+        return None
+    task = entry["task"]
+    if task.done():
         _qr_login_tasks.pop(user_id, None)
         return None
     return task
@@ -152,13 +162,18 @@ async def clear_pending_input(context: ContextTypes.DEFAULT_TYPE, user_id: int,
     await cancel_pending_phone(user_id, bot=bot)
 
 
-def _register_qr_login_task(user_id: int, task: asyncio.Task) -> None:
+def _register_qr_login_task(
+    user_id: int, task: asyncio.Task,
+    sensitive_msg_ids: list[int] | None = None, chat_id: int | None = None,
+) -> None:
     """Регистрирует фоновую QR-задачу до её завершения."""
-    _qr_login_tasks[user_id] = task
+    _qr_login_tasks[user_id] = {
+        "task": task, "sensitive_msg_ids": sensitive_msg_ids or [], "chat_id": chat_id,
+    }
 
     def _cleanup(done_task: asyncio.Task) -> None:
-        current_task = _qr_login_tasks.get(user_id)
-        if current_task is done_task:
+        entry = _qr_login_tasks.get(user_id)
+        if entry is not None and entry["task"] is done_task:
             _qr_login_tasks.pop(user_id, None)
 
     task.add_done_callback(_cleanup)
@@ -290,16 +305,20 @@ async def _start_qr_flow(
         keyboard = InlineKeyboardMarkup([
             [InlineKeyboardButton(cancel_label, callback_data="connect:cancel")]
         ])
-        await bot.send_photo(chat_id=chat_id, photo=buf, caption=msg, reply_markup=keyboard)
+        qr_sent = await bot.send_photo(chat_id=chat_id, photo=buf, caption=msg, reply_markup=keyboard)
+        sensitive_msg_ids = []
+        qr_mid = getattr(qr_sent, "message_id", None)
+        if qr_mid is not None:
+            sensitive_msg_ids.append(qr_mid)
 
         if DEBUG_PRINT:
             print(f"{get_timestamp()} [CONNECT_QR] QR sent to user {user_id}, waiting for scan...")
 
         # Запускаем polling в фоне
         task = asyncio.create_task(
-            _poll_qr_login(client, user_id, language_code, bot, chat_id)
+            _poll_qr_login(client, user_id, language_code, bot, chat_id, sensitive_msg_ids=sensitive_msg_ids)
         )
-        _register_qr_login_task(user_id, task)
+        _register_qr_login_task(user_id, task, sensitive_msg_ids=sensitive_msg_ids, chat_id=chat_id)
         client = None
 
     except Exception as e:
@@ -584,7 +603,7 @@ async def on_connect_cancel_callback(update: Update, context: ContextTypes.DEFAU
     u = update.effective_user
     language_code = u.language_code
 
-    # Отменяем QR-flow
+    # Отменяем QR-flow; cleanup выполнит сама фоновая задача в finally
     qr_task = _get_qr_login_task(u.id)
     if qr_task is not None:
         qr_task.cancel()
@@ -779,7 +798,10 @@ async def _finalize_phone_login(
             await _delete_sensitive_messages(bot, chat_id, sensitive_msg_ids)
 
 
-async def _poll_qr_login(client: Client, user_id: int, language_code: str, bot: Bot, chat_id: int) -> None:
+async def _poll_qr_login(
+    client: Client, user_id: int, language_code: str, bot: Bot, chat_id: int,
+    sensitive_msg_ids: list[int] | None = None,
+) -> None:
     """Фоновая задача: ожидает сканирования QR-кода (до 2 минут)."""
     try:
         authorized = False
@@ -910,6 +932,9 @@ async def _poll_qr_login(client: Client, user_id: int, language_code: str, bot:
         # Не отключаем клиент, если он передан в _pending_2fa (ожидает 2FA-пароль)
         if user_id not in _pending_2fa:
             await _safe_disconnect_temp_client(client, user_id)
+        # Удаляем чувствительные сообщения QR-flow (QR-фото и др.)
+        if sensitive_msg_ids:
+            await _delete_sensitive_messages(bot, chat_id, sensitive_msg_ids)
 
 
 async def handle_2fa_password(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:

tests/test_connect_flow.py

@@ -387,11 +387,13 @@ async def test_cancels_phone_flow(self):
 
     @pytest.mark.asyncio
     async def test_cancels_qr_task(self):
-        """Отмена QR-flow → отменяет задачу."""
+        """Отмена QR-flow → отменяет задачу и убирает её из реестра."""
         user_id = 123456
         mock_task = MagicMock()
         mock_task.done.return_value = False
-        _qr_login_tasks[user_id] = mock_task
+        _qr_login_tasks[user_id] = {
+            "task": mock_task, "sensitive_msg_ids": [500], "chat_id": user_id,
+        }
 
         update = _make_callback_update(user_id=user_id)
         context = _make_context()
@@ -401,6 +403,7 @@ async def test_cancels_qr_task(self):
 
         mock_task.cancel.assert_called_once()
         assert user_id not in _qr_login_tasks
+        context.bot.delete_message.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_cancels_2fa_flow(self):

tests/test_handlers.py

@@ -1,5 +1,6 @@
 # tests/test_handlers.py — Тесты для handlers/pyrogram_handlers.py
 
+import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -36,6 +37,7 @@
 
 from config import DEFAULT_STYLE, STYLE_TO_EMOJI
 TYPING_TEXT = SYSTEM_MESSAGES["draft_typing"].format(emoji=STYLE_TO_EMOJI[DEFAULT_STYLE])
+REAL_ASYNCIO_SLEEP = asyncio.sleep
 
 
 def _close_coroutine_task(coro):
@@ -337,13 +339,15 @@ async def test_poll_qr_login_success_saves_session_and_starts_listening(self, mo
              patch("handlers.connect_handler.get_system_message", new_callable=AsyncMock, return_value="Connected"):
             mock_pc.start_listening = AsyncMock(return_value=True)
 
-            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456)
+            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
 
         mock_client.export_session_string.assert_called_once()
         mock_client.disconnect.assert_called_once()
         mock_save_session.assert_called_once_with(123, "session-123")
         mock_pc.start_listening.assert_called_once_with(123, "session-123")
         mock_bot.send_message.assert_called_once_with(chat_id=456, text="Connected")
+        # QR-сообщение удалено
+        mock_bot.delete_message.assert_called_once_with(chat_id=456, message_id=500)
 
     @pytest.mark.asyncio
     async def test_poll_qr_login_stops_when_save_session_fails(self, mock_bot):
@@ -362,7 +366,7 @@ async def test_poll_qr_login_stops_when_save_session_fails(self, mock_bot):
              patch("handlers.connect_handler.get_system_message", new_callable=AsyncMock, return_value="Connect failed"):
             mock_pc.start_listening = AsyncMock(return_value=True)
 
-            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456)
+            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
 
         mock_pc.start_listening.assert_not_called()
         mock_bot.send_message.assert_called_once_with(chat_id=456, text="Connect failed")
@@ -385,7 +389,7 @@ async def test_poll_qr_login_clears_session_when_listener_start_fails(self, mock
              patch("handlers.connect_handler.get_system_message", new_callable=AsyncMock, return_value="Connect failed"):
             mock_pc.start_listening = AsyncMock(return_value=False)
 
-            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456)
+            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
 
         mock_clear.assert_called_once_with(123)
         mock_bot.send_message.assert_called_once_with(chat_id=456, text="Connect failed")
@@ -409,7 +413,7 @@ async def test_poll_qr_login_stores_user_from_success_result(self, mock_bot):
              patch("handlers.connect_handler.get_system_message", new_callable=AsyncMock, return_value="OK"):
             mock_pc.start_listening = AsyncMock(return_value=True)
 
-            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456)
+            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
 
         mock_client.storage.user_id.assert_called_with(777)
         mock_client.storage.is_bot.assert_called_with(False)
@@ -437,17 +441,67 @@ async def test_poll_qr_migration_2fa_enters_pending_2fa(self, mock_bot):
             mock_auth_cls.return_value.create = AsyncMock(return_value=b"new_key")
             mock_pyro_session.return_value = AsyncMock()
 
-            await _poll_qr_login(mock_client, 42, "en", mock_bot, 456)
+            await _poll_qr_login(mock_client, 42, "en", mock_bot, 456, sensitive_msg_ids=[500])
 
         # Клиент сохранён в _pending_2fa (не отключен)
         assert 42 in _pending_2fa
         assert _pending_2fa[42]["client"] is mock_client
         mock_bot.send_message.assert_called_once_with(chat_id=456, text="Enter 2FA password")
         mock_client.disconnect.assert_not_called()
+        # QR-сообщение удалено даже при переходе в 2FA
+        mock_bot.delete_message.assert_called_once_with(chat_id=456, message_id=500)
 
         # Cleanup
         _pending_2fa.pop(42, None)
 
+    @pytest.mark.asyncio
+    async def test_poll_qr_login_deletes_qr_message_on_timeout(self, mock_bot):
+        """QR-сообщение удаляется при таймауте."""
+        # Все poll-ы возвращают LoginToken (не авторизован)
+        login_token = type("LoginToken", (), {"token": b"tok"})()
+
+        mock_client = AsyncMock()
+        mock_client.invoke = AsyncMock(return_value=login_token)
+        mock_client.disconnect = AsyncMock()
+
+        with patch("handlers.connect_handler.asyncio.sleep", new_callable=AsyncMock), \
+             patch("handlers.connect_handler.get_system_message", new_callable=AsyncMock, return_value="QR expired"):
+
+            await _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
+
+        # QR-сообщение удалено при таймауте
+        mock_bot.delete_message.assert_called_once_with(chat_id=456, message_id=500)
+
+    @pytest.mark.asyncio
+    async def test_poll_qr_login_deletes_qr_message_on_task_cancel(self, mock_bot):
+        """При отмене фоновой QR-задачи cleanup удаляет QR-сообщение ровно один раз."""
+        login_token = type("LoginToken", (), {"token": b"tok"})()
+
+        mock_client = AsyncMock()
+        mock_client.invoke = AsyncMock(return_value=login_token)
+        mock_client.disconnect = AsyncMock()
+
+        sleep_started = asyncio.Event()
+        release_sleep = asyncio.Event()
+
+        async def blocked_sleep(_: int) -> None:
+            sleep_started.set()
+            await release_sleep.wait()
+
+        with patch("handlers.connect_handler.asyncio.sleep", side_effect=blocked_sleep):
+            task = asyncio.create_task(
+                _poll_qr_login(mock_client, 123, "en", mock_bot, 456, sensitive_msg_ids=[500])
+            )
+            await sleep_started.wait()
+            await REAL_ASYNCIO_SLEEP(0)
+            task.cancel()
+
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+        mock_client.disconnect.assert_called_once()
+        mock_bot.delete_message.assert_called_once_with(chat_id=456, message_id=500)
+
 
 class TestOnPyrogramMessage:
     """Тесты для on_pyrogram_message()."""