xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-10-22 08:25:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fluo-rite
2026-10-13 02:29:07 -04:00
parent 688e7761e7
commit 607a94f324
3 changed files with 136 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/platform/sources/telegram.py
+22 -7
View File
@@ -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'
tests/unit_tests/platform/test_telegram_adapter.py
+107 -2
View File
@@ -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.
tests/unit_tests/platform/test_telegram_eba_adapter.py
+7 -1
View File
@@ -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)