Merge pull request #4186 from pipecat-ai/mb/fix-websocket-disconnect-race-condition

Fix FastAPI WebSocket disconnect race condition

changelog/4186.fixed.md

@@ -0,0 +1 @@
+- Fixed `FastAPIWebsocketTransport` intermittently hanging on shutdown when the remote side (e.g. Twilio) disconnects while audio is being sent. A race condition between the send and receive paths could cause the `on_client_disconnected` callback to be skipped, leaving the pipeline waiting for a disconnect signal that never came.

src/pipecat/transports/websocket/fastapi.py

@@ -150,17 +150,9 @@ class FastAPIWebsocketClient:
                 else:
                     await self._websocket.send_text(data)
         except Exception as e:
-            logger.error(
+            logger.warning(
                 f"{self} exception sending data: {e.__class__.__name__} ({e}), application_state: {self._websocket.application_state}"
             )
-            # For some reason the websocket is disconnected, and we are not able to send data
-            # So let's properly handle it and disconnect the transport if it is not already disconnecting
-            if (
-                self._websocket.application_state == WebSocketState.DISCONNECTED
-                and not self.is_closing
-            ):
-                logger.warning("Closing already disconnected websocket!")
-                self._closing = True
 
     async def disconnect(self):
         """Disconnect the WebSocket client."""
@@ -189,7 +181,11 @@ class FastAPIWebsocketClient:
 
     def _can_send(self):
         """Check if data can be sent through the WebSocket."""
-        return self.is_connected and not self.is_closing
+        return (
+            self.is_connected
+            and not self.is_closing
+            and self._websocket.application_state != WebSocketState.DISCONNECTED
+        )
 
     @property
     def is_connected(self) -> bool:

tests/test_fastapi_websocket.py

@@ -4,10 +4,17 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+import asyncio
 import unittest
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, PropertyMock
 
-from pipecat.transports.websocket.fastapi import _WebSocketMessageIterator
+from starlette.websockets import WebSocketState
+
+from pipecat.transports.websocket.fastapi import (
+    FastAPIWebsocketCallbacks,
+    FastAPIWebsocketClient,
+    _WebSocketMessageIterator,
+)
 
 
 class TestWebSocketMessageIterator(unittest.IsolatedAsyncioTestCase):
@@ -66,5 +73,134 @@ class TestWebSocketMessageIterator(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(len(messages), 0)
 
 
+class TestSendDisconnectRace(unittest.IsolatedAsyncioTestCase):
+    """Tests for the race condition in issue #3912.
+
+    When the remote side disconnects while send() is in flight, send() should
+    not set _closing = True, because that flag means "we initiated the close."
+    Setting it from send() prevents the receive loop from firing
+    on_client_disconnected, which can cause the pipeline to hang.
+    """
+
+    def _make_client(self, mock_ws):
+        callbacks = FastAPIWebsocketCallbacks(
+            on_client_connected=AsyncMock(),
+            on_client_disconnected=AsyncMock(),
+            on_session_timeout=AsyncMock(),
+        )
+        client = FastAPIWebsocketClient(mock_ws, callbacks)
+        return client, callbacks
+
+    async def test_send_disconnect_does_not_set_closing(self):
+        """send() should not set _closing when the remote side disconnects."""
+        mock_ws = AsyncMock()
+        type(mock_ws).client_state = PropertyMock(return_value=WebSocketState.CONNECTED)
+        type(mock_ws).application_state = PropertyMock(return_value=WebSocketState.DISCONNECTED)
+        mock_ws.send_bytes.side_effect = Exception("connection closed")
+
+        client, _ = self._make_client(mock_ws)
+
+        await client.send(b"audio data")
+
+        self.assertFalse(client.is_closing)
+
+    async def test_send_suppressed_after_disconnect(self):
+        """After a failed send, _can_send() returns False via application_state.
+
+        Simulates real Starlette behavior: application_state starts CONNECTED,
+        transitions to DISCONNECTED when send_bytes raises (Starlette does this
+        internally on OSError before re-raising as WebSocketDisconnect).
+        """
+        mock_ws = AsyncMock()
+        type(mock_ws).client_state = PropertyMock(return_value=WebSocketState.CONNECTED)
+
+        # application_state transitions from CONNECTED → DISCONNECTED on send failure
+        app_state = {"state": WebSocketState.CONNECTED}
+        type(mock_ws).application_state = PropertyMock(side_effect=lambda: app_state["state"])
+
+        def fail_and_transition(data):
+            app_state["state"] = WebSocketState.DISCONNECTED
+            raise Exception("connection closed")
+
+        mock_ws.send_bytes.side_effect = fail_and_transition
+
+        client, _ = self._make_client(mock_ws)
+
+        # First send: _can_send() passes (app_state CONNECTED), send_bytes raises,
+        # Starlette sets app_state to DISCONNECTED
+        await client.send(b"audio data")
+        # Second send: _can_send() returns False (app_state now DISCONNECTED)
+        await client.send(b"more audio")
+
+        # send_bytes was only called once (the first attempt)
+        mock_ws.send_bytes.assert_called_once()
+
+    async def test_disconnect_callback_fires_when_send_races_receive(self):
+        """Regression test for issue #3912.
+
+        The receive loop is blocked waiting for the next message. Meanwhile,
+        send() is called and hits an exception because the remote side closed.
+        Then the receive loop unblocks and sees the disconnect.
+
+        on_client_disconnected must still fire, because the remote side
+        initiated the close — not us.
+        """
+        send_done = asyncio.Event()
+
+        mock_ws = AsyncMock()
+        type(mock_ws).client_state = PropertyMock(return_value=WebSocketState.CONNECTED)
+        type(mock_ws).application_state = PropertyMock(return_value=WebSocketState.DISCONNECTED)
+        mock_ws.send_bytes.side_effect = Exception("connection closed")
+
+        # receive() blocks until send has completed, then returns disconnect.
+        # This enforces the exact ordering that causes the bug.
+        async def mock_receive():
+            await send_done.wait()
+            return {"type": "websocket.disconnect"}
+
+        mock_ws.receive = mock_receive
+
+        client, callbacks = self._make_client(mock_ws)
+
+        # Simulate the _receive_messages logic from FastAPIWebsocketInputTransport
+        async def receive_loop():
+            try:
+                async for _ in _WebSocketMessageIterator(mock_ws):
+                    pass
+            except Exception:
+                pass
+            if not client.is_closing:
+                await client.trigger_client_disconnected()
+
+        recv_task = asyncio.create_task(receive_loop())
+
+        # Let the receive loop start and block on receive()
+        await asyncio.sleep(0)
+
+        # send() races — hits exception but does NOT set _closing
+        await client.send(b"audio data")
+        self.assertFalse(client.is_closing)
+
+        # Unblock the receive loop — it sees the disconnect
+        send_done.set()
+        await recv_task
+
+        # The callback fires because _closing was not poisoned by send()
+        callbacks.on_client_disconnected.assert_called_once()
+
+    async def test_send_text_disconnect_does_not_set_closing(self):
+        """Same as test_send_disconnect_does_not_set_closing but with text data."""
+        mock_ws = AsyncMock()
+        type(mock_ws).client_state = PropertyMock(return_value=WebSocketState.CONNECTED)
+        type(mock_ws).application_state = PropertyMock(return_value=WebSocketState.DISCONNECTED)
+        mock_ws.send_text.side_effect = Exception("connection closed")
+
+        client, _ = self._make_client(mock_ws)
+
+        await client.send("text data")
+
+        self.assertFalse(client.is_closing)
+
+
 if __name__ == "__main__":
     unittest.main()