Group Telegram tool streams by response
Anmol Malik committed
May 28, 2026 at 23:00 UTC
a17cef5b03f500dee1e9e9c3fa0bcdb0c03f3eee
2 files changed
+101
-2
plugins/_telegram_integration/helpers/draft_stream.py
+19
-2
@@ -29,10 +29,10 @@ TOOL_EMOJIS: dict[str, str] = {
29
"code": "⌨️",
30
"code_execution_tool": "⌨️",
31
"duckduckgo_search": "🔎",
32
+ "read_file": "📖",
33
"file": "📄",
34
"knowledge_tool": "📚",
35
"memory": "🧠",
35
- "read_file": "📖",
36
"search": "🔎",
37
"search_engine": "🔎",
38
"search_files": "🔎",
@@ -89,6 +89,12 @@ async def send_intermediate_response(
89
response_text: str,
90
keyboard: list[list[dict]] | None = None,
91
) -> bool:
92
+ if context.data.get(CTX_TG_RESPONSE_MESSAGE_ID):
93
+ sent = await finalize_response(context, response_text, keyboard)
94
+ if sent:
95
+ _reset_progress_group(context)
96
+ return sent
97
+
98
html = _format_response(response_text)
99
if not html:
100
return False
@@ -104,7 +110,10 @@ async def send_intermediate_response(
110
parse_mode="HTML",
111
reply_markup=_keyboard_markup(keyboard),
112
)
107
- return bool(sent_id)
113
+ sent = bool(sent_id)
114
+ if sent:
115
+ _reset_progress_group(context)
116
+ return sent
117
except Exception as e:
118
PrintStyle.debug(f"Telegram intermediate response failed: {e}")
119
return False
@@ -142,6 +151,11 @@ def clear(context: AgentContext) -> None:
151
context.data.pop(key, None)
152
153
154
+def _reset_progress_group(context: AgentContext) -> None:
155
+ context.data.pop(CTX_TG_PROGRESS_LINES, None)
156
+ context.data.pop(CTX_TG_PROGRESS_MESSAGE_ID, None)
157
+
158
+
159
def _stream_enabled(context: AgentContext) -> bool:
160
value = context.get_data(CTX_TG_STREAM_ENABLED)
161
return True if value is None else bool(value)
@@ -208,6 +222,7 @@ async def _update_response_message(
222
keyboard: list[list[dict]] | None = None,
223
force: bool = False,
224
) -> bool:
225
+ had_message = bool(context.data.get(CTX_TG_RESPONSE_MESSAGE_ID))
226
message_id = await _ensure_response_message(context, text)
227
bot = _bot_instance(context)
228
chat_id = context.data.get(CTX_TG_CHAT_ID)
@@ -215,6 +230,8 @@ async def _update_response_message(
230
return False
231
markup = _keyboard_markup(keyboard)
232
html = _format_response(text)
233
+ if not had_message and not markup and not force:
234
+ return True
235
try:
236
ok = await tc.raw_edit_text(
237
bot.bot.token,
tests/test_telegram_intermediate_response.py
+82
@@ -4,6 +4,8 @@ from plugins._telegram_integration.helpers import draft_stream
4
from plugins._telegram_integration.helpers.constants import (
5
CTX_TG_BOT,
6
CTX_TG_CHAT_ID,
7
+ CTX_TG_PROGRESS_LINES,
8
+ CTX_TG_PROGRESS_MESSAGE_ID,
9
CTX_TG_REPLY_TO,
10
CTX_TG_RESPONSE_MESSAGE_ID,
11
)
@@ -68,3 +70,83 @@ def test_intermediate_response_sends_separate_non_reply_message(monkeypatch):
70
}
71
]
72
assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data
73
+
74
+
75
+def test_intermediate_response_finalizes_active_stream_and_starts_next_tool_group(monkeypatch):
76
+ calls = []
77
+ next_message_id = 100
78
+
79
+ async def fake_send(token, chat_id, text, reply_to_message_id=None, parse_mode="HTML", reply_markup=None):
80
+ nonlocal next_message_id
81
+ calls.append(
82
+ {
83
+ "method": "send",
84
+ "text": text,
85
+ "reply_to_message_id": reply_to_message_id,
86
+ "parse_mode": parse_mode,
87
+ "reply_markup": reply_markup,
88
+ "message_id": next_message_id,
89
+ }
90
+ )
91
+ next_message_id += 1
92
+ return next_message_id - 1
93
+
94
+ async def fake_edit(token, chat_id, message_id, text, parse_mode="HTML", reply_markup=None):
95
+ calls.append(
96
+ {
97
+ "method": "edit",
98
+ "message_id": message_id,
99
+ "text": text,
100
+ "parse_mode": parse_mode,
101
+ "reply_markup": reply_markup,
102
+ }
103
+ )
104
+ return True
105
+
106
+ context = FakeContext()
107
+ monkeypatch.setattr(draft_stream, "_bot_instance", lambda ctx: FakeBot())
108
+ monkeypatch.setattr(draft_stream.tc, "raw_send_text", fake_send)
109
+ monkeypatch.setattr(draft_stream.tc, "raw_edit_text", fake_edit)
110
+
111
+ asyncio.run(draft_stream.add_tool_start(context, "search_engine", {"query": "telegram bot api"}))
112
+ asyncio.run(draft_stream.update_response(context, "Found the docs."))
113
+ sent = asyncio.run(draft_stream.send_intermediate_response(context, "Found the docs."))
114
+ asyncio.run(draft_stream.add_tool_start(context, "read_file", {"path": "notes.md"}))
115
+
116
+ assert sent is True
117
+ assert calls == [
118
+ {
119
+ "method": "send",
120
+ "text": "🔎 search engine: telegram bot api",
121
+ "reply_to_message_id": None,
122
+ "parse_mode": None,
123
+ "reply_markup": None,
124
+ "message_id": 100,
125
+ },
126
+ {
127
+ "method": "send",
128
+ "text": "Found the docs.",
129
+ "reply_to_message_id": 456,
130
+ "parse_mode": "HTML",
131
+ "reply_markup": None,
132
+ "message_id": 101,
133
+ },
134
+ {
135
+ "method": "edit",
136
+ "message_id": 101,
137
+ "text": "Found the docs.",
138
+ "parse_mode": "HTML",
139
+ "reply_markup": None,
140
+ },
141
+ {
142
+ "method": "send",
143
+ "text": "📖 read file: notes.md",
144
+ "reply_to_message_id": None,
145
+ "parse_mode": None,
146
+ "reply_markup": None,
147
+ "message_id": 102,
148
+ },
149
+ ]
150
+ assert context.data[CTX_TG_PROGRESS_MESSAGE_ID] == 102
151
+ assert context.data[CTX_TG_PROGRESS_LINES] == ["📖 read file: notes.md"]
152
+ assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data