11from __future__ import annotations
22
3+ import asyncio
34import contextlib
45import json
56import logging
5455class AsyncWebsocketClient (BaseWebsocketClient ):
5556 _session : AsyncSession
5657 _ws : AsyncExtensionFromHTTP
58+ _reader_task : asyncio .Task [None ]
5759
5860 def __init__ (
5961 self ,
@@ -65,6 +67,9 @@ def __init__(
6567 ) -> None :
6668 super ().__init__ (api_url , token , max_size = max_size )
6769 self ._session = session if session is not None else AsyncSession ()
70+ self ._new_message = asyncio .Condition ()
71+ self ._fatal_error : BaseException | None = None
72+ self ._id_errors : dict [int , BaseException ] = {}
6873
6974 async def __aenter__ (self ) -> Self :
7075 await self ._session .__aenter__ ()
@@ -74,8 +79,11 @@ async def __aenter__(self) -> Self:
7479 msg = "Server did not upgrade to WebSocket"
7580 raise ReceivingError (msg )
7681 self ._ws = resp .extension
82+ # Authenticate before starting the reader task: these messages have no "id"
83+ # and aren't dispatched through it, so nothing else may read the socket yet.
7784 okay = await self .authentication_phase ()
7885 logger .info ("Authenticated with Home Assistant (%s)" , okay .ha_version )
86+ self ._reader_task = asyncio .create_task (self ._reader_loop ())
7987 await self .supported_features_phase ()
8088 return self
8189
@@ -85,10 +93,50 @@ async def __aexit__(
8593 exc_value : BaseException | None ,
8694 traceback : TracebackType | None ,
8795 ) -> None :
96+ self ._reader_task .cancel ()
97+ with contextlib .suppress (asyncio .CancelledError ):
98+ await self ._reader_task
8899 await self ._ws .close ()
89100 del self ._ws
90101 await self ._session .__aexit__ (exc_type , exc_value , traceback )
91102
103+ async def _reader_loop (self ) -> None :
104+ """
105+ Continuously read messages off the socket and dispatch them by id.
106+
107+ This is the only task allowed to touch the socket directly, so reads can
108+ never interleave or starve a waiting `recv()` call.
109+ """
110+ while True :
111+ try :
112+ data = await self ._async_recv ()
113+ except Exception as e : # noqa: BLE001
114+ # No message id is available yet, so there's no single owner to
115+ # blame this on: the connection is dead for every pending recv().
116+ async with self ._new_message :
117+ self ._fatal_error = e
118+ self ._new_message .notify_all ()
119+ return
120+ await self ._dispatch (data )
121+
122+ async def _dispatch (self , data : dict [str , Any ]) -> None :
123+ """Hand a received message to `handle_recv`, routing any failure to its owner."""
124+ try :
125+ self .handle_recv (data )
126+ except Exception as e : # noqa: BLE001
127+ data_id = data .get ("id" )
128+ async with self ._new_message :
129+ if isinstance (data_id , int ):
130+ # The message carried its own id, so only the caller waiting
131+ # on that id needs to see the failure.
132+ self ._id_errors [data_id ] = e
133+ else :
134+ self ._fatal_error = e
135+ self ._new_message .notify_all ()
136+ return
137+ async with self ._new_message :
138+ self ._new_message .notify_all ()
139+
92140 async def _async_send (self , data : dict [str , Any ]) -> None :
93141 """Send a message to the websocket server."""
94142 if data .get ("type" ) != "auth" :
@@ -147,20 +195,26 @@ async def recv(
147195 msg_id : int ,
148196 ) -> EventResponse | ResultResponse | PingResponse | None :
149197 """Receive a response to a message from the websocket server."""
150- while True :
151- ## have we received a message with the id we're looking for?
152- if self ._result_responses .get (msg_id ) is not None :
153- return self ._result_responses .pop (msg_id )
154- if self ._event_responses .get (msg_id , []):
155- return self ._event_responses [msg_id ].pop (0 )
156- if (
157- self ._ping_responses .get (msg_id ) is not None
158- and self ._ping_responses [msg_id ].end is not None
159- ):
160- return self ._ping_responses .pop (msg_id )
161-
162- ## if not, keep receiving messages until we do
163- self .handle_recv (await self ._async_recv ())
198+ async with self ._new_message :
199+ while True :
200+ ## have we received a message with the id we're looking for?
201+ if self ._result_responses .get (msg_id ) is not None :
202+ return self ._result_responses .pop (msg_id )
203+ if self ._event_responses .get (msg_id , []):
204+ return self ._event_responses [msg_id ].pop (0 )
205+ if (
206+ self ._ping_responses .get (msg_id ) is not None
207+ and self ._ping_responses [msg_id ].end is not None
208+ ):
209+ return self ._ping_responses .pop (msg_id )
210+ ## did our own request fail, or did the connection die outright?
211+ if msg_id in self ._id_errors :
212+ raise self ._id_errors .pop (msg_id )
213+ if self ._fatal_error is not None :
214+ raise self ._fatal_error
215+
216+ ## if not, wait for the reader task to dispatch another message
217+ await self ._new_message .wait ()
164218
165219 async def recv_result (self , msg_id : int ) -> ResultResponse :
166220 """Receive a ResultResponse, raising TypeError if the response is not a ResultResponse."""
0 commit comments