Add Telegram long-run status updates
Anmol Malik committed
May 28, 2026 at 23:14 UTC
99207d60a2ce3a086f124ad9d27d096b09971657
6 files changed
+182
-6
plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py
+2
-1
@@ -27,7 +27,7 @@ class TelegramFriendlyError(Extension):
27
28
context.data[CTX_TG_ERROR_SENT] = True
29
30
- from plugins._telegram_integration.helpers import draft_stream, error_ui
30
+ from plugins._telegram_integration.helpers import draft_stream, error_ui, heartbeat
31
from plugins._telegram_integration.helpers.handler import send_telegram_reply
32
33
text = error_ui.friendly_error_message(exception)
@@ -38,5 +38,6 @@ class TelegramFriendlyError(Extension):
38
typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None)
39
if typing_stop:
40
typing_stop.set()
41
+ await heartbeat.stop(context)
42
draft_stream.clear(context)
43
context.data.pop(CTX_TG_REPLY_TO, None)
plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py
+2
-1
@@ -12,6 +12,7 @@ class TelegramDraftStart(Extension):
12
if not context.data.get(CTX_TG_BOT):
13
return
14
15
- from plugins._telegram_integration.helpers import draft_stream
15
+ from plugins._telegram_integration.helpers import draft_stream, heartbeat
16
17
await draft_stream.start(context)
18
+ await heartbeat.start(context)
plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py
+4
-4
@@ -26,14 +26,13 @@ class TelegramAutoReply(Extension):
26
return
27
28
response_text = _extract_last_response(context)
29
- if not response_text:
30
- return
29
30
attachments = context.data.pop(CTX_TG_ATTACHMENTS, [])
31
keyboard = context.data.pop(CTX_TG_KEYBOARD, None)
32
33
try:
36
- await self._send_reply(context, response_text, attachments, keyboard)
34
+ if response_text:
35
+ await self._send_reply(context, response_text, attachments, keyboard)
36
except Exception as e:
37
PrintStyle.error(f"Telegram auto-reply error: {format_error(e)}")
38
finally:
@@ -41,8 +40,9 @@ class TelegramAutoReply(Extension):
40
typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None)
41
if typing_stop:
42
typing_stop.set()
44
- from plugins._telegram_integration.helpers import draft_stream
43
+ from plugins._telegram_integration.helpers import draft_stream, heartbeat
44
45
+ await heartbeat.stop(context)
46
draft_stream.clear(context)
47
context.data.pop(CTX_TG_REPLY_TO, None)
48
plugins/_telegram_integration/helpers/constants.py
+2
@@ -23,3 +23,5 @@ CTX_TG_RESPONSE_MESSAGE_ID = "_telegram_response_message_id"
23
CTX_TG_RESPONSE_TEXT = "_telegram_response_text"
24
CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update"
25
CTX_TG_ERROR_SENT = "_telegram_error_sent"
26
+CTX_TG_HEARTBEAT_TASK = "_telegram_heartbeat_task"
27
+CTX_TG_HEARTBEAT_STOP = "_telegram_heartbeat_stop"
plugins/_telegram_integration/helpers/heartbeat.py
new
+101
@@ -0,0 +1,101 @@
1
+from __future__ import annotations
2
+
3
+import asyncio
4
+import re
5
+import time
6
+from contextlib import suppress
7
+
8
+from agent import AgentContext
9
+from helpers.print_style import PrintStyle
10
+from plugins._telegram_integration.helpers import telegram_client as tc
11
+from plugins._telegram_integration.helpers.bot_manager import get_bot
12
+from plugins._telegram_integration.helpers.constants import (
13
+ CTX_TG_BOT,
14
+ CTX_TG_CHAT_ID,
15
+ CTX_TG_HEARTBEAT_STOP,
16
+ CTX_TG_HEARTBEAT_TASK,
17
+)
18
+
19
+
20
+HEARTBEAT_INTERVAL_SECONDS = 180
21
+
22
+
23
+async def start(context: AgentContext) -> None:
24
+ task = context.data.get(CTX_TG_HEARTBEAT_TASK)
25
+ if task and not task.done():
26
+ return
27
+
28
+ stop_event = asyncio.Event()
29
+ context.data[CTX_TG_HEARTBEAT_STOP] = stop_event
30
+ context.data[CTX_TG_HEARTBEAT_TASK] = asyncio.create_task(
31
+ _heartbeat_loop(context, stop_event, time.monotonic())
32
+ )
33
+
34
+
35
+async def stop(context: AgentContext) -> None:
36
+ stop_event = context.data.pop(CTX_TG_HEARTBEAT_STOP, None)
37
+ if stop_event:
38
+ stop_event.set()
39
+
40
+ task = context.data.pop(CTX_TG_HEARTBEAT_TASK, None)
41
+ if task and not task.done():
42
+ task.cancel()
43
+ with suppress(asyncio.CancelledError):
44
+ await task
45
+
46
+
47
+async def _heartbeat_loop(
48
+ context: AgentContext,
49
+ stop_event: asyncio.Event,
50
+ started_at: float,
51
+) -> None:
52
+ try:
53
+ while True:
54
+ try:
55
+ await asyncio.wait_for(stop_event.wait(), timeout=HEARTBEAT_INTERVAL_SECONDS)
56
+ return
57
+ except asyncio.TimeoutError:
58
+ await _send_heartbeat(context, time.monotonic() - started_at)
59
+ except asyncio.CancelledError:
60
+ raise
61
+ except Exception as exc:
62
+ PrintStyle.debug(f"Telegram heartbeat stopped: {exc}")
63
+
64
+
65
+async def _send_heartbeat(context: AgentContext, elapsed_seconds: float) -> bool:
66
+ bot_name = context.data.get(CTX_TG_BOT)
67
+ chat_id = context.data.get(CTX_TG_CHAT_ID)
68
+ bot = get_bot(bot_name) if bot_name else None
69
+ if not bot or not chat_id:
70
+ return False
71
+
72
+ text = heartbeat_text(elapsed_seconds, _current_reason(context))
73
+ sent_id = await tc.raw_send_text(
74
+ bot.bot.token,
75
+ int(chat_id),
76
+ text,
77
+ parse_mode=None,
78
+ )
79
+ return bool(sent_id)
80
+
81
+
82
+def heartbeat_text(elapsed_seconds: float, reason: str = "") -> str:
83
+ minutes = max(1, int(round(elapsed_seconds / 60)))
84
+ detail = reason or "working on your request"
85
+ return f"Still working... ({minutes} min elapsed - {detail})"
86
+
87
+
88
+def _current_reason(context: AgentContext) -> str:
89
+ progress = str(getattr(getattr(context, "log", None), "progress", "") or "")
90
+ progress = _clean_progress(progress)
91
+ if not progress or progress.lower() == "waiting for input":
92
+ return "working on your request"
93
+ return progress
94
+
95
+
96
+def _clean_progress(text: str) -> str:
97
+ value = re.sub(r"icon://[a-zA-Z0-9_]+(?:\[[^\]]*\])?\s*", "", text)
98
+ value = re.sub(r"\s+", " ", value).strip()
99
+ if len(value) > 80:
100
+ value = f"{value[:77].rstrip()}..."
101
+ return value
tests/test_telegram_heartbeat.py
new
+71
@@ -0,0 +1,71 @@
1
+import asyncio
2
+
3
+from plugins._telegram_integration.helpers import heartbeat
4
+from plugins._telegram_integration.helpers.constants import (
5
+ CTX_TG_BOT,
6
+ CTX_TG_CHAT_ID,
7
+ CTX_TG_HEARTBEAT_STOP,
8
+ CTX_TG_HEARTBEAT_TASK,
9
+)
10
+
11
+
12
+class FakeBot:
13
+ class Bot:
14
+ token = "token"
15
+
16
+ bot = Bot()
17
+
18
+
19
+class FakeLog:
20
+ progress = "icon://search[Search] A0: Searching official docs"
21
+
22
+
23
+class FakeContext:
24
+ def __init__(self):
25
+ self.data = {
26
+ CTX_TG_BOT: "main",
27
+ CTX_TG_CHAT_ID: 123,
28
+ }
29
+ self.log = FakeLog()
30
+
31
+
32
+def test_heartbeat_text_omits_iteration_count():
33
+ text = heartbeat.heartbeat_text(180, "waiting for provider response")
34
+
35
+ assert text == "Still working... (3 min elapsed - waiting for provider response)"
36
+ assert "iteration" not in text
37
+
38
+
39
+def test_heartbeat_sends_periodic_status_and_stops(monkeypatch):
40
+ calls = []
41
+
42
+ async def fake_send(token, chat_id, text, parse_mode=None, **kwargs):
43
+ calls.append(
44
+ {
45
+ "token": token,
46
+ "chat_id": chat_id,
47
+ "text": text,
48
+ "parse_mode": parse_mode,
49
+ }
50
+ )
51
+ return len(calls)
52
+
53
+ context = FakeContext()
54
+ monkeypatch.setattr(heartbeat, "HEARTBEAT_INTERVAL_SECONDS", 0.01)
55
+ monkeypatch.setattr(heartbeat, "get_bot", lambda name: FakeBot())
56
+ monkeypatch.setattr(heartbeat.tc, "raw_send_text", fake_send)
57
+
58
+ async def run():
59
+ await heartbeat.start(context)
60
+ await asyncio.sleep(0.025)
61
+ await heartbeat.stop(context)
62
+
63
+ asyncio.run(run())
64
+
65
+ assert calls
66
+ assert calls[0]["chat_id"] == 123
67
+ assert calls[0]["parse_mode"] is None
68
+ assert "Still working..." in calls[0]["text"]
69
+ assert "A0: Searching official docs" in calls[0]["text"]
70
+ assert CTX_TG_HEARTBEAT_TASK not in context.data
71
+ assert CTX_TG_HEARTBEAT_STOP not in context.data