fix(tgapp): preserve streaming overflow segments (#330)
Keep Telegram streaming overflow messages tied to the primary live message, reuse them across edits, and delete stale overflow replies when streamed content shrinks.
This commit is contained in:
manaskarra authored and GitHub committed
1 parent
ffeae5e59b
commit
c3a1b24175
2 files changed
+255
-3
No files matched your search
+55
-3
@@ -380,6 +380,7 @@ class _TelegramStreamSession:
|
||||
self.sent_segments = 0
|
||||
self.active_display = ""
|
||||
self.pending_display = ""
|
||||
self._edit_overflow_msgs = {}
|
||||
self.retry_until = 0.0
|
||||
self.last_update_at = 0.0
|
||||
self.last_update_raw_len = 0
|
||||
@@ -600,15 +601,66 @@ class _TelegramStreamSession:
|
||||
raise
|
||||
return updated if hasattr(updated, "edit_text") else msg
|
||||
|
||||
def _message_key(self, msg):
|
||||
chat_id = getattr(getattr(msg, "chat", None), "id", None)
|
||||
message_id = getattr(msg, "message_id", None)
|
||||
if chat_id is not None and message_id is not None:
|
||||
return (chat_id, message_id)
|
||||
if message_id is not None:
|
||||
return ("message", message_id)
|
||||
return ("object", id(msg))
|
||||
|
||||
async def _delete_text_once(self, msg):
|
||||
delete = getattr(msg, "delete", None)
|
||||
if delete is None:
|
||||
return
|
||||
try:
|
||||
result = delete()
|
||||
if hasattr(result, "__await__"):
|
||||
await result
|
||||
except RetryAfter as exc:
|
||||
self._set_retry_after(exc)
|
||||
raise
|
||||
except Exception as exc:
|
||||
print(f"[TG stale overflow delete error] {type(exc).__name__}: {exc}", flush=True)
|
||||
|
||||
async def _delete_text(self, msg, wait_retry=True):
|
||||
if wait_retry:
|
||||
await self._retry_call(self._delete_text_once, msg)
|
||||
else:
|
||||
await self._delete_text_once(msg)
|
||||
|
||||
async def _edit_text(self, msg, text, wait_retry=True):
|
||||
segments = _markdown_safe_segments(text) or ["..."]
|
||||
old_key = self._message_key(msg)
|
||||
overflow_msgs = self._edit_overflow_msgs.get(old_key, [])
|
||||
if wait_retry:
|
||||
updated = await self._retry_call(self._edit_text_once, msg, segments[0])
|
||||
else:
|
||||
updated = await self._edit_text_once(msg, segments[0])
|
||||
for segment in segments[1:]:
|
||||
updated = await self._reply_text(segment, wait_retry=wait_retry)
|
||||
return updated if hasattr(updated, "edit_text") else msg
|
||||
primary_msg = updated if hasattr(updated, "edit_text") else msg
|
||||
self._edit_overflow_msgs.pop(old_key, None)
|
||||
|
||||
new_overflow_msgs = []
|
||||
for index, segment in enumerate(segments[1:]):
|
||||
if index < len(overflow_msgs):
|
||||
overflow_msg = overflow_msgs[index]
|
||||
if wait_retry:
|
||||
edited_overflow = await self._retry_call(self._edit_text_once, overflow_msg, segment)
|
||||
else:
|
||||
edited_overflow = await self._edit_text_once(overflow_msg, segment)
|
||||
new_overflow_msgs.append(
|
||||
edited_overflow if hasattr(edited_overflow, "edit_text") else overflow_msg
|
||||
)
|
||||
else:
|
||||
new_overflow_msgs.append(await self._reply_text(segment, wait_retry=wait_retry))
|
||||
|
||||
for stale_msg in overflow_msgs[len(new_overflow_msgs):]:
|
||||
await self._delete_text(stale_msg, wait_retry=wait_retry)
|
||||
|
||||
if new_overflow_msgs:
|
||||
self._edit_overflow_msgs[self._message_key(primary_msg)] = new_overflow_msgs
|
||||
return primary_msg
|
||||
|
||||
async def _upsert_live_message(self, text, wait_retry=True):
|
||||
if self.live_msg is None:
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
FRONTENDS = REPO_ROOT / "frontends"
|
||||
for path in (str(REPO_ROOT), str(FRONTENDS)):
|
||||
if path not in sys.path:
|
||||
sys.path.insert(0, path)
|
||||
|
||||
|
||||
def _install_import_stubs():
|
||||
telegram = types.ModuleType("telegram")
|
||||
telegram.BotCommand = object
|
||||
telegram.InlineKeyboardButton = object
|
||||
telegram.InlineKeyboardMarkup = object
|
||||
|
||||
constants = types.ModuleType("telegram.constants")
|
||||
|
||||
class ChatType:
|
||||
PRIVATE = "private"
|
||||
|
||||
class MessageLimit:
|
||||
MAX_TEXT_LENGTH = 4096
|
||||
|
||||
class ParseMode:
|
||||
MARKDOWN_V2 = "MarkdownV2"
|
||||
|
||||
constants.ChatType = ChatType
|
||||
constants.MessageLimit = MessageLimit
|
||||
constants.ParseMode = ParseMode
|
||||
|
||||
error = types.ModuleType("telegram.error")
|
||||
|
||||
class RetryAfter(Exception):
|
||||
def __init__(self, retry_after=0):
|
||||
super().__init__("retry after")
|
||||
self.retry_after = retry_after
|
||||
|
||||
error.RetryAfter = RetryAfter
|
||||
|
||||
ext = types.ModuleType("telegram.ext")
|
||||
ext.ApplicationBuilder = object
|
||||
ext.CallbackQueryHandler = object
|
||||
ext.MessageHandler = object
|
||||
ext.ContextTypes = types.SimpleNamespace(DEFAULT_TYPE=object)
|
||||
ext.filters = types.SimpleNamespace(
|
||||
COMMAND=object(),
|
||||
PHOTO=object(),
|
||||
TEXT=object(),
|
||||
Document=types.SimpleNamespace(ALL=object()),
|
||||
)
|
||||
|
||||
helpers = types.ModuleType("telegram.helpers")
|
||||
helpers.escape_markdown = lambda text, version=2, entity_type=None: text or ""
|
||||
|
||||
request = types.ModuleType("telegram.request")
|
||||
request.HTTPXRequest = object
|
||||
|
||||
agentmain = types.ModuleType("agentmain")
|
||||
|
||||
class GeneraticAgent:
|
||||
def __init__(self):
|
||||
self.verbose = False
|
||||
self.inc_out = False
|
||||
|
||||
agentmain.GeneraticAgent = GeneraticAgent
|
||||
|
||||
chatapp_common = types.ModuleType("chatapp_common")
|
||||
chatapp_common.FILE_HINT = ""
|
||||
chatapp_common.HELP_TEXT = ""
|
||||
chatapp_common.TELEGRAM_MENU_COMMANDS = []
|
||||
chatapp_common.clean_reply = lambda text: text
|
||||
chatapp_common.ensure_single_instance = lambda *args, **kwargs: None
|
||||
chatapp_common.extract_files = lambda text: []
|
||||
chatapp_common.format_restore = lambda: (([], "", 0), None)
|
||||
chatapp_common.redirect_log = lambda *args, **kwargs: None
|
||||
chatapp_common.require_runtime = lambda *args, **kwargs: None
|
||||
chatapp_common.split_text = lambda text, limit: [text[i : i + limit] for i in range(0, len(text), limit)]
|
||||
|
||||
continue_cmd = types.ModuleType("continue_cmd")
|
||||
continue_cmd.handle_frontend_command = lambda *args, **kwargs: ""
|
||||
continue_cmd.reset_conversation = lambda *args, **kwargs: ""
|
||||
|
||||
llmcore = types.ModuleType("llmcore")
|
||||
llmcore.mykeys = {}
|
||||
|
||||
sys.modules.update(
|
||||
{
|
||||
"telegram": telegram,
|
||||
"telegram.constants": constants,
|
||||
"telegram.error": error,
|
||||
"telegram.ext": ext,
|
||||
"telegram.helpers": helpers,
|
||||
"telegram.request": request,
|
||||
"agentmain": agentmain,
|
||||
"chatapp_common": chatapp_common,
|
||||
"continue_cmd": continue_cmd,
|
||||
"llmcore": llmcore,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_install_import_stubs()
|
||||
tgapp = importlib.import_module("tgapp")
|
||||
|
||||
|
||||
class FakeChat:
|
||||
type = tgapp.ChatType.PRIVATE
|
||||
id = 123
|
||||
|
||||
|
||||
class FakeMessage:
|
||||
_next_id = 1
|
||||
|
||||
def __init__(self, root, text=""):
|
||||
self.root = root
|
||||
self.chat = FakeChat()
|
||||
self.message_id = FakeMessage._next_id
|
||||
FakeMessage._next_id += 1
|
||||
self.text = text
|
||||
self.parse_mode = None
|
||||
self.edit_count = 0
|
||||
self.deleted = False
|
||||
|
||||
async def reply_text(self, text, parse_mode=None):
|
||||
msg = FakeMessage(self.root, text)
|
||||
msg.parse_mode = parse_mode
|
||||
self.root.messages.append(msg)
|
||||
return msg
|
||||
|
||||
async def edit_text(self, text, parse_mode=None):
|
||||
self.text = text
|
||||
self.parse_mode = parse_mode
|
||||
self.edit_count += 1
|
||||
return self
|
||||
|
||||
async def delete(self):
|
||||
self.deleted = True
|
||||
|
||||
|
||||
class FakeRoot(FakeMessage):
|
||||
def __init__(self):
|
||||
self.root = self
|
||||
self.chat = FakeChat()
|
||||
self.message_id = 0
|
||||
self.messages = []
|
||||
|
||||
|
||||
class TelegramStreamSegmentTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_edit_text_keeps_primary_live_message_when_text_overflows(self):
|
||||
root = FakeRoot()
|
||||
session = tgapp._TelegramStreamSession(root)
|
||||
primary = await root.reply_text("thinking...")
|
||||
|
||||
live_msg = await session._edit_text(primary, "a" * 5000, wait_retry=False)
|
||||
|
||||
self.assertIs(live_msg, primary)
|
||||
self.assertEqual(len(root.messages), 2)
|
||||
self.assertEqual(root.messages[0], primary)
|
||||
self.assertEqual(primary.text, "a" * 4096)
|
||||
self.assertEqual(root.messages[1].text, "a" * 904)
|
||||
|
||||
async def test_repeated_overflow_edit_reuses_overflow_message_without_duplicates(self):
|
||||
root = FakeRoot()
|
||||
session = tgapp._TelegramStreamSession(root)
|
||||
primary = await root.reply_text("thinking...")
|
||||
|
||||
live_msg = await session._edit_text(primary, "a" * 5000, wait_retry=False)
|
||||
overflow = root.messages[1]
|
||||
live_msg = await session._edit_text(live_msg, "b" * 6000, wait_retry=False)
|
||||
|
||||
self.assertIs(live_msg, primary)
|
||||
self.assertEqual(len(root.messages), 2)
|
||||
self.assertIs(root.messages[1], overflow)
|
||||
self.assertEqual(primary.text, "b" * 4096)
|
||||
self.assertEqual(overflow.text, "b" * 1904)
|
||||
self.assertEqual(primary.edit_count, 2)
|
||||
self.assertEqual(overflow.edit_count, 1)
|
||||
|
||||
async def test_shrinking_edited_text_deletes_stale_overflow_segment(self):
|
||||
root = FakeRoot()
|
||||
session = tgapp._TelegramStreamSession(root)
|
||||
primary = await root.reply_text("thinking...")
|
||||
|
||||
live_msg = await session._edit_text(primary, "a" * 5000, wait_retry=False)
|
||||
overflow = root.messages[1]
|
||||
live_msg = await session._edit_text(live_msg, "short", wait_retry=False)
|
||||
|
||||
self.assertIs(live_msg, primary)
|
||||
self.assertEqual(primary.text, "short")
|
||||
self.assertTrue(overflow.deleted)
|
||||
self.assertEqual(session._edit_overflow_msgs, {})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user