Group Telegram tool streams by response

plugins/_telegram_integration/helpers/draft_stream.py

@@ -29,10 +29,10 @@ TOOL_EMOJIS: dict[str, str] = {
     "code": "⌨️",
     "code_execution_tool": "⌨️",
     "duckduckgo_search": "🔎",
+    "read_file": "📖",
     "file": "📄",
     "knowledge_tool": "📚",
     "memory": "🧠",
-    "read_file": "📖",
     "search": "🔎",
     "search_engine": "🔎",
     "search_files": "🔎",
@@ -89,6 +89,12 @@ async def send_intermediate_response(
     response_text: str,
     keyboard: list[list[dict]] | None = None,
 ) -> bool:
+    if context.data.get(CTX_TG_RESPONSE_MESSAGE_ID):
+        sent = await finalize_response(context, response_text, keyboard)
+        if sent:
+            _reset_progress_group(context)
+        return sent
+
     html = _format_response(response_text)
     if not html:
         return False
@@ -104,7 +110,10 @@ async def send_intermediate_response(
             parse_mode="HTML",
             reply_markup=_keyboard_markup(keyboard),
         )
-        return bool(sent_id)
+        sent = bool(sent_id)
+        if sent:
+            _reset_progress_group(context)
+        return sent
     except Exception as e:
         PrintStyle.debug(f"Telegram intermediate response failed: {e}")
         return False
@@ -142,6 +151,11 @@ def clear(context: AgentContext) -> None:
         context.data.pop(key, None)
 
 
+def _reset_progress_group(context: AgentContext) -> None:
+    context.data.pop(CTX_TG_PROGRESS_LINES, None)
+    context.data.pop(CTX_TG_PROGRESS_MESSAGE_ID, None)
+
+
 def _stream_enabled(context: AgentContext) -> bool:
     value = context.get_data(CTX_TG_STREAM_ENABLED)
     return True if value is None else bool(value)
@@ -208,6 +222,7 @@ async def _update_response_message(
     keyboard: list[list[dict]] | None = None,
     force: bool = False,
 ) -> bool:
+    had_message = bool(context.data.get(CTX_TG_RESPONSE_MESSAGE_ID))
     message_id = await _ensure_response_message(context, text)
     bot = _bot_instance(context)
     chat_id = context.data.get(CTX_TG_CHAT_ID)
@@ -215,6 +230,8 @@ async def _update_response_message(
         return False
     markup = _keyboard_markup(keyboard)
     html = _format_response(text)
+    if not had_message and not markup and not force:
+        return True
     try:
         ok = await tc.raw_edit_text(
             bot.bot.token,

tests/test_telegram_intermediate_response.py

@@ -4,6 +4,8 @@ from plugins._telegram_integration.helpers import draft_stream
 from plugins._telegram_integration.helpers.constants import (
     CTX_TG_BOT,
     CTX_TG_CHAT_ID,
+    CTX_TG_PROGRESS_LINES,
+    CTX_TG_PROGRESS_MESSAGE_ID,
     CTX_TG_REPLY_TO,
     CTX_TG_RESPONSE_MESSAGE_ID,
 )
@@ -68,3 +70,83 @@ def test_intermediate_response_sends_separate_non_reply_message(monkeypatch):
         }
     ]
     assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data
+
+
+def test_intermediate_response_finalizes_active_stream_and_starts_next_tool_group(monkeypatch):
+    calls = []
+    next_message_id = 100
+
+    async def fake_send(token, chat_id, text, reply_to_message_id=None, parse_mode="HTML", reply_markup=None):
+        nonlocal next_message_id
+        calls.append(
+            {
+                "method": "send",
+                "text": text,
+                "reply_to_message_id": reply_to_message_id,
+                "parse_mode": parse_mode,
+                "reply_markup": reply_markup,
+                "message_id": next_message_id,
+            }
+        )
+        next_message_id += 1
+        return next_message_id - 1
+
+    async def fake_edit(token, chat_id, message_id, text, parse_mode="HTML", reply_markup=None):
+        calls.append(
+            {
+                "method": "edit",
+                "message_id": message_id,
+                "text": text,
+                "parse_mode": parse_mode,
+                "reply_markup": reply_markup,
+            }
+        )
+        return True
+
+    context = FakeContext()
+    monkeypatch.setattr(draft_stream, "_bot_instance", lambda ctx: FakeBot())
+    monkeypatch.setattr(draft_stream.tc, "raw_send_text", fake_send)
+    monkeypatch.setattr(draft_stream.tc, "raw_edit_text", fake_edit)
+
+    asyncio.run(draft_stream.add_tool_start(context, "search_engine", {"query": "telegram bot api"}))
+    asyncio.run(draft_stream.update_response(context, "Found the docs."))
+    sent = asyncio.run(draft_stream.send_intermediate_response(context, "Found the docs."))
+    asyncio.run(draft_stream.add_tool_start(context, "read_file", {"path": "notes.md"}))
+
+    assert sent is True
+    assert calls == [
+        {
+            "method": "send",
+            "text": "🔎 search engine: telegram bot api",
+            "reply_to_message_id": None,
+            "parse_mode": None,
+            "reply_markup": None,
+            "message_id": 100,
+        },
+        {
+            "method": "send",
+            "text": "Found the docs.",
+            "reply_to_message_id": 456,
+            "parse_mode": "HTML",
+            "reply_markup": None,
+            "message_id": 101,
+        },
+        {
+            "method": "edit",
+            "message_id": 101,
+            "text": "Found the docs.",
+            "parse_mode": "HTML",
+            "reply_markup": None,
+        },
+        {
+            "method": "send",
+            "text": "📖 read file: notes.md",
+            "reply_to_message_id": None,
+            "parse_mode": None,
+            "reply_markup": None,
+            "message_id": 102,
+        },
+    ]
+    assert context.data[CTX_TG_PROGRESS_MESSAGE_ID] == 102
+    assert context.data[CTX_TG_PROGRESS_LINES] == ["📖 read file: notes.md"]
+    assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data