3333from typing import Any
3434
3535try :
36- from websocket import WebSocket , create_connection
36+ from websocket import WebSocket , WebSocketException , create_connection
3737except ImportError :
3838 raise ImportError (
3939 "cup_remote requires the 'websocket-client' package.\n "
@@ -67,10 +67,27 @@ def __init__(self, url: str, *, timeout: float = 30.0) -> None:
6767
6868 def connect (self ) -> dict :
6969 """Connect to the remote CUP server and return its info."""
70- self ._ws = create_connection (self .url , timeout = self .timeout )
70+ self ._ws = create_connection (
71+ self .url ,
72+ timeout = self .timeout ,
73+ enable_multithread = True ,
74+ )
7175 self .info = self ._call ("info" )
7276 return self .info
7377
78+ def _reconnect (self ) -> None :
79+ """Reconnect after a dropped connection."""
80+ try :
81+ if self ._ws :
82+ self ._ws .close ()
83+ except Exception :
84+ pass
85+ self ._ws = create_connection (
86+ self .url ,
87+ timeout = self .timeout ,
88+ enable_multithread = True ,
89+ )
90+
7491 def close (self ) -> None :
7592 if self ._ws :
7693 self ._ws .close ()
@@ -95,19 +112,26 @@ def _call(self, method: str, **params: Any) -> Any:
95112 raise RuntimeError ("Not connected. Call .connect() first." )
96113
97114 with self ._lock :
98- self ._msg_id += 1
99- msg_id = self ._msg_id
115+ for attempt in range (2 ):
116+ try :
117+ self ._msg_id += 1
118+ msg_id = self ._msg_id
100119
101- request = {"id" : msg_id , "method" : method , "params" : params }
102- self ._ws .send (json .dumps (request ))
120+ request = {"id" : msg_id , "method" : method , "params" : params }
121+ self ._ws .send (json .dumps (request ))
103122
104- raw = self ._ws .recv ()
105- response = json .loads (raw )
123+ raw = self ._ws .recv ()
124+ response = json .loads (raw )
106125
107- if "error" in response and response ["error" ] is not None :
108- raise RuntimeError (f"Remote error: { response ['error' ]} " )
126+ if "error" in response and response ["error" ] is not None :
127+ raise RuntimeError (f"Remote error: { response ['error' ]} " )
109128
110- return response .get ("result" )
129+ return response .get ("result" )
130+ except (ConnectionError , OSError , WebSocketException ) as e :
131+ if attempt == 1 :
132+ raise
133+ # Connection dropped — reconnect and retry once
134+ self ._reconnect ()
111135
112136 # -- Session API (mirrors cup.Session) ----------------------------------
113137
0 commit comments