From 607a94f324d03e5b26969617dda932df3e26852f Mon Sep 17 00:00:00 2001 From: fluo-rite <137024811+fluo-rite@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:28:58 -0400 Subject: [PATCH] fix(telegram): respect entity boundaries when parsing bot mentions --- src/langbot/pkg/platform/sources/telegram.py | 29 +++-- .../platform/test_telegram_adapter.py | 109 +++++++++++++++++- .../platform/test_telegram_eba_adapter.py | 8 +- 3 files changed, 136 insertions(+), 10 deletions(-) diff --git a/src/langbot/pkg/platform/sources/telegram.py b/src/langbot/pkg/platform/sources/telegram.py index f78b56893..e2cd99741 100644 --- a/src/langbot/pkg/platform/sources/telegram.py +++ b/src/langbot/pkg/platform/sources/telegram.py @@ -149,23 +149,38 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte async def target2yiri(message: telegram.Message, bot: telegram.Bot, bot_account_id: str): message_components = [] - def parse_message_text(text: str) -> list[platform_message.MessageComponent]: + def parse_message_text( + text: str, entities: dict[telegram.MessageEntity, str] + ) -> list[platform_message.MessageComponent]: msg_components = [] + mention = f'@{bot_account_id}' + bot_mentions = [ + entity + for entity, value in entities.items() + if (entity.type == telegram.MessageEntity.MENTION and value == mention) + or (entity.type == telegram.MessageEntity.BOT_COMMAND and value.endswith(mention)) + ] - if f'@{bot_account_id}' in text: + if bot_mentions: + # Telegram entity offsets use UTF-16 code units, not Python character indexes. + encoded = text.encode('utf-16-le') + for entity in sorted(bot_mentions, key=lambda entity: entity.offset, reverse=True): + end = (entity.offset + entity.length) * 2 + start = end - len(mention.encode('utf-16-le')) + encoded = encoded[:start] + encoded[end:] msg_components.append(platform_message.At(target=bot_account_id)) - text = text.replace(f'@{bot_account_id}', '') + text = encoded.decode('utf-16-le') msg_components.append(platform_message.Plain(text=text)) return msg_components if message.text: message_text = message.text - message_components.extend(parse_message_text(message_text)) + message_components.extend(parse_message_text(message_text, message.parse_entities())) if message.photo: if message.caption: - message_components.extend(parse_message_text(message.caption)) + message_components.extend(parse_message_text(message.caption, message.parse_caption_entities())) file = await message.photo[-1].get_file() @@ -192,7 +207,7 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte if message.voice: if message.caption: - message_components.extend(parse_message_text(message.caption)) + message_components.extend(parse_message_text(message.caption, message.parse_caption_entities())) file = await message.voice.get_file() @@ -215,7 +230,7 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte if message.document: if message.caption: - message_components.extend(parse_message_text(message.caption)) + message_components.extend(parse_message_text(message.caption, message.parse_caption_entities())) file = await message.document.get_file() file_name = message.document.file_name or 'document' diff --git a/tests/unit_tests/platform/test_telegram_adapter.py b/tests/unit_tests/platform/test_telegram_adapter.py index bc7af1bf5..3cc6cafd4 100644 --- a/tests/unit_tests/platform/test_telegram_adapter.py +++ b/tests/unit_tests/platform/test_telegram_adapter.py @@ -1,11 +1,12 @@ -"""Tests for Telegram Dify form callback helpers.""" +"""Tests for Telegram message conversion and Dify form callback helpers.""" import json from contextlib import asynccontextmanager +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest -from telegram import ForceReply +from telegram import ForceReply, Message import langbot_plugin.api.entities.builtin.platform.entities as platform_entities import langbot_plugin.api.entities.builtin.platform.events as platform_events @@ -33,6 +34,110 @@ TELEGRAM_BOT_TOKEN = '123456789:AAExampleBotTokenThatMustNotLeak' TELEGRAM_FILE_URL = f'https://api.telegram.org/file/bot{TELEGRAM_BOT_TOKEN}/photos/file_0.jpg' +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('text', 'entities', 'expected_text', 'mentioned'), + [ + ('hello', [], 'hello', False), + ('hello @test_bot', [('mention', 6, 9)], 'hello ', True), + ('@test_bot_extra hello', [('mention', 0, 15)], '@test_bot_extra hello', False), + ('@test_bot', [('code', 0, 9)], '@test_bot', False), + ('@test_bot', [('pre', 0, 9)], '@test_bot', False), + ('https://example.com/@test_bot', [('url', 0, 29)], 'https://example.com/@test_bot', False), + ('/help@test_bot', [('bot_command', 0, 14)], '/help', True), + ('/help@test_bot_extra', [('bot_command', 0, 20)], '/help@test_bot_extra', False), + ('😀 @test_bot!', [('mention', 3, 9)], '😀 !', True), + ( + '@test_bot @test_bot', + [('mention', 0, 9), ('code', 10, 9)], + ' @test_bot', + True, + ), + ( + '@test_bot 😀 @test_bot', + [('mention', 0, 9), ('mention', 13, 9)], + ' 😀 ', + True, + ), + ], +) +async def test_telegram_mentions_use_entity_boundaries(text, entities, expected_text, mentioned): + message = Message.de_json( + { + 'message_id': 10, + 'date': 1700000000, + 'chat': {'id': 123, 'type': 'private'}, + 'text': text, + 'entities': [{'type': kind, 'offset': offset, 'length': length} for kind, offset, length in entities], + }, + None, + ) + + chain = await TelegramMessageConverter.target2yiri(message, None, 'test_bot') + + expected = [platform_message.Plain(text=expected_text)] + if mentioned: + expected.insert(0, platform_message.At(target='test_bot')) + assert list(chain) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize('media_type', ['photo', 'voice', 'document']) +@pytest.mark.parametrize('entity_type', ['mention', 'code']) +async def test_telegram_caption_mentions_preserve_media(monkeypatch, media_type, entity_type): + media = {'file_id': 'file-1', 'file_unique_id': 'unique-1'} + if media_type == 'photo': + media = [{**media, 'width': 1, 'height': 1}] + elif media_type == 'voice': + media = {**media, 'duration': 1} + else: + media = {**media, 'file_name': 'report.txt', 'file_size': 5} + message = Message.de_json( + { + 'message_id': 10, + 'date': 1700000000, + 'chat': {'id': 123, 'type': 'private'}, + 'caption': '😀 @test_bot', + 'caption_entities': [{'type': entity_type, 'offset': 3, 'length': 9}], + media_type: media, + }, + SimpleNamespace(get_file=AsyncMock(return_value=SimpleNamespace(file_path=TELEGRAM_FILE_URL))), + ) + + response = MagicMock(headers={}) + + async def iter_chunked(_chunk_size): + yield b'media' + + response.content.iter_chunked = iter_chunked + + @asynccontextmanager + async def fake_get(url): + assert url == TELEGRAM_FILE_URL + yield response + + monkeypatch.setattr( + 'langbot.pkg.platform.sources.telegram.httpclient.get_session', + lambda **kwargs: MagicMock(get=fake_get), + ) + + chain = list(await TelegramMessageConverter.target2yiri(message, None, 'test_bot')) + + if entity_type == 'mention': + assert chain.pop(0) == platform_message.At(target='test_bot') + assert chain.pop(0) == platform_message.Plain(text='😀 ') + else: + assert chain.pop(0) == platform_message.Plain(text='😀 @test_bot') + assert len(chain) == 1 + expected_type = { + 'photo': platform_message.Image, + 'voice': platform_message.Voice, + 'document': platform_message.File, + } + assert isinstance(chain[0], expected_type[media_type]) + assert chain[0].base64.endswith('bWVkaWE=') + + @pytest.mark.asyncio async def test_telegram_photo_does_not_expose_bot_token_in_image_url(): """Regression test for the Telegram bot-token leak. diff --git a/tests/unit_tests/platform/test_telegram_eba_adapter.py b/tests/unit_tests/platform/test_telegram_eba_adapter.py index eca17ca4a..773a2c9f8 100644 --- a/tests/unit_tests/platform/test_telegram_eba_adapter.py +++ b/tests/unit_tests/platform/test_telegram_eba_adapter.py @@ -311,7 +311,13 @@ async def test_telegram_interaction_callback_answers_and_dispatches_once(): @pytest.mark.asyncio async def test_telegram_converter_maps_message_and_edited_message_events(): - update = make_update({'message': base_message_payload(text='hello @test_bot')}) + update = make_update( + { + 'message': base_message_payload( + text='hello @test_bot', entities=[{'type': 'mention', 'offset': 6, 'length': 9}] + ) + } + ) event = await TelegramEventConverter.target2yiri(update, make_adapter().bot, 'test_bot') assert isinstance(event, platform_events.MessageReceivedEvent)