fix: reconnect to platform events websocket after connection drop

Author: vdusek

src/apify/events/_apify_event_manager.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Annotated, Self
 
 import websockets.asyncio.client
+import websockets.exceptions
 from pydantic import Discriminator, TypeAdapter
 from typing_extensions import Unpack, override
 
@@ -91,49 +92,70 @@ async def __aexit__(
         exc_value: BaseException | None,
         exc_traceback: TracebackType | None,
     ) -> None:
-        if self._platform_events_websocket:
-            await self._platform_events_websocket.close()
-
+        # Cancel the message-processing task first so that closing the websocket below is not treated
+        # as a dropped connection and followed by a reconnect attempt.
         if self._process_platform_messages_task and not self._process_platform_messages_task.done():
             self._process_platform_messages_task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await self._process_platform_messages_task
 
+        if self._platform_events_websocket:
+            await self._platform_events_websocket.close()
+
         await super().__aexit__(exc_type, exc_value, exc_traceback)
 
     async def _process_platform_messages(self, ws_url: str) -> None:
+        def process_exception(exc: Exception) -> Exception | None:
+            # Until the first connection succeeds, treat every error as fatal so that `__aenter__` fails fast.
+            # Afterwards, treat every error as transient — the reconnect iterator keeps retrying with backoff
+            # so that platform events (e.g. `MIGRATING`) are not missed for the rest of the run.
+            if self._connected_to_platform_websocket is None or not self._connected_to_platform_websocket.done():
+                return exc
+            return None
+
         try:
-            async with websockets.asyncio.client.connect(ws_url) as websocket:
+            async for websocket in websockets.asyncio.client.connect(ws_url, process_exception=process_exception):
                 self._platform_events_websocket = websocket
-                if self._connected_to_platform_websocket is not None:
-                    self._connected_to_platform_websocket.set_result(True)
-
-                async for message in websocket:
-                    try:
-                        parsed_message = event_data_adapter.validate_json(message)
-
-                        if isinstance(parsed_message, DeprecatedEvent):
-                            continue
-
-                        if isinstance(parsed_message, UnknownEvent):
-                            logger.info(
-                                f'Unknown message received: event_name={parsed_message.name}, '
-                                f'event_data={parsed_message.data}'
+                connected_future = self._connected_to_platform_websocket
+                if connected_future is not None and not connected_future.done():
+                    connected_future.set_result(True)
+
+                try:
+                    async for message in websocket:
+                        try:
+                            parsed_message = event_data_adapter.validate_json(message)
+
+                            if isinstance(parsed_message, DeprecatedEvent):
+                                continue
+
+                            if isinstance(parsed_message, UnknownEvent):
+                                logger.info(
+                                    f'Unknown message received: event_name={parsed_message.name}, '
+                                    f'event_data={parsed_message.data}'
+                                )
+                                continue
+
+                            self.emit(
+                                event=parsed_message.name,
+                                event_data=parsed_message.data
+                                if not isinstance(parsed_message.data, SystemInfoEventData)
+                                else parsed_message.data.to_crawlee_format(self._configuration.dedicated_cpus or 1),
                             )
-                            continue
-
-                        self.emit(
-                            event=parsed_message.name,
-                            event_data=parsed_message.data
-                            if not isinstance(parsed_message.data, SystemInfoEventData)
-                            else parsed_message.data.to_crawlee_format(self._configuration.dedicated_cpus or 1),
-                        )
-
-                        if parsed_message.name == Event.MIGRATING:
-                            await self._emit_persist_state_event_rec_task.stop()
-                            self.emit(event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=True))
-                    except Exception:
-                        logger.exception('Cannot parse Actor event', extra={'raw_message': message})
+
+                            if parsed_message.name == Event.MIGRATING:
+                                await self._emit_persist_state_event_rec_task.stop()
+                                self.emit(
+                                    event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=True)
+                                )
+                        except Exception:
+                            logger.exception('Cannot parse Actor event', extra={'raw_message': message})
+                except websockets.exceptions.ConnectionClosed:
+                    pass
+
+                logger.warning(
+                    f'Connection to platform events websocket was closed '
+                    f'(code={websocket.close_code}, reason={websocket.close_reason!r}), reconnecting...'
+                )
         except Exception:
             logger.exception('Error in websocket connection')
             if self._connected_to_platform_websocket is not None and not self._connected_to_platform_websocket.done():

tests/unit/events/test_apify_event_manager.py

@@ -12,7 +12,6 @@
 import pytest
 import websockets
 import websockets.asyncio.server
-import websockets.exceptions
 
 from crawlee.events._types import Event
 
@@ -320,38 +319,32 @@ def migrating_listener(data: Any) -> None:
         assert len(migration_persist_events) >= 1
 
 
-async def test_websocket_mid_stream_disconnect_does_not_raise_invalid_state_error(
-    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
-) -> None:
-    """Regression: a mid-stream websocket disconnect after a successful connect must not raise InvalidStateError.
-
-    The `_connected_to_platform_websocket` future is resolved to `True` on successful connect. If the websocket
-    later drops, the outer `except` in `_process_platform_messages` must not call `set_result(False)` on the
-    already-resolved future.
-    """
+async def test_websocket_reconnects_after_connection_drop(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that after a mid-stream websocket drop, the manager reconnects and keeps receiving platform events."""
     async with (
         _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected),
         ApifyEventManager(Configuration.get_global_configuration()) as event_manager,
     ):
         await client_connected.wait()
+        aborting_calls: list[Any] = []
 
-        # Force an abnormal close from the server so the client's `async for` raises ConnectionClosedError.
+        def listener(data: Any) -> None:
+            aborting_calls.append(data)
+
+        event_manager.on(event=Event.ABORTING, listener=listener)
+
+        # Drop the connection abnormally from the server side.
+        client_connected.clear()
         for ws in list(connected_ws_clients):
             await ws.close(code=1011, reason='Simulated server error')
 
-        task = event_manager._process_platform_messages_task
-        assert task is not None
-        await asyncio.wait_for(asyncio.shield(task), timeout=2.0)
-
-        exc = task.exception()
-        assert not isinstance(exc, asyncio.InvalidStateError), f'Task raised InvalidStateError: {exc}'
+        # The event manager should reconnect on its own.
+        await asyncio.wait_for(client_connected.wait(), timeout=5.0)
 
-        # Confirm the test actually exercised the disconnect path — the outer `except` in
-        # `_process_platform_messages` should have logged a `ConnectionClosedError`.
-        logged_exc_types = [
-            record.exc_info[0] for record in caplog.records if record.exc_info and record.exc_info[0] is not None
-        ]
-        assert any(issubclass(exc_type, websockets.exceptions.ConnectionClosedError) for exc_type in logged_exc_types)
+        # Events sent over the new connection must still be received.
+        websockets.broadcast(connected_ws_clients, json.dumps({'name': 'aborting'}))
+        await poll_until_condition(lambda: bool(aborting_calls), poll_interval=0.05)
+        assert len(aborting_calls) == 1
 
 
 async def test_malformed_message_logs_exception(