изменил структуру _pending для хранения дополнительных данных и обновил обработку исключений

Author: ink-developer

src/pymax/core.py

@@ -137,7 +137,7 @@ def __init__(
         self._outgoing: asyncio.Queue[dict[str, Any]] | None = None
         self._recv_task: asyncio.Task[Any] | None = None
         self._outgoing_task: asyncio.Task[Any] | None = None
-        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._pending: dict[int, tuple[asyncio.Future[dict[str, Any]], int, int | None]] = {}
         self._file_upload_waiters: dict[int, asyncio.Future[dict[str, Any]]] = {}
         self._background_tasks: set[asyncio.Task[Any]] = set()
         self._stop_event = asyncio.Event()
@@ -414,8 +414,9 @@ async def _cleanup_client(self):
             self._outgoing_task = None
 
         for fut in self._pending.values():
-            if not fut.done():
-                fut.set_exception(SocketNotConnectedError())
+            old_fut = fut[0]
+            if not old_fut.done():
+                old_fut.set_exception(SocketNotConnectedError())
         self._pending.clear()
 
         if self._socket:

src/pymax/interfaces.py

@@ -114,10 +114,11 @@ async def _cleanup_client(self) -> None:
                 await self._outgoing_task
             self._outgoing_task = None
 
-        for fut in self._pending.values():
-            if not fut.done():
-                fut.set_exception(WebSocketNotConnectedError())
-        self._pending.clear()
+            for pending in self._pending.values():
+                fut = pending[0]
+                if not fut.done():
+                    fut.set_exception(WebSocketNotConnectedError())
+            self._pending.clear()
 
         if self._ws:
             try:
@@ -281,14 +282,30 @@ def _parse_json(self, raw: Any) -> dict[str, Any] | None:
             self.logger.warning("JSON parse error", exc_info=True)
             return None
 
-    def _handle_pending(self, seq: int | None, data: dict) -> bool:
-        if isinstance(seq, int):
-            fut = self._pending.get(seq)
-            if fut and not fut.done():
-                fut.set_result(data)
-                self.logger.debug("Matched response for pending seq=%s", seq)
-                return True
-        return False
+    def _handle_pending(self, seq: int | None, data: dict[str, Any]) -> bool:
+        if seq is None:
+            return False
+
+        pending = self._pending.get(seq)
+        if not pending:
+            return False
+
+        fut, expected_cmd, expected_opcode = pending
+
+        cmd = data.get("cmd")
+        opcode = data.get("opcode")
+
+        if cmd != expected_cmd:
+            return False
+
+        if expected_opcode is not None and opcode != expected_opcode:
+            return False
+
+        if not fut.done():
+            fut.set_result(data)
+
+        self.logger.debug("Matched response for pending seq=%s", seq)
+        return True
 
     async def _handle_incoming_queue(self, data: dict[str, Any]) -> None:
         if self._incoming:
@@ -311,18 +328,6 @@ async def _handle_file_upload(self, data: dict[str, Any]) -> None:
                     fut.set_result(data)
                     self.logger.debug("Fulfilled file upload waiter for %s=%s", key, id_)
 
-    async def _send_notification_response(self, chat_id: int, message_id: str) -> None:
-        if self._socket is not None:
-            return
-        await self._queue_message(
-            opcode=Opcode.NOTIF_MESSAGE,
-            payload={"chatId": chat_id, "messageId": message_id},
-            cmd=1,
-        )
-        self.logger.debug(
-            "Sent NOTIF_MESSAGE_RECEIVED for chat_id=%s message_id=%s", chat_id, message_id
-        )
-
     async def _handle_message_notifications(self, data: dict) -> None:
         if data.get("opcode") != Opcode.NOTIF_MESSAGE.value:
             return
@@ -331,9 +336,6 @@ async def _handle_message_notifications(self, data: dict) -> None:
         if not msg:
             return
 
-        if msg.chat_id and msg.id:
-            await self._send_notification_response(msg.chat_id, str(msg.id))
-
         handlers_map = {
             MessageStatus.EDITED: self._on_message_edit_handlers,
             MessageStatus.REMOVED: self._on_message_delete_handlers,
@@ -538,6 +540,12 @@ async def _sync(self, user_agent: UserAgentPayload | None = None) -> None:
 
             if error := raw_payload.get("error"):
                 MixinsUtils.handle_error(data)
+
+            ping_task = self._create_safe_task(
+                self._send_interactive_ping(),
+                name="interactive_ping",
+            )
+
             chat_marker = raw_payload.get("chatMarker")
             if chat_marker:
                 self.chat_marker = chat_marker

src/pymax/mixins/socket.py

@@ -274,8 +274,9 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             raise
 
         for fut in list(self._pending.values()):
-            if not fut.done():
-                fut.set_exception(SocketNotConnectedError())
+            old_fut = fut[0]
+            if not old_fut.done():
+                old_fut.set_exception(SocketNotConnectedError())
         self._pending.clear()
 
         self.is_connected = True
@@ -419,8 +420,9 @@ async def _recv_loop(self) -> None:
                 self.is_connected = False
 
                 for fut in list(self._pending.values()):
-                    if not fut.done():
-                        fut.set_exception(SocketNotConnectedError())
+                    old_fut = fut[0]
+                    if not old_fut.done():
+                        old_fut.set_exception(SocketNotConnectedError())
 
                 self._pending.clear()
 
@@ -513,8 +515,9 @@ async def _send_and_wait(
             self.logger.warning("Connection lost while sending: %s", conn_err)
             self.is_connected = False
             for pending_fut in list(self._pending.values()):
-                if not pending_fut.done():
-                    pending_fut.set_exception(SocketNotConnectedError())
+                old_fut = pending_fut[0]
+                if not old_fut.done():
+                    old_fut.set_exception(SocketNotConnectedError())
             self._pending.clear()
 
             if not fut.done():

src/pymax/mixins/websocket.py

@@ -49,12 +49,18 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             self.uri,
             origin=WEBSOCKET_ORIGIN,
             user_agent_header=user_agent.header_user_agent,
+            compression="deflate",
+            ping_interval=None,
+            open_timeout=10,
+            close_timeout=10,
             proxy=self.proxy,
+            max_size=None,
         )
 
         for fut in list(self._pending.values()):
-            if not fut.done():
-                fut.set_exception(WebSocketNotConnectedError())
+            old_fut = fut[0]
+            if not old_fut.done():
+                old_fut.set_exception(WebSocketNotConnectedError())
         self._pending.clear()
 
         self.is_connected = True
@@ -68,12 +74,11 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             self._outgoing_loop(), name="outgoing_loop websocket task"
         )
         self.logger.debug("is_connected=%s before starting ping", self.is_connected)
-        ping_task = self._create_safe_task(
-            self._send_interactive_ping(),
-            name="interactive_ping",
-        )
+
         self.logger.info("WebSocket connected, starting handshake")
-        return await self._handshake(user_agent)
+        resp = await self._handshake(user_agent)
+
+        return resp
 
     async def _recv_loop(self) -> None:
         if self._ws is None:
@@ -101,9 +106,10 @@ async def _recv_loop(self) -> None:
                 self.logger.exception(
                     f"WebSocket connection closed with error: {e.code}, {e.reason}; exiting recv loop"
                 )
-                for fut in self._pending.values():
+                for pending in self._pending.values():
+                    fut = pending[0]
                     if not fut.done():
-                        fut.set_exception(WebSocketNotConnectedError)
+                        fut.set_exception(WebSocketNotConnectedError())
                 self._pending.clear()
 
                 self.is_connected = False
@@ -130,11 +136,13 @@ async def _send_and_wait(
         fut: asyncio.Future[dict[str, Any]] = loop.create_future()
         seq_key = msg["seq"]
 
-        old_fut = self._pending.get(seq_key)
-        if old_fut and not old_fut.done():
-            old_fut.cancel()
+        old = self._pending.get(seq_key)
+        if old:
+            old_fut = old[0]
+            if not old_fut.done():
+                old_fut.cancel()
 
-        self._pending[seq_key] = fut
+        self._pending[seq_key] = (fut, 1, int(opcode))
 
         try:
             self.logger.debug(

src/pymax/protocols.py

@@ -60,7 +60,7 @@ def __init__(self, logger: Logger) -> None:
         self.reconnect_delay: float
         self._ws: websockets.ClientConnection | None = None
         self._seq: int = 0
-        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._pending: dict[int, tuple[asyncio.Future[dict[str, Any]], int, int | None]] = {}
         self._recv_task: asyncio.Task[Any] | None = None
         self._incoming: asyncio.Queue[dict[str, Any]] | None = None
         self._file_upload_waiters: dict[