From e37dc772f3f19b72c4eca141065ae29716ac5877 Mon Sep 17 00:00:00 2001 From: QueryPlanner Date: Wed, 19 Aug 2026 14:26:08 +0530 Subject: [PATCH] fix: harden album buffer error paths - Resolve album.future inside _handle_album_turn itself, not only via the caller's finally block - Roll back the buffer entry if scheduling the max-wait watchdog task raises during add_message - Log the suppressed album-wait exception at warning, not debug - Document the sync invariant that keeps debounce/max-wait flush race-free Co-Authored-By: Claude Sonnet 5 --- src/blacki/telegram/album_buffer.py | 20 ++++++++- src/blacki/telegram/bot.py | 23 ++++++++++- tests/test_telegram_album_buffer.py | 63 +++++++++++++++++++++++++++++ tests/test_telegram_bot.py | 54 +++++++++++++++++++++++-- 4 files changed, 153 insertions(+), 7 deletions(-) diff --git a/src/blacki/telegram/album_buffer.py b/src/blacki/telegram/album_buffer.py index e129c9a..07189c3 100644 --- a/src/blacki/telegram/album_buffer.py +++ b/src/blacki/telegram/album_buffer.py @@ -84,7 +84,13 @@ async def add_message(self, message: Message, seq: int) -> None: ) self._buffers[key] = album - max_wait_task = asyncio.create_task(self._max_wait(album)) + try: + max_wait_task = asyncio.create_task(self._max_wait(album)) + except Exception: + # Scheduling failed before any timer exists to flush this + # album later, so don't leave it orphaned in self._buffers. + self._buffers.pop(key, None) + raise self._background_tasks.add(max_wait_task) max_wait_task.add_done_callback(self._background_tasks.discard) album.max_wait_task = max_wait_task @@ -119,7 +125,17 @@ def _on_debounce_expired(self, album: _BufferedAlbum) -> None: self._flush(album) def _flush(self, album: _BufferedAlbum) -> None: - """Mark an album processed, remove it from the buffer, and flush it.""" + """Mark an album processed, remove it from the buffer, and flush it. + + Both the debounce callback and the max-wait task call this method, + so the check-then-set on ``album.processed`` below is the only guard + against double-flushing. It is race-free only because this method + is synchronous with no ``await`` before ``album.processed = True``: + the event loop cannot interleave the check and the set. Do not add + an ``await`` before that line, or make this method ``async``, + without replacing the guard with something that is still atomic + against both callers. + """ if album.processed: return album.processed = True diff --git a/src/blacki/telegram/bot.py b/src/blacki/telegram/bot.py index 27410bc..138a9d2 100644 --- a/src/blacki/telegram/bot.py +++ b/src/blacki/telegram/bot.py @@ -303,7 +303,13 @@ async def _safe_handle_update(self, update: Update) -> None: if current_task is not None and current_task.cancelling() > 0: raise except Exception as exc: - logger.debug("Album buffer wait suppressed error: %s", exc) + # A failed album turn already reports its own user-facing + # error via _send_photo_error inside _handle_album_turn. + # This wait only unblocks processing of the next message + # in this conversation, so it must not raise on the + # album's behalf, but it is logged at warning level (not + # debug) so an unexpected failure here stays visible. + logger.warning("Album buffer wait suppressed error: %s", exc) await self._run_sequenced_turn( conversation_key, current_seq, self._handle_update(update) @@ -371,7 +377,20 @@ async def _process_flushed_album(self, album: _BufferedAlbum) -> None: album.future.set_result(None) async def _handle_album_turn(self, album: _BufferedAlbum) -> None: - """Download album images and run a single ADK turn.""" + """Download album images and run a single ADK turn. + + Resolves ``album.future`` before returning on every exit path, so + this method is safe to call directly and does not rely on a caller + (``_process_flushed_album``) to resolve the future on its behalf. + """ + try: + await self._run_album_turn(album) + finally: + if album.future is not None and not album.future.done(): + album.future.set_result(None) + + async def _run_album_turn(self, album: _BufferedAlbum) -> None: + """Validate, download, and process a flushed album's photos.""" chat_id = album.chat_id message_thread_id = album.message_thread_id messages = album.messages diff --git a/tests/test_telegram_album_buffer.py b/tests/test_telegram_album_buffer.py index 269563e..5600eb4 100644 --- a/tests/test_telegram_album_buffer.py +++ b/tests/test_telegram_album_buffer.py @@ -99,6 +99,46 @@ async def test_buffer_album_message_branches(buffer: AlbumBuffer) -> None: album.debounce_handle.cancel() +@pytest.mark.asyncio +async def test_add_message_cleans_up_on_watchdog_schedule_failure( + buffer: AlbumBuffer, +) -> None: + """If scheduling the max-wait watchdog raises, the album must not be + left orphaned in self._buffers with no timer to ever flush it + (Github issue #163, item 2).""" + msg = Message.model_validate( + { + "message_id": 1, + "date": "2024-01-01T00:00:00Z", + "chat": {"id": 123, "type": "private"}, + "media_group_id": "schedule-fail", + "photo": [ + { + "file_id": "p1", + "file_unique_id": "u1", + "width": 10, + "height": 10, + } + ], + } + ) + + def _fail_to_schedule(coro): + coro.close() # avoid a "coroutine was never awaited" warning + raise RuntimeError("event loop is shutting down") + + with ( + patch( + "blacki.telegram.album_buffer.asyncio.create_task", + side_effect=_fail_to_schedule, + ), + pytest.raises(RuntimeError, match="event loop is shutting down"), + ): + await buffer.add_message(msg, 1) + + assert (123, None, "schedule-fail") not in buffer._buffers + + @pytest.mark.asyncio async def test_album_max_wait_triggers_flush(buffer: AlbumBuffer) -> None: """Test _max_wait completing its sleep and triggering flush.""" @@ -142,3 +182,26 @@ async def test_flush_album_branches(buffer: AlbumBuffer) -> None: ) buffer._flush(album_none_handles) assert album_none_handles.processed is True + + +def test_flush_is_idempotent_against_debounce_max_wait_race( + buffer: AlbumBuffer, on_flush: MagicMock +) -> None: + """Both the debounce callback and the max-wait task call _flush on the + same album; this pins _flush's check-then-set as a single synchronous + call with no ``await`` in between, so calling it twice back-to-back + (as if both timers fired in the same event-loop iteration) must flush + exactly once (Github issue #163, item 4).""" + album = _BufferedAlbum( + chat_id=123, + message_thread_id=None, + media_group_id="race", + chat_type=ChatType.PRIVATE, + messages=[], + ) + buffer._buffers[(123, None, "race")] = album + + buffer._on_debounce_expired(album) + buffer._flush(album) + + assert on_flush.call_count == 1 diff --git a/tests/test_telegram_bot.py b/tests/test_telegram_bot.py index 72d73d0..7559ce7 100644 --- a/tests/test_telegram_bot.py +++ b/tests/test_telegram_bot.py @@ -4707,6 +4707,46 @@ async def test_album_empty_photos_list_rejected( assert len(runtime_recorder.run_user_turn_calls) == 0 assert "failed to process" in mock_api.send_message.await_args.kwargs["text"] + @pytest.mark.asyncio + async def test_album_turn_resolves_future_on_validation_early_return( + self, + telegram_config: TelegramConfig, + runtime_recorder: RecordingRuntime, + ) -> None: + """_handle_album_turn resolves album.future itself on every early + return, even when called directly rather than through + _process_flushed_album (Github issue #163, item 1).""" + bot = TelegramBot(telegram_config, cast(AdkRuntime, runtime_recorder)) + mock_api = create_autospec(TelegramApiClient, instance=True) + mock_api.send_message = AsyncMock() + bot._api = mock_api + + loop = asyncio.get_running_loop() + future: asyncio.Future[None] = loop.create_future() + album = _BufferedAlbum( + chat_id=123, + message_thread_id=None, + media_group_id="empty-photos-future", + chat_type=ChatType.PRIVATE, + messages=[ + Message.model_validate( + { + "message_id": 1, + "date": "2024-01-01T00:00:00Z", + "chat": {"id": 123, "type": "private"}, + "media_group_id": "empty-photos-future", + "photo": [], + } + ) + ], + future=future, + ) + + await bot._handle_album_turn(album) + + assert future.done() + assert future.result() is None + @pytest.mark.asyncio async def test_album_followed_by_text_during_debounce_preserves_order( self, @@ -4922,6 +4962,7 @@ async def test_safe_handle_update_active_album_future_branches( self, telegram_config: TelegramConfig, runtime_recorder: RecordingRuntime, + caplog: pytest.LogCaptureFixture, ) -> None: """Test _safe_handle_update active album buffer with None future or error.""" bot = TelegramBot(telegram_config, cast(AdkRuntime, runtime_recorder)) @@ -4968,10 +5009,17 @@ async def test_safe_handle_update_active_album_future_branches( ) bot._album_buffer._buffers[(123, None, "err-fut")] = album_err_future - await bot._safe_handle_update( - Update.model_validate({"update_id": 2, "message": text_msg.model_dump()}) - ) + with caplog.at_level(logging.WARNING, logger="blacki.telegram.bot"): + await bot._safe_handle_update( + Update.model_validate( + {"update_id": 2, "message": text_msg.model_dump()} + ) + ) assert len(runtime_recorder.run_user_turn_calls) == 2 + assert any( + "Album buffer wait suppressed error" in record.message + for record in caplog.records + ) # Case 3: active album future cancelled while current_task is cancelling cancel_future: asyncio.Future[None] = loop.create_future()