142 lines
4.6 KiB
Python
142 lines
4.6 KiB
Python
"""One ``sendMessageDraft`` stream per running turn (§3.8).
|
|
|
|
The client folds a draft into the message that follows only when their
|
|
texts are identical - so the last push is the final text itself, rendered
|
|
exactly as the outbox will send it, with no status line.
|
|
|
|
A draft is ephemeral and lives 30 s, Telegram throttles edits to about one
|
|
per second per chat, and thinking or a tool call would otherwise look like a
|
|
hang - so the draft opens with a status line straight away, is refreshed on
|
|
a timer rather than on every delta, and is kept alive while nothing changes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import time
|
|
import zlib
|
|
from typing import TYPE_CHECKING
|
|
|
|
from aiogram.exceptions import TelegramAPIError, TelegramBadRequest
|
|
|
|
from beaver_gateway.frontends.telegram.render import to_html, to_html_tail
|
|
|
|
if TYPE_CHECKING:
|
|
from aiogram import Bot
|
|
|
|
__all__ = ["Draft"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.frontends.telegram.drafts")
|
|
|
|
_TAIL = 3500
|
|
_KEEPALIVE = 20.0
|
|
|
|
|
|
class Draft:
|
|
def __init__(
|
|
self,
|
|
bot: Bot,
|
|
*,
|
|
chat_id: int,
|
|
thread_id: int | None,
|
|
turn_id: str,
|
|
interval: float = 0.7,
|
|
status: str = "⏳ думаю",
|
|
) -> None:
|
|
self._bot = bot
|
|
self._chat_id = chat_id
|
|
self._thread_id = thread_id
|
|
self._draft_id = (zlib.crc32(turn_id.encode()) & 0x7FFFFFFF) or 1
|
|
self._interval = interval
|
|
self.status = status
|
|
self.text = ""
|
|
self._dirty = True
|
|
self._broken = False
|
|
self._last_sent = 0.0
|
|
self._task: asyncio.Task[None] | None = None
|
|
self._inflight: asyncio.Future[None] | None = None
|
|
|
|
def start(self) -> None:
|
|
if self._task is None:
|
|
self._task = asyncio.create_task(self._run())
|
|
|
|
def set_status(self, status: str) -> None:
|
|
if status != self.status:
|
|
self.status = status
|
|
self._dirty = True
|
|
|
|
def append(self, text: str) -> None:
|
|
if text:
|
|
self.text += text
|
|
self._dirty = True
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop pushing and wait for the push in flight.
|
|
|
|
The final ``sendMessage`` must reach Telegram after the last draft,
|
|
or the draft lands on top of the message and lingers for its 30 s.
|
|
"""
|
|
if self._task is not None:
|
|
self._task.cancel()
|
|
await asyncio.gather(self._task, return_exceptions=True)
|
|
self._task = None
|
|
if self._inflight is not None:
|
|
with contextlib.suppress(Exception):
|
|
await self._inflight
|
|
self._inflight = None
|
|
|
|
async def finish(self, text: str) -> None:
|
|
"""Last push: the final text verbatim, so the message replaces the draft."""
|
|
await self.stop()
|
|
if self._broken or not text.strip():
|
|
return
|
|
try:
|
|
try:
|
|
await self._send(to_html(text), "HTML")
|
|
except TelegramBadRequest:
|
|
await self._send(text, None)
|
|
except TelegramAPIError as exc:
|
|
_log.warning(
|
|
"final draft to %s/%s failed: %s", self._chat_id, self._thread_id, exc
|
|
)
|
|
|
|
async def _run(self) -> None:
|
|
while not self._broken:
|
|
if self._dirty or time.monotonic() - self._last_sent > _KEEPALIVE:
|
|
self._inflight = asyncio.ensure_future(self._push())
|
|
await asyncio.shield(self._inflight)
|
|
await asyncio.sleep(self._interval)
|
|
|
|
async def _push(self) -> None:
|
|
self._dirty = False
|
|
self._last_sent = time.monotonic()
|
|
rendered, plain = self._render()
|
|
try:
|
|
try:
|
|
await self._send(rendered, "HTML")
|
|
except TelegramBadRequest:
|
|
await self._send(plain, None)
|
|
except TelegramAPIError as exc:
|
|
self._broken = True
|
|
_log.warning(
|
|
"draft to %s/%s stopped: %s", self._chat_id, self._thread_id, exc
|
|
)
|
|
|
|
async def _send(self, text: str, parse_mode: str | None) -> None:
|
|
await self._bot.send_message_draft(
|
|
chat_id=self._chat_id,
|
|
draft_id=self._draft_id,
|
|
message_thread_id=self._thread_id,
|
|
text=text,
|
|
parse_mode=parse_mode,
|
|
)
|
|
|
|
def _render(self) -> tuple[str, str]:
|
|
tail = self.text[-_TAIL:]
|
|
if not tail.strip():
|
|
return to_html(self.status), self.status
|
|
rendered = f"{to_html(self.status)}\n\n{to_html_tail(self.text, _TAIL)}"
|
|
return rendered, f"{self.status}\n\n{tail}"
|