фиксы сокета (не опять а снова)

Author: ink-developer

src/pymax/core.py

@@ -234,13 +234,7 @@ async def _post_login_tasks(self, sync: bool = True) -> None:
         if sync:
             await self._sync()
 
-        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",
-        )
-
-        start_scheduled_task = self._create_safe_task(
+        start_scheaduled_task = self._create_safe_task(
             self._start_scheduled_tasks(),
             name="start_scheduled_tasks",
         )

src/pymax/exceptions.py

@@ -35,8 +35,11 @@ class SocketSendError(Exception):
     Исключение, вызываемое при ошибке отправки данных через сокет.
     """
 
-    def __init__(self) -> None:
-        super().__init__("Send and wait failed (socket)")
+    def __init__(self, message: str | None = None) -> None:
+        if not message:
+            super().__init__("Send and wait failed (socket)")
+        else:
+            super().__init__(message)
 
 
 class ResponseError(Exception):

src/pymax/mixins/channel.py

@@ -30,6 +30,7 @@ async def resolve_channel_by_name(self, name: str) -> Channel | None:
         ).model_dump(by_alias=True)
 
         data = await self._send_and_wait(opcode=Opcode.LINK_INFO, payload=payload)
+
         if data.get("payload", {}).get("error"):
             MixinsUtils.handle_error(data)
 

src/pymax/mixins/socket.py

@@ -145,7 +145,7 @@ def _create_socket_with_proxy(self, proxy: str) -> socket.socket:
             auth_req = b"\x01" + bytes([len(u)]) + u + bytes([len(p)]) + p
             sock.sendall(auth_req)
 
-            auth_resp = self._ssl_read_exactly(sock, 2)
+            auth_resp = self._recv_exactly_plain(sock, 2)
             if auth_resp != b"\x01\x00":
                 sock.close()
                 raise ConnectionError("SOCKS5 authentication failed")
@@ -220,6 +220,7 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
         :rtype: dict[str, Any] | None
         """
         if user_agent is None or self.headers is None:
+            self.logger.debug("No user agent provided, using default")
             user_agent = self.headers or UserAgentPayload()
 
         if sys.version_info[:2] == (3, 12):
@@ -253,6 +254,11 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             self.logger.error("SSL handshake timeout")
             raise
 
+        for fut in list(self._pending.values()):
+            if not fut.done():
+                fut.set_exception(SocketNotConnectedError())
+        self._pending.clear()
+
         self.is_connected = True
         self._incoming = asyncio.Queue()
         self._outgoing = asyncio.Queue()
@@ -273,6 +279,12 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             )
         )
 
+        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._outgoing_task = self._create_safe_task(
             self._outgoing_loop(), name="outgoing_loop socket task"
         )
@@ -384,6 +396,12 @@ 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())
+
+                self._pending.clear()
+
                 self._socket = None
 
                 if self.reconnect and consecutive_errors < max_consecutive_errors:
@@ -462,8 +480,13 @@ async def _send_and_wait(
             return data
 
         except (ssl.SSLEOFError, ssl.SSLError, ConnectionError, BrokenPipeError) as conn_err:
-            self.logger.warning("Connection lost: %s, attempting reconnect...", conn_err)
+            self.logger.warning("Connection lost while sending: %s", conn_err)
             self.is_connected = False
+            if not fut.done():
+                fut.set_exception(SocketSendError("connection lost during send"))
+
+            self._socket = None
+            raise SocketSendError("Connection lost during send") from conn_err
 
         except asyncio.TimeoutError:
             self.logger.exception("Send and wait failed (opcode=%s, seq=%s)", opcode, msg["seq"])
@@ -473,7 +496,7 @@ async def _send_and_wait(
             raise SocketSendError from exc
 
         finally:
-            self._pending.pop(msg["seq"], None)
+            self._pending.pop(seq_key, None)
 
     @override
     async def _get_chat(self, chat_id: int) -> Chat | None:
@@ -498,4 +521,4 @@ async def send_task():
                 loop = asyncio.get_running_loop()
                 await loop.run_in_executor(None, lambda: self._socket.sendall(packet))
 
-        asyncio.create_task(send_task())
+        self._create_safe_task(send_task())

src/pymax/mixins/websocket.py

@@ -51,6 +51,12 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             user_agent_header=user_agent.header_user_agent,
             proxy=self.proxy,
         )
+
+        for fut in list(self._pending.values()):
+            if not fut.done():
+                fut.set_exception(WebSocketNotConnectedError())
+        self._pending.clear()
+
         self.is_connected = True
         self._incoming = asyncio.Queue()
         self._outgoing = asyncio.Queue()
@@ -61,6 +67,11 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
         self._outgoing_task = self._create_safe_task(
             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)