| 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 |