Compare commits

...

1 Commits

Author SHA1 Message Date
James Hush
ae8b9f0756 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.
2025-10-01 10:23:24 +08:00

View File

@@ -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."""