Merge pull request #2656 from fluo-rite/fix/telegram-mention-entities

fix(telegram): respect entity boundaries when parsing bot mentions

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'

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.

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)