diff --git a/dcfs/core/api/message/__init__.py b/dcfs/core/api/message/__init__.py index fbe864f..f355151 100644 --- a/dcfs/core/api/message/__init__.py +++ b/dcfs/core/api/message/__init__.py @@ -10,11 +10,13 @@ DeleteMessagesReq, DownloadFileReq, DownloadFileResp, + EditMessageMediaReq, EditMessageTextReq, GetPinnedMessageReq, MessageResp, PinMessageReq, SearchMessageReq, + SendFileReq, SendTextReq, ) from dcfs.utils.chained_async_iterator import ChainedAsyncIterator @@ -27,6 +29,10 @@ # Discord's bulk delete caps at 100 messages per request. DELETE_BATCH_SIZE = 100 +DISCORD_MSG_LIMIT = 4000 +OVERFLOW_SENTINEL = "DCFS_OVERFLOW" +OVERFLOW_FILENAME = "overflow.json" + rate = Rate(20, Duration.SECOND) bucket = InMemoryBucket([rate]) limiter = Limiter(bucket, max_delay=60 * 1000) # 60 seconds max delay @@ -42,24 +48,63 @@ def __try_acquire(name: str): async def send_text(self, message: str) -> int: self.__try_acquire("MessageApi.send_text") + + if len(message) <= DISCORD_MSG_LIMIT: + return ( + await self.discord_api.next_bot.send_text( + SendTextReq(chat=self.private_file_channel, text=message) + ) + ).message_id + return ( - await self.discord_api.next_bot.send_text( - SendTextReq(chat=self.private_file_channel, text=message) + await self.discord_api.next_bot.send_file( + SendFileReq( + chat=self.private_file_channel, + name=OVERFLOW_FILENAME, + caption=OVERFLOW_SENTINEL, + buffer=message.encode("utf-8"), + ) ) ).message_id async def edit_message_text(self, message_id: int, message: str) -> int: self.__try_acquire("MessageApi.edit_message_text") + + if len(message) <= DISCORD_MSG_LIMIT: + return ( + await self.discord_api.next_bot.edit_message_text( + EditMessageTextReq( + chat=self.private_file_channel, + message_id=message_id, + text=message, + ) + ) + ).message_id + return ( - await self.discord_api.next_bot.edit_message_text( - EditMessageTextReq( + await self.discord_api.next_bot.edit_message_media( + EditMessageMediaReq( chat=self.private_file_channel, message_id=message_id, - text=message, + name=OVERFLOW_FILENAME, + text=OVERFLOW_SENTINEL, + buffer=message.encode("utf-8"), ) ) ).message_id + async def get_text(self, message: MessageResp) -> str: + """Retrieve the text content of a message, handling overflows.""" + if message.text != OVERFLOW_SENTINEL or not message.document: + return message.text + + logger.debug(f"Retrieving overflowed content for message {message.message_id}") + resp = await self.download_file(message.message_id, 0, -1) + full = bytearray() + async for chunk in resp.chunks: + full.extend(chunk) + return full.decode("utf-8") + async def delete_messages(self, message_ids: Iterable[int]) -> None: """Best-effort deletion of channel messages. diff --git a/dcfs/core/repository/impl/fd/dc_msg.py b/dcfs/core/repository/impl/fd/dc_msg.py index 4b3e5de..c313c09 100644 --- a/dcfs/core/repository/impl/fd/dc_msg.py +++ b/dcfs/core/repository/impl/fd/dc_msg.py @@ -183,7 +183,8 @@ async def get( self._cache_invalidate(fr.message_id) return DCFSFileDesc.empty(fr.name) - fd = DCFSFileDesc.from_dict(json.loads(message.text), name=fr.name) + text = await self._message_api.get_text(message) + fd = DCFSFileDesc.from_dict(json.loads(text), name=fr.name) self._cache_put(fr.message_id, fd) if not validate: diff --git a/dcfs/discord/impl/discord_bot.py b/dcfs/discord/impl/discord_bot.py index cbcb0b0..e3bdb3b 100644 --- a/dcfs/discord/impl/discord_bot.py +++ b/dcfs/discord/impl/discord_bot.py @@ -68,7 +68,7 @@ async def send_file(self, req: SendFileReq) -> SendMessageResp: channel_id = self._parse_channel_id(req.chat) channel = await self._get_channel(channel_id) f = discord.File(io.BytesIO(req.buffer), filename=req.name) - msg = await channel.send(file=f) + msg = await channel.send(content=req.caption, file=f) return SendMessageResp(message_id=msg.id) async def get_messages(self, req: GetMessagesReq) -> GetMessagesResp: @@ -148,7 +148,7 @@ async def edit_message_media(self, req: EditMessageMediaReq) -> Message: except discord.NotFound: raise MessageNotFound(req.message_id) f = discord.File(io.BytesIO(req.buffer), filename=req.name) - await msg.edit(attachments=[f]) + await msg.edit(content=req.text, attachments=[f]) return Message(message_id=msg.id) async def download_file(self, req: DownloadFileReq) -> DownloadFileResp: diff --git a/dcfs/reqres.py b/dcfs/reqres.py index 5207473..5cab32e 100644 --- a/dcfs/reqres.py +++ b/dcfs/reqres.py @@ -96,6 +96,7 @@ class SendFileReq(Chat, FileAttr): class EditMessageMediaReq(Chat, Message): buffer: bytes name: str + text: str = "" @dataclass diff --git a/tests/dcfs/core/api/message/test_overflow.py b/tests/dcfs/core/api/message/test_overflow.py new file mode 100644 index 0000000..d5d840a --- /dev/null +++ b/tests/dcfs/core/api/message/test_overflow.py @@ -0,0 +1,80 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock +from dcfs.core.api.message import MessageApi, OVERFLOW_SENTINEL, OVERFLOW_FILENAME +from dcfs.reqres import SendMessageResp, MessageResp, Document + +@pytest.mark.asyncio +async def test_send_text_overflow(mocker): + # Setup + discord_api = MagicMock() + bot = AsyncMock() + discord_api.next_bot = bot + message_api = MessageApi(discord_api, private_file_channel=123) + + # Large message > 4000 chars + large_text = "a" * 4001 + bot.send_file.return_value = SendMessageResp(message_id=999) + + # Execute + msg_id = await message_api.send_text(large_text) + + # Verify + assert msg_id == 999 + bot.send_file.assert_called_once() + args, _ = bot.send_file.call_args + req = args[0] + assert req.chat == 123 + assert req.name == OVERFLOW_FILENAME + assert req.caption == OVERFLOW_SENTINEL + assert req.buffer == large_text.encode("utf-8") + +@pytest.mark.asyncio +async def test_get_text_overflow(mocker): + # Setup + discord_api = MagicMock() + bot = AsyncMock() + discord_api.next_bot = bot + message_api = MessageApi(discord_api, private_file_channel=123) + + overflow_text = "this is a very large message content" + + # Mock download_file (which is called by get_text for overflowed messages) + async def mock_chunks(): + yield overflow_text.encode("utf-8") + + mock_download = mocker.patch.object( + message_api, + "download_file", + return_value=AsyncMock(chunks=mock_chunks()) + ) + + message = MessageResp( + message_id=999, + text=OVERFLOW_SENTINEL, + document=Document(name=OVERFLOW_FILENAME, size=len(overflow_text)) + ) + + # Execute + retrieved_text = await message_api.get_text(message) + + # Verify + assert retrieved_text == overflow_text + mock_download.assert_called_once_with(999, 0, -1) + +@pytest.mark.asyncio +async def test_get_text_normal(mocker): + # Setup + discord_api = MagicMock() + message_api = MessageApi(discord_api, private_file_channel=123) + + message = MessageResp( + message_id=999, + text="normal message", + document=None + ) + + # Execute + retrieved_text = await message_api.get_text(message) + + # Verify + assert retrieved_text == "normal message"