Work outside main thread and add auto reconnect in WebSockets

spotapi/status.py

@@ -3,7 +3,7 @@
 from spotapi.login import Login
 from spotapi.types.annotations import enforce
 from spotapi.types.data import PlayerState, Devices, Track
-from spotapi.websocket import WebsocketStreamer
+from spotapi.websocket import WebsocketStreamer, WebSocketError
 from typing import Dict, Any, Callable, List, ParamSpec, TypeVar
 
 __all__ = [
@@ -47,8 +47,15 @@ def __init__(self, login: Login, s_device_id: str | None = None) -> None:
 
     def renew_state(self) -> None:
         self._device_dump = self.connect_device()
-        self._state = self._device_dump["player_state"]
-        self._devices = self._device_dump["devices"]
+
+        if type(self._device_dump) != dict:
+            raise WebSocketError("Invalid device dump received", error=str(self._device_dump))
+
+        self._state = self._device_dump.get("player_state")
+        if self._state == None:
+            raise WebSocketError("Could not obtain 'player_state' from connect_device response", error=str(self._device_dump))
+
+        self._devices = self._device_dump.get("devices", {})
 
     @functools.cached_property
     def saved_state(self) -> PlayerState:

spotapi/websocket.py

@@ -5,6 +5,7 @@
 import json
 import time
 import signal
+import traceback
 from typing import Any, Dict
 from spotapi.login import Login
 from spotapi.client import BaseClient
@@ -36,35 +37,84 @@ class WebsocketStreamer:
         "ws_dump",
         "connection_id",
         "keep_alive_thread",
+        "supervisor_thread",
     )
 
     def __init__(self, login: Login) -> None:
         if not login.logged_in:
             raise ValueError("Must be logged in")
+
+        self.login = login
 
-        self.base = BaseClient(login.client)
-        self.client = self.base.client
+        self.device_id = random_hex_string(32)
 
-        self.base.get_session()
-        self.base.get_client_token()
+        self.rlock = threading.Lock()
+        self.ws_dump = None
+        self.ws = None
+        self.connect()
 
-        self.device_id = random_hex_string(32)
+        self.keep_alive_thread = threading.Thread(
+            target=self.keep_alive,
+            daemon=True,
+        )
+        self.keep_alive_thread.start()
 
+
+        self.supervisor_thread = threading.Thread(
+            target=self._supervise,
+            daemon=True,
+        )
+        self.supervisor_thread.start()
+
+        try:
+            signal.signal(signal.SIGINT, self.handle_interrupt)
+        except ValueError:  #< Not running in the main thread
+            pass
+
+        def _cleanup():
+            print("Websockets closing due to program ending")
+            self.disconnect()
+        atexit.register(_cleanup)
+
+    def _create_websocket(self) -> None:
         uri = f"wss://dealer.spotify.com/?access_token={self.base.access_token}"
+
         self.ws = connect(
             uri,
-            user_agent_header="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
+            user_agent_header=(
+                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                "AppleWebKit/537.36 (KHTML, like Gecko) "
+                "Chrome/120.0.0.0 Safari/537.36"
+            ),
         )
 
-        self.rlock = threading.Lock()
-        self.ws_dump: Dict[Any, Any] | None = None
         self.connection_id = self.get_init_packet()
+
+    def connect(self):
+        with self.rlock:
+            self.base = BaseClient(self.login.client)
+            self.client = self.base.client
 
-        self.keep_alive_thread = threading.Thread(target=self.keep_alive, daemon=True)
-        self.keep_alive_thread.start()
+            self.base.get_session()
+            self.base.get_client_token()
 
-        atexit.register(self.ws.close)
-        signal.signal(signal.SIGINT, self.handle_interrupt)
+            self._create_websocket()
+
+            self.register_device()
+
+    def disconnect(self):
+        with self.rlock:
+            if self.ws == None:
+                return
+            try:
+                self.ws.close()
+                self.ws = None
+            except Exception as e:
+                print(f"Failed to close websocket: {e}")
+
+    def reconnect(self) -> None:
+        self.disconnect()
+        self.connect()
 
     def register_device(self) -> None:
         url = f"https://gue1-spclient.spotify.com/track-playback/v1/devices"
@@ -107,7 +157,19 @@ def register_device(self) -> None:
         resp = self.client.post(url, json=payload, authenticate=True)
 
         if resp.fail:
-            raise WebSocketError("Could not register device", error=resp.error.string)
+            try:
+                print("\nREGISTER DEVICE FAILED")
+                print(f"device_id     = {self.device_id}")
+                print(f"connection_id = {self.connection_id}")
+                print(f"error         = {resp.error.string}")
+                print(f"response      = {resp.response}")
+            except Exception:
+                pass
+
+            raise WebSocketError(
+                "Could not register device",
+                error=resp.error.string
+            )
 
     def connect_device(self) -> Dict[str, Any]:
         url = f"https://gue1-spclient.spotify.com/connect-state/v1/devices/hobs_{self.device_id}"
@@ -130,38 +192,102 @@ def connect_device(self) -> Dict[str, Any]:
         resp = self.client.put(url, json=payload, authenticate=True, headers=headers)
 
         if resp.fail:
-            raise WebSocketError("Could not connect device", error=resp.error.string)
+            try:
+                print("\nCONNECT DEVICE FAILED")
+                print(f"device_id     = {self.device_id}")
+                print(f"connection_id = {self.connection_id}")
+                print(f"error         = {resp.error.string}")
+                print(f"response      = {resp.response}")
+            except Exception:
+                pass
+
+            raise WebSocketError(
+                "Could not connect device",
+                error=resp.error.string
+            )
 
         return resp.response
 
     def keep_alive(self) -> None:
         while True:
             try:
                 time.sleep(60)
+
                 with self.rlock:
                     self.ws.send('{"type":"ping"}')
-            except (ConnectionError, KeyboardInterrupt):
-                break
 
-    def get_packet(self) -> Dict[Any, Any]:
-        with self.rlock:
-            ws_dump = dict(json.loads(self.ws.recv()))
-            self.ws_dump = ws_dump
-            return self.ws_dump
+            except Exception as e:
+                print("Websocket disconnected, reconnecting...")
+
+                try:
+                    self.reconnect()
+                except Exception as reconnectError:
+                    print(f"Reconnect failed: {reconnectError}")
+                time.sleep(5)
+
+    def get_packet(self):
+        while True:
+            try:
+                with self.rlock:
+                    ws_dump = dict(json.loads(self.ws.recv()))
+
+                self.ws_dump = ws_dump
+                return ws_dump
+
+            except Exception as e:
+                print(e)
+                try:
+                    self.reconnect()
+                except Exception as e:
+                    print("Failed to reconnect!")
+
+                time.sleep(1)
 
     def get_init_packet(self) -> str:
         """Gets the Spotify Connection ID in the init packet"""
-        packet = self.get_packet()
+        self.ws_dump = dict(json.loads(self.ws.recv()))
 
         if (
-            packet.get("headers") is None
-            or dict(packet["headers"]).get("Spotify-Connection-Id") is None
+            self.ws_dump.get("headers") is None
+            or dict(self.ws_dump["headers"]).get("Spotify-Connection-Id") is None
         ):
             raise ValueError("Invalid init packet")
 
-        return packet["headers"]["Spotify-Connection-Id"]
+        return self.ws_dump["headers"]["Spotify-Connection-Id"]
 
     def handle_interrupt(self, signum: int, frame: Any) -> None:
         """Handle interrupt signal (Ctrl+C)"""
-        self.ws.close()
+        print("Ctrl+C detected, exiting spotapi")
+        self.disconnect()
         exit(0)
+
+    def _supervise(self) -> None:
+        """Monitor websocket and threads, attempt reconnects when down."""
+        backoff = 1
+        while True:
+            try:
+                need_reconnect = False
+
+                if self.ws == None:
+                    need_reconnect = True
+                else:
+                    try:
+                        closed_attr = self.ws.closed
+                        if closed_attr is True:
+                            need_reconnect = True
+                    except Exception:
+                        need_reconnect = True
+
+                if need_reconnect:
+                    try:
+                        self.reconnect()
+                        backoff = 1
+                    except Exception as e:
+                        print(f"Spotapi failed to reconnect: {e}")
+                        traceback.print_exc()
+                        time.sleep(backoff)
+                        backoff = min(backoff * 2, 60)
+
+                time.sleep(5)
+            except Exception:
+                time.sleep(5)