Return serialized container MsgId on finalize

client/src/telethon/_impl/mtproto/mtp/encrypted.py

@@ -204,9 +204,9 @@ class Encrypted(Mtp):
     def _get_current_salt(self) -> int:
         return self._salts[-1].salt if self._salts else 0
 
-    def _finalize_plain(self) -> bytes:
+    def _finalize_plain(self) -> Optional[Tuple[MsgId, bytes]]:
         if not self._msg_count:
-            return b""
+            return None
 
         if self._msg_count == 1:
             del self._buffer[:CONTAINER_HEADER_LEN]
@@ -216,7 +216,7 @@ class Encrypted(Mtp):
         )
 
         if self._msg_count == 1:
-            container_msg_id: Union[Type[Single], int] = Single
+            container_msg_id = self._last_msg_id
         else:
             container_msg_id = self._get_new_msg_id()
             self._buffer[HEADER_LEN : HEADER_LEN + CONTAINER_HEADER_LEN] = struct.pack(
@@ -235,7 +235,7 @@ class Encrypted(Mtp):
         self._msg_count = 0
         result = bytes(self._buffer)
         self._buffer.clear()
-        return result
+        return MsgId(container_msg_id), result
 
     def _process_message(self, message: Message) -> None:
         if message_requires_ack(message):
@@ -465,12 +465,13 @@ class Encrypted(Mtp):
 
         return self._serialize_msg(body, True)
 
-    def finalize(self) -> bytes:
-        buffer = self._finalize_plain()
-        if not buffer:
-            return buffer
-        else:
-            return encrypt_data_v2(buffer, self._auth_key)
+    def finalize(self) -> Optional[Tuple[MsgId, bytes]]:
+        result = self._finalize_plain()
+        if not result:
+            return None
+
+        msg_id, buffer = result
+        return msg_id, encrypt_data_v2(buffer, self._auth_key)
 
     def deserialize(self, payload: bytes) -> Deserialization:
         check_message_buffer(payload)

client/src/telethon/_impl/mtproto/mtp/plain.py

@@ -1,5 +1,5 @@
 import struct
-from typing import Optional
+from typing import Optional, Tuple
 
 from ..utils import check_message_buffer
 from .types import Deserialization, MsgId, Mtp
@@ -23,10 +23,13 @@ class Plain(Mtp):
         self._buffer += request  # message_data
         return msg_id
 
-    def finalize(self) -> bytes:
+    def finalize(self) -> Optional[Tuple[MsgId, bytes]]:
+        if not self._buffer:
+            return None
+
         result = bytes(self._buffer)
         self._buffer.clear()
-        return result
+        return MsgId(0), result
 
     def deserialize(self, payload: bytes) -> Deserialization:
         check_message_buffer(payload)

client/src/telethon/_impl/mtproto/mtp/types.py

@@ -165,12 +165,23 @@ class Deserialization:
 class Mtp(ABC):
     @abstractmethod
     def push(self, request: bytes) -> Optional[MsgId]:
-        pass
+        """
+        Push a request's body to the internal buffer.
+
+        On success, return the serialized message identifier.
+        """
 
     @abstractmethod
-    def finalize(self) -> bytes:
-        pass
+    def finalize(self) -> Optional[Tuple[MsgId, bytes]]:
+        """
+        Finalize the buffer of serialized requests.
+
+        If the buffer is empty, :data:`None` is returned.
+        Otherwise, the message identifier for the entire buffer and the serialized buffer are returned.
+        """
 
     @abstractmethod
     def deserialize(self, payload: bytes) -> Deserialization:
-        pass
+        """
+        Deserialize incoming buffer payload.
+        """

client/src/telethon/_impl/mtsender/sender.py

@@ -271,8 +271,9 @@ class Sender:
                 else:
                     break
 
-        mtp_buffer = self._mtp.finalize()
-        if mtp_buffer:
+        result = self._mtp.finalize()
+        if result:
+            _, mtp_buffer = result
             self._transport.pack(mtp_buffer, self._writer.write)
             self._write_drain_pending = True