Fix: Ensure on_client_disconnected event always fires in FastAPIWebsocketTransport
This fix addresses an issue where the on_client_disconnected callback would not fire in approximately 5% of calls when using FastAPIWebsocketTransport. The problem was caused by a race condition introduced after v0.0.86 when event handlers were changed to run in parallel. When stop() or cancel() initiated the disconnect, the _closing flag would be set, preventing trigger_client_disconnected() from being called in _receive_messages(). Now disconnect() always calls trigger_client_disconnected() when closing the WebSocket, ensuring the event fires reliably whether the disconnect is initiated by the transport or the remote client. Fixes the same issue as commit 019c1a6d from origin/fastapi_disconnect_issue branch.
This commit is contained in:
@@ -278,6 +278,13 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Main message receiving loop for WebSocket messages."""
|
||||
|
||||
async def trigger_disconnect_if_needed():
|
||||
# Trigger `on_client_disconnected` if the client actually disconnects,
|
||||
# that is, we are not the ones disconnecting.
|
||||
if not self._client.is_closing:
|
||||
await self._client.trigger_client_disconnected()
|
||||
|
||||
try:
|
||||
async for message in self._client.receive():
|
||||
if not self._params.serializer:
|
||||
@@ -294,11 +301,14 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
|
||||
await self.push_frame(frame)
|
||||
except Exception as e:
|
||||
logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})")
|
||||
|
||||
# Trigger `on_client_disconnected` if the client actually disconnects,
|
||||
# that is, we are not the ones disconnecting.
|
||||
if not self._client.is_closing:
|
||||
await self._client.trigger_client_disconnected()
|
||||
finally:
|
||||
# Use shield to prevent cancellation from stopping the disconnect callback
|
||||
try:
|
||||
await asyncio.shield(trigger_disconnect_if_needed())
|
||||
except asyncio.CancelledError:
|
||||
# Even if we're cancelled, try to trigger the disconnect
|
||||
await trigger_disconnect_if_needed()
|
||||
raise
|
||||
|
||||
async def _monitor_websocket(self):
|
||||
"""Wait for self._params.session_timeout seconds, if the websocket is still open, trigger timeout event."""
|
||||
|
||||
Reference in New Issue
Block a user