fix: prevent duplicate platform websocket when nesting ApifyEventManager

src/apify/events/_apify_event_manager.py

@@ -92,9 +92,19 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp
         self._connected_to_platform_websocket: asyncio.Future[bool] | None = None
         """Future that resolves when the connection to the platform websocket is established."""
 
+        self._context_depth = 0
+        """Nesting depth of active contexts; the outermost enter/exit (0 <-> 1) starts/tears down the websocket."""
+
     @override
     async def __aenter__(self) -> Self:
         await super().__aenter__()
+
+        # Track nesting depth ourselves rather than reading the parent's private ref count. Only the outermost
+        # enter (0 -> 1) starts the websocket machinery; a nested enter reuses the existing connection.
+        self._context_depth += 1
+        if self._context_depth > 1:
+            return self
+
         self._connected_to_platform_websocket = asyncio.Future()
 
         # Run tasks but don't await them
@@ -119,16 +129,20 @@ async def __aexit__(
         exc_value: BaseException | None,
         exc_traceback: TracebackType | None,
     ) -> None:
-        # Cancel the task before closing the websocket so that the closed connection is not treated as a drop
-        # 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()
-
+        # Only the outermost exit (1 -> 0) tears down the websocket machinery; a nested exit must leave the
+        # connection and its processing task intact for the still-active outer context.
+        if self._context_depth == 1:
+            # Cancel the task before closing the websocket so that the closed connection is not treated as a drop
+            # 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()
+
+        self._context_depth -= 1
         await super().__aexit__(exc_type, exc_value, exc_traceback)
 
     def _process_connection_exception(self, exc: Exception) -> Exception | None:

tests/unit/events/test_apify_event_manager.py

@@ -277,6 +277,69 @@ async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
         assert len(connected_ws_clients) == 1
 
 
+async def test_reentrant_context_reuses_single_platform_connection(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A nested `async with` on the same event manager reuses the one platform connection, not a second."""
+    async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+        async with event_manager:
+            await client_connected.wait()
+            assert len(connected_ws_clients) == 1
+            outer_task = event_manager._process_platform_messages_task
+
+            event_calls: list[Any] = []
+            event_manager.on(event=Event.SYSTEM_INFO, listener=event_calls.append)
+
+            async with event_manager:
+                # The nested enter must not replace the message-processing task nor open a second connection.
+                assert event_manager._process_platform_messages_task is outer_task
+                await asyncio.sleep(0.2)
+                assert len(connected_ws_clients) == 1
+
+                # A single platform event must be delivered exactly once, not once per connection.
+                websockets.broadcast(
+                    connected_ws_clients, json.dumps({'name': 'systemInfo', 'data': DUMMY_SYSTEM_INFO})
+                )
+                await poll_until_condition(lambda: len(event_calls) >= 1, poll_interval=0.05)
+                await asyncio.sleep(0.2)
+                assert len(event_calls) == 1
+
+        # The outermost exit tears down the single shared connection and its processing task; nothing is leaked.
+        assert outer_task is not None
+        assert outer_task.done()
+        await poll_until_condition(lambda: len(connected_ws_clients) == 0, poll_interval=0.05)
+        assert len(connected_ws_clients) == 0
+
+
+async def test_reentrant_exit_leaves_outer_context_functional(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Exiting a nested context leaves the outer connection alive and working; only the final exit tears it down."""
+    async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+        async with event_manager:
+            await client_connected.wait()
+            outer_task = event_manager._process_platform_messages_task
+            assert outer_task is not None
+
+            async with event_manager:
+                pass
+
+            # The nested exit must not cancel the outer context's processing task.
+            assert event_manager.active is True
+            assert not outer_task.done()
+
+            # Events must still be delivered on the surviving connection.
+            event_calls: list[Any] = []
+            event_manager.on(event=Event.SYSTEM_INFO, listener=event_calls.append)
+            websockets.broadcast(connected_ws_clients, json.dumps({'name': 'systemInfo', 'data': DUMMY_SYSTEM_INFO}))
+            await poll_until_condition(lambda: len(event_calls) == 1, poll_interval=0.05)
+            assert len(event_calls) == 1
+
+        # The final exit tears everything down; no connection or task is leaked.
+        assert event_manager.active is False
+        assert outer_task.done()
+        await poll_until_condition(lambda: len(connected_ws_clients) == 0, poll_interval=0.05)
+        assert len(connected_ws_clients) == 0
+
+
 async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
     async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):