diff --git a/src/langbot/pkg/platform/sources/telegram.py b/src/langbot/pkg/platform/sources/telegram.py index f78b56893..92aa91705 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}'.lower() + bot_mentions = [ + entity + for entity, value in entities.items() + if (entity.type == telegram.MessageEntity.MENTION and value.lower() == mention) + or (entity.type == telegram.MessageEntity.BOT_COMMAND and value.lower().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..cf661d638 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,127 @@ 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!', [('mention', 3, 9)], '😀 !', True), + ('/help@TeSt_BoT', [('bot_command', 0, 14)], '/help', 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( + ('caption', 'caption_entities', 'expected_text', 'mentioned'), + [ + ('😀 @test_bot', [('mention', 3, 9)], '😀 ', True), + ('😀 @TEST_BOT', [('mention', 3, 9)], '😀 ', True), + ('😀 @test_bot', [('code', 3, 9)], '😀 @test_bot', False), + ('😀 @test_bot', [('pre', 3, 9)], '😀 @test_bot', False), + ('https://example.com/@test_bot', [('url', 0, 29)], 'https://example.com/@test_bot', False), + ('@test_bot_extra', [('mention', 0, 15)], '@test_bot_extra', False), + ('ordinary caption', [], 'ordinary caption', False), + ('😀 @TEST_BOT @test_bot', [('mention', 3, 9), ('mention', 13, 9)], '😀 ', True), + ('😀 @TEST_BOT @test_bot', [('mention', 3, 9), ('code', 13, 9)], '😀 @test_bot', True), + ], +) +async def test_telegram_caption_mentions_preserve_media( + monkeypatch, media_type, caption, caption_entities, expected_text, mentioned +): + 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': caption, + 'caption_entities': [ + {'type': kind, 'offset': offset, 'length': length} for kind, offset, length in caption_entities + ], + 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 mentioned: + assert chain.pop(0) == platform_message.At(target='test_bot') + assert chain.pop(0) == platform_message.Plain(text=expected_text) + 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 24761fa8e..3c2a0a15f 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)