add typing simulation

assistant_bot/agent.py

@@ -6,7 +6,6 @@ from typing import Any
 from zoneinfo import ZoneInfo
 
 from telegram import Update
-from telegram.constants import ChatAction
 from telegram.ext import ContextTypes
 
 from .ai import AIClientError
@@ -16,16 +15,14 @@ from .prompts import ASSISTANT_SYSTEM_PROMPT
 from .reminders import parse_reminder
 from .storage import AssistantStorage
 from .telegram_utils import (
-    edit_markdown,
-    edit_or_reply,
     escape_markdown_text,
     get_ai_client,
     get_storage,
     get_tz,
-    markdown_code,
     parse_positive_int,
     reply_markdown,
     require_user_id,
+    typing_action,
 )
 from .time_utils import format_local_dt, to_utc_iso, utc_now
 
@@ -435,19 +432,12 @@ async def run_agent_prompt(
         {"role": "user", "content": prompt},
     ]
 
-    await context.bot.send_chat_action(
-        chat_id=chat_id,
-        action=ChatAction.TYPING,
-    )
-    placeholder = await reply_markdown(
-        update.effective_message,
-        f"💭 *Думаю через {escape_markdown_text(ai_client.display_name)} / "
-        f"{markdown_code(model)}…*",
-    )
-
     try:
         last_tool_results: list[dict[str, Any]] = []
         reset_context = False
+        final_answer: str | None = None
+
+        async with typing_action(context.bot, chat_id):
             for _step in range(MAX_AGENT_STEPS):
                 raw_answer = await ai_client.chat(model, messages, json_mode=True)
                 decision = parse_agent_decision(raw_answer)
@@ -476,7 +466,9 @@ async def run_agent_prompt(
                     messages.append(
                         {
                             "role": "assistant",
-                        "content": json_dumps({"tool_calls": decision.tool_calls}),
+                            "content": json_dumps(
+                                {"tool_calls": decision.tool_calls}
+                            ),
                         }
                     )
                     messages.append(
@@ -493,19 +485,12 @@ async def run_agent_prompt(
                     continue
 
                 if decision.final:
-                await edit_or_reply(placeholder, decision.final)
+                    final_answer = decision.final
-                save_conversation_exchange(
+                    break
-                    storage,
-                    user_id,
-                    chat_id,
-                    prompt,
-                    decision.final,
-                    reset_context,
-                )
-                return
 
                 messages.append({"role": "assistant", "content": raw_answer})
 
+            if final_answer is None:
                 final_answer = await ai_client.chat(
                     model,
                     [
@@ -520,7 +505,8 @@ async def run_agent_prompt(
                         },
                     ],
                 )
-        await edit_or_reply(placeholder, final_answer)
+
+        await reply_markdown(update.effective_message, final_answer)
         save_conversation_exchange(
             storage,
             user_id,
@@ -536,8 +522,8 @@ async def run_agent_prompt(
             hint = (
                 "Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY и доступ к выбранной модели."
             )
-        await edit_markdown(
+        await reply_markdown(
-            placeholder,
+            update.effective_message,
             f"⚠️ **Не получилось вызвать "
             f"{escape_markdown_text(ai_client.display_name)}.**\n\n"
             f"{escape_markdown_text(exc)}\n\n"
@@ -546,8 +532,8 @@ async def run_agent_prompt(
         return
     except Exception:
         logger.exception("Failed to run AI prompt")
-        await edit_markdown(
+        await reply_markdown(
-            placeholder,
+            update.effective_message,
             "⚠️ **Произошла внутренняя ошибка** при запросе к модели.",
         )
 

assistant_bot/telegram_utils.py

@@ -1,10 +1,13 @@
+import asyncio
 import logging
 import re
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
 from uuid import uuid4
 from zoneinfo import ZoneInfo
 
 from telegram import Bot, InlineQueryResultArticle, InputTextMessageContent, Message, Update
-from telegram.constants import ParseMode
+from telegram.constants import ChatAction, ParseMode
 from telegram.error import BadRequest
 from telegram.ext import ContextTypes
 from telegramify_markdown import (
@@ -24,6 +27,7 @@ from .storage import AssistantStorage
 logger = logging.getLogger(__name__)
 
 FORMATTING_FALLBACK_NOTICE = "⚠️ Не удалось применить форматирование."
+TYPING_REFRESH_INTERVAL_SECONDS = 4.0
 _MARKDOWN_SPECIAL_CHARS = re.compile(r"([\\`*_[\]{}()#+\-.!|>~])")
 _RENDER_CONFIG = RenderConfig()
 _RENDER_CONFIG.markdown_symbol.heading_level_1 = ""
@@ -32,6 +36,51 @@ _RENDER_CONFIG.markdown_symbol.heading_level_3 = ""
 _RENDER_CONFIG.markdown_symbol.heading_level_4 = ""
 
 
+async def _refresh_typing(
+    bot: Bot,
+    chat_id: int,
+    interval: float,
+) -> None:
+    while True:
+        await asyncio.sleep(interval)
+        try:
+            await bot.send_chat_action(
+                chat_id=chat_id,
+                action=ChatAction.TYPING,
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.warning("Failed to refresh Telegram typing action", exc_info=True)
+
+
+@asynccontextmanager
+async def typing_action(
+    bot: Bot,
+    chat_id: int,
+    interval: float = TYPING_REFRESH_INTERVAL_SECONDS,
+) -> AsyncIterator[None]:
+    """Keep Telegram's typing indicator active while work is in progress."""
+    try:
+        await bot.send_chat_action(
+            chat_id=chat_id,
+            action=ChatAction.TYPING,
+        )
+    except Exception:
+        logger.warning("Failed to send Telegram typing action", exc_info=True)
+
+    refresh_task = asyncio.create_task(
+        _refresh_typing(bot, chat_id, interval),
+        name=f"telegram-typing-{chat_id}",
+    )
+    try:
+        yield
+    finally:
+        refresh_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await refresh_task
+
+
 def escape_markdown_text(value: object) -> str:
     """Escape dynamic text before embedding it into ordinary Markdown."""
     return _MARKDOWN_SPECIAL_CHARS.sub(r"\\\1", str(value))

tests/test_telegram_utils.py

@@ -1,8 +1,9 @@
+import asyncio
 import unittest
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
 
-from telegram.constants import ParseMode
+from telegram.constants import ChatAction, ParseMode
 from telegram.error import BadRequest
 
 from assistant_bot.telegram_utils import (
@@ -11,6 +12,7 @@ from assistant_bot.telegram_utils import (
     escape_markdown_text,
     render_markdown_chunks,
     reply_markdown,
+    typing_action,
 )
 
 
@@ -112,5 +114,33 @@ class MarkdownDeliveryTests(unittest.IsolatedAsyncioTestCase):
         )
 
 
+class TypingActionTests(unittest.IsolatedAsyncioTestCase):
+    async def test_refreshes_typing_until_context_exits(self) -> None:
+        refreshed = asyncio.Event()
+        call_count = 0
+
+        async def send_chat_action(**_kwargs) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count >= 2:
+                refreshed.set()
+
+        bot = SimpleNamespace(
+            send_chat_action=AsyncMock(side_effect=send_chat_action)
+        )
+
+        async with typing_action(bot, chat_id=42, interval=0.001):
+            await asyncio.wait_for(refreshed.wait(), timeout=0.5)
+
+        calls_after_exit = call_count
+        await asyncio.sleep(0.01)
+
+        self.assertGreaterEqual(calls_after_exit, 2)
+        self.assertEqual(call_count, calls_after_exit)
+        for call in bot.send_chat_action.await_args_list:
+            self.assertEqual(call.kwargs["chat_id"], 42)
+            self.assertEqual(call.kwargs["action"], ChatAction.TYPING)
+
+
 if __name__ == "__main__":
     unittest.main()