diff --git a/src/ccbot/bot.py b/src/ccbot/bot.py index 809efcc9d..5172ec7bd 100644 --- a/src/ccbot/bot.py +++ b/src/ccbot/bot.py @@ -1802,6 +1802,12 @@ async def handle_new_message(msg: NewMessage, bot: Bot) -> None: # Enqueue content message task # Note: tool_result editing is handled inside _process_content_task # to ensure sequential processing with tool_use message sending + # + # Everything except the final assistant text answer (thinking, + # tool_use, tool_result, local_command, user-message echo) is + # "secondary" — it gets deleted once the next message for this + # topic arrives, so only the actual answer is left behind. + is_secondary = not (msg.role == "assistant" and msg.content_type == "text") await enqueue_content_message( bot=bot, user_id=user_id, @@ -1812,6 +1818,7 @@ async def handle_new_message(msg: NewMessage, bot: Bot) -> None: text=msg.text, thread_id=thread_id, image_data=msg.image_data, + is_secondary=is_secondary, ) # Update user's read offset to current file position diff --git a/src/ccbot/handlers/cleanup.py b/src/ccbot/handlers/cleanup.py index cd7e75805..4751841be 100644 --- a/src/ccbot/handlers/cleanup.py +++ b/src/ccbot/handlers/cleanup.py @@ -12,7 +12,11 @@ from telegram import Bot from .interactive_ui import clear_interactive_msg -from .message_queue import clear_status_msg_info, clear_tool_msg_ids_for_topic +from .message_queue import ( + clear_secondary_msg_info, + clear_status_msg_info, + clear_tool_msg_ids_for_topic, +) async def clear_topic_state( @@ -30,6 +34,7 @@ async def clear_topic_state( Cleans up: - _status_msg_info (status message tracking) - _tool_msg_ids (tool_use → message_id mapping) + - _secondary_msg_info (last deletable secondary message per topic) - _interactive_msgs and _interactive_mode (interactive UI state) - user_data pending state (_pending_thread_id, _pending_thread_text) """ @@ -39,6 +44,9 @@ async def clear_topic_state( # Clear tool message ID tracking clear_tool_msg_ids_for_topic(user_id, thread_id) + # Clear secondary-message tracking + clear_secondary_msg_info(user_id, thread_id) + # Clear interactive UI state (also deletes message from chat) await clear_interactive_msg(user_id, bot, thread_id) diff --git a/src/ccbot/handlers/message_queue.py b/src/ccbot/handlers/message_queue.py index 536c21fee..db7cd8521 100644 --- a/src/ccbot/handlers/message_queue.py +++ b/src/ccbot/handlers/message_queue.py @@ -15,6 +15,10 @@ - Message queue worker: Background task processing user's queue - Content task processing with tool_use/tool_result handling - Status message tracking and conversion (keyed by (user_id, thread_id)) + - Secondary message tracking: the last non-final (thinking/tool_use/ + tool_result/local_command/user-echo) message per topic is deleted as + soon as the next message arrives, so only the final assistant answer + is left behind once a turn completes. """ import asyncio @@ -64,6 +68,10 @@ class MessageTask: content_type: str = "text" thread_id: int | None = None # Telegram topic thread_id for targeted send image_data: list[tuple[str, bytes]] | None = None # From tool_result images + # True for everything except the final assistant text answer (thinking, + # tool_use, tool_result, local_command, user-message echo). Secondary + # messages are deleted as soon as the next message for the topic arrives. + is_secondary: bool = False # Per-user message queues and worker tasks @@ -78,6 +86,13 @@ class MessageTask: # Status message tracking: (user_id, thread_id_or_0) -> (message_id, window_id, last_text) _status_msg_info: dict[tuple[int, int], tuple[int, str, str]] = {} +# Last secondary (non-final) content message(s) per topic, deleted once the +# next message arrives: (user_id, thread_id_or_0) -> (message_ids, window_id). +# A list because merged tasks (several secondary entries queued together get +# folded into one MessageTask by _merge_content_tasks) or a paginated +# local_command can produce more than one Telegram message per task. +_secondary_msg_info: dict[tuple[int, int], tuple[list[int], str]] = {} + # Flood control: user_id -> monotonic time when ban expires _flood_until: dict[int, float] = {} @@ -130,7 +145,9 @@ def _can_merge_tasks(base: MessageTask, candidate: MessageTask) -> bool: return False if candidate.content_type in ("tool_use", "tool_result"): return False - return True + # Never fold a secondary (deletable) message together with the final + # answer — the merged task would inherit only one is_secondary value. + return base.is_secondary == candidate.is_secondary async def _merge_content_tasks( @@ -192,6 +209,7 @@ async def _merge_content_tasks( tool_use_id=first.tool_use_id, content_type=first.content_type, thread_id=first.thread_id, + is_secondary=first.is_secondary, ), merge_count, ) @@ -386,9 +404,38 @@ async def _process_content_task(bot: Bot, user_id: int, task: MessageTask) -> No logger.debug(f"Failed to edit tool msg {edit_msg_id}, sending new") # Fall through to send as new message - # 2. Send content messages, converting status message to first content part + # 2. A new message is arriving for this topic — delete whatever secondary + # (non-final) message was left over from the previous one. Applies + # whether this new task is itself secondary or the final answer: either + # way the old one is now superseded and safe to remove. + # + # Remaining (not-yet-deleted) ids are written back before re-raising + # RetryAfter, so _process_content_with_retry's retry of this whole + # task resumes deleting exactly where it left off instead of losing + # track of ids already popped — same reasoning as for a single id, + # just extended to a list (see the flood-control regression test). + skey = (user_id, tid) + old_secondary = _secondary_msg_info.get(skey) + if old_secondary: + old_msg_ids, old_wid = old_secondary + remaining_ids = list(old_msg_ids) + while remaining_ids: + old_msg_id = remaining_ids[0] + try: + await bot.delete_message(chat_id=chat_id, message_id=old_msg_id) + remaining_ids.pop(0) + except RetryAfter: + _secondary_msg_info[skey] = (remaining_ids, old_wid) + raise + except Exception as e: + logger.debug("Failed to delete secondary message %d: %s", old_msg_id, e) + remaining_ids.pop(0) + _secondary_msg_info.pop(skey, None) + + # 3. Send content messages, converting status message to first content part first_part = True last_msg_id: int | None = None + sent_msg_ids: list[int] = [] for part in task.parts: sent = None @@ -404,6 +451,7 @@ async def _process_content_task(bot: Bot, user_id: int, task: MessageTask) -> No ) if converted_msg_id is not None: last_msg_id = converted_msg_id + sent_msg_ids.append(converted_msg_id) continue sent = await send_with_fallback( @@ -415,15 +463,25 @@ async def _process_content_task(bot: Bot, user_id: int, task: MessageTask) -> No if sent: last_msg_id = sent.message_id + sent_msg_ids.append(sent.message_id) - # 3. Record tool_use message ID for later editing + # 4. Record tool_use message ID for later editing if last_msg_id and task.tool_use_id and task.content_type == "tool_use": _tool_msg_ids[(task.tool_use_id, user_id, tid)] = last_msg_id - # 4. Send images if present (from tool_result with base64 image blocks) + # 5. Track every message sent for this task as the topic's current + # secondary message(s), so all of them get deleted once the next + # message (of any kind) arrives — not just the last one. A merged + # task (several secondary entries queued together, folded into one + # MessageTask by _merge_content_tasks) or a paginated local_command + # can produce more than one Telegram message here. + if sent_msg_ids and task.is_secondary: + _secondary_msg_info[skey] = (sent_msg_ids, wid) + + # 6. Send images if present (from tool_result with base64 image blocks) await _send_task_images(bot, chat_id, task) - # 5. After content, check and send status + # 7. After content, check and send status await _check_and_send_status(bot, user_id, wid, task.thread_id) @@ -643,6 +701,7 @@ async def enqueue_content_message( text: str | None = None, thread_id: int | None = None, image_data: list[tuple[str, bytes]] | None = None, + is_secondary: bool = False, ) -> None: """Enqueue a content message task.""" logger.debug( @@ -662,6 +721,7 @@ async def enqueue_content_message( content_type=content_type, thread_id=thread_id, image_data=image_data, + is_secondary=is_secondary, ) queue.put_nowait(task) @@ -709,6 +769,12 @@ def clear_status_msg_info(user_id: int, thread_id: int | None = None) -> None: _status_msg_info.pop(skey, None) +def clear_secondary_msg_info(user_id: int, thread_id: int | None = None) -> None: + """Clear secondary-message tracking for a user (and optionally a specific thread).""" + skey = (user_id, thread_id or 0) + _secondary_msg_info.pop(skey, None) + + def clear_tool_msg_ids_for_topic(user_id: int, thread_id: int | None = None) -> None: """Clear tool message ID tracking for a specific topic. diff --git a/src/ccbot/markdown_v2.py b/src/ccbot/markdown_v2.py index d01bbdf86..ff3c7e7a8 100644 --- a/src/ccbot/markdown_v2.py +++ b/src/ccbot/markdown_v2.py @@ -124,6 +124,12 @@ def _escape_mdv2(text: str) -> str: def _render_expandable_quote(m: re.Match[str]) -> str: """Render an expandable blockquote block in raw MarkdownV2. + Telegram's expandable-quote grammar requires the first quoted line to + start with "**>" (not plain ">"); otherwise the trailing "||" is an + unmatched spoiler entity and the whole sendMessage call is rejected + with "Can't parse entities", silently degrading the message to + unformatted plain text via the caller's fallback path. + Truncates the rendered output to _EXPQUOTE_MAX_RENDERED chars to ensure the final message fits within Telegram's 4096 limit. """ @@ -148,9 +154,10 @@ def _render_expandable_quote(m: re.Match[str]) -> str: break built.append(f">{line}") total_len += line_cost - if truncated: - return "\n".join(built) + suffix - return "\n".join(built) + "||" + body = "\n".join(built) + (suffix if truncated else "||") + # Mark the block as expandable by promoting the first ">" to "**>". + idx = body.index(">") + return body[:idx] + "**" + body[idx:] def _markdownify(text: str) -> str: diff --git a/tests/ccbot/handlers/test_message_queue.py b/tests/ccbot/handlers/test_message_queue.py new file mode 100644 index 000000000..81a4a1b58 --- /dev/null +++ b/tests/ccbot/handlers/test_message_queue.py @@ -0,0 +1,283 @@ +"""Tests for message_queue — secondary-message rolling delete-on-next behavior. + +Covers: _can_merge_tasks refusing to fold secondary and final-answer tasks +together, _process_content_task deleting the previous secondary message(s) +when the next one arrives (including multi-part tasks from merging or +pagination), the tool_use/tool_result edit-in-place path being left +untouched by that deletion, flood-control resilience, and +clear_secondary_msg_info. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from telegram.error import RetryAfter + +from ccbot.handlers.message_queue import ( + MessageTask, + _can_merge_tasks, + _process_content_task, + _secondary_msg_info, + _tool_msg_ids, + clear_secondary_msg_info, +) + +USER_ID = 1 +THREAD_ID = 42 +CHAT_ID = 555 + + +@pytest.fixture +def mock_bot(): + bot = AsyncMock() + sent = MagicMock() + sent.message_id = 100 + bot.send_message.return_value = sent + return bot + + +@pytest.fixture(autouse=True) +def _clear_queue_state(): + _secondary_msg_info.clear() + _tool_msg_ids.clear() + yield + _secondary_msg_info.clear() + _tool_msg_ids.clear() + + +def _task( + content_type: str, + is_secondary: bool, + parts: list[str] | None = None, + tool_use_id: str | None = None, +) -> MessageTask: + return MessageTask( + task_type="content", + window_id="@1", + parts=parts or ["hello"], + content_type=content_type, + thread_id=THREAD_ID, + is_secondary=is_secondary, + tool_use_id=tool_use_id, + ) + + +class TestCanMergeSecondary: + def test_secondary_and_final_do_not_merge(self) -> None: + base = _task("thinking", is_secondary=True) + candidate = _task("text", is_secondary=False) + assert _can_merge_tasks(base, candidate) is False + + def test_same_secondary_merges(self) -> None: + base = _task("thinking", is_secondary=True) + candidate = _task("thinking", is_secondary=True) + assert _can_merge_tasks(base, candidate) is True + + +@pytest.mark.asyncio +class TestSecondaryMessageRollingDelete: + async def test_secondary_message_is_tracked(self, mock_bot: AsyncMock) -> None: + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([100], "@1") + mock_bot.delete_message.assert_not_called() + + async def test_next_secondary_deletes_previous(self, mock_bot: AsyncMock) -> None: + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + second_sent = MagicMock() + second_sent.message_id = 101 + mock_bot.send_message.return_value = second_sent + await _process_content_task(mock_bot, USER_ID, _task("tool_result", True)) + + mock_bot.delete_message.assert_called_once_with(chat_id=CHAT_ID, message_id=100) + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([101], "@1") + + async def test_final_answer_deletes_previous_but_is_not_tracked( + self, mock_bot: AsyncMock + ) -> None: + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + final_sent = MagicMock() + final_sent.message_id = 102 + mock_bot.send_message.return_value = final_sent + await _process_content_task( + mock_bot, USER_ID, _task("text", False, parts=["Fixed it"]) + ) + + mock_bot.delete_message.assert_called_once_with(chat_id=CHAT_ID, message_id=100) + assert (USER_ID, THREAD_ID) not in _secondary_msg_info + + async def test_tool_result_edit_in_place_does_not_delete( + self, mock_bot: AsyncMock + ) -> None: + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task( + mock_bot, USER_ID, _task("tool_use", True, tool_use_id="tu1") + ) + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([100], "@1") + + await _process_content_task( + mock_bot, USER_ID, _task("tool_result", True, tool_use_id="tu1") + ) + + mock_bot.delete_message.assert_not_called() + mock_bot.edit_message_text.assert_called_once() + # Same message_id throughout — edited in place, tracking stays valid. + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([100], "@1") + + async def test_multi_part_secondary_task_tracks_every_message( + self, mock_bot: AsyncMock + ) -> None: + """Regression test for merged secondary tasks losing all but the last id. + + _merge_content_tasks folds several secondary entries that arrived + close together (e.g. task-notification bursts) into one MessageTask + with several `parts`. Each part is sent as its own Telegram message; + all of them — not just the last — must be tracked so the next + message deletes every one of them instead of leaking all but one. + """ + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + sent_ids = iter([201, 202, 203]) + + def _send(*_a, **_kw): + m = MagicMock() + m.message_id = next(sent_ids) + return m + + mock_bot.send_message.side_effect = _send + + merged = _task("text", True, parts=["notif 1", "notif 2", "notif 3"]) + await _process_content_task(mock_bot, USER_ID, merged) + + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([201, 202, 203], "@1") + assert mock_bot.send_message.call_count == 3 + + # Next message must delete all three, not just the last. + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + next_sent = MagicMock() + next_sent.message_id = 204 + mock_bot.send_message.side_effect = None + mock_bot.send_message.return_value = next_sent + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + deleted_ids = {c.kwargs["message_id"] for c in mock_bot.delete_message.call_args_list} + assert deleted_ids == {201, 202, 203} + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([204], "@1") + + async def test_retry_after_on_delete_keeps_tracking_for_retry( + self, mock_bot: AsyncMock + ) -> None: + """A flood-controlled delete must not orphan the tracked message. + + Regression test: popping _secondary_msg_info before the delete call + succeeded meant a RetryAfter (bare `except Exception: pass`) silently + dropped the delete and forgot the message — it was never cleaned up + by a later retry. The entry must survive so a retried call can still + find and delete it. + """ + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([100], "@1") + + mock_bot.delete_message.side_effect = RetryAfter(5) + with pytest.raises(RetryAfter): + await _process_content_task(mock_bot, USER_ID, _task("tool_result", True)) + + # Still tracked (not silently dropped) so a retry can delete it. + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([100], "@1") + + async def test_retry_after_partway_through_a_multi_id_delete_keeps_remainder( + self, mock_bot: AsyncMock + ) -> None: + """RetryAfter partway through deleting several tracked ids must not + re-attempt (or lose track of) the ones already handled.""" + _secondary_msg_info[(USER_ID, THREAD_ID)] = ([301, 302, 303], "@1") + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + # First delete succeeds, second is flood-controlled. + mock_bot.delete_message.side_effect = [None, RetryAfter(5)] + with pytest.raises(RetryAfter): + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + assert mock_bot.delete_message.call_count == 2 + # 301 already deleted — must not be retried; 302/303 remain for the retry. + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([302, 303], "@1") + + async def test_non_retryable_delete_failure_still_clears_tracking( + self, mock_bot: AsyncMock + ) -> None: + with ( + patch("ccbot.handlers.message_queue.session_manager") as mock_sm, + patch("ccbot.handlers.message_queue.tmux_manager") as mock_tmux, + ): + mock_sm.resolve_chat_id.return_value = CHAT_ID + mock_tmux.find_window_by_id = AsyncMock(return_value=None) + + await _process_content_task(mock_bot, USER_ID, _task("thinking", True)) + + mock_bot.delete_message.side_effect = Exception("message to delete not found") + second_sent = MagicMock() + second_sent.message_id = 101 + mock_bot.send_message.return_value = second_sent + await _process_content_task(mock_bot, USER_ID, _task("tool_result", True)) + + # Gives up cleanly rather than retrying forever, and still tracks the new message. + assert _secondary_msg_info[(USER_ID, THREAD_ID)] == ([101], "@1") + + +class TestClearSecondaryMsgInfo: + def test_clears_tracked_entry(self) -> None: + _secondary_msg_info[(USER_ID, THREAD_ID)] = ([100], "@1") + clear_secondary_msg_info(USER_ID, THREAD_ID) + assert (USER_ID, THREAD_ID) not in _secondary_msg_info diff --git a/tests/ccbot/test_markdown_v2.py b/tests/ccbot/test_markdown_v2.py index 3628b185f..ceb77a3a5 100644 --- a/tests/ccbot/test_markdown_v2.py +++ b/tests/ccbot/test_markdown_v2.py @@ -46,13 +46,34 @@ def test_expandable_quote_sentinels(self) -> None: result = convert_markdown(text) assert EXP_START not in result assert EXP_END not in result - assert ">quoted content||" in result + # Telegram requires "**>" (not plain ">") to open an *expandable* + # blockquote — otherwise the trailing "||" is an unmatched spoiler + # entity and Telegram rejects the whole message. + assert "**>quoted content||" in result def test_mixed_text_and_expandable_quote(self) -> None: text = f"before {EXP_START}inside quote{EXP_END} after" result = convert_markdown(text) assert EXP_START not in result assert EXP_END not in result - assert ">inside quote||" in result + assert "**>inside quote||" in result assert "before" in result assert "after" in result + + def test_multiline_expandable_quote_only_first_line_marked(self) -> None: + text = f"{EXP_START}line one\nline two\nline three{EXP_END}" + result = convert_markdown(text) + assert "**>line one" in result + assert "\n>line two" in result + assert "\n>line three||" in result + # Only the first line carries the "**" expandability marker. + assert result.count("**>") == 1 + + def test_expandable_quote_truncation_keeps_expandable_marker(self) -> None: + text = f"{EXP_START}{'x' * 10000}{EXP_END}" + result = convert_markdown(text) + assert "**>" in result + assert result.index("**>") < 10 + assert "\\(truncated\\)||" in result + # Rendered quote block must still fit Telegram's message limit. + assert len(result) < 4096