fix: don't block StatusMessageWatcher exit with 6s sleep on exception

src/apify_client/_status_message_watcher.py

@@ -115,7 +115,8 @@ async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """Cancel the logging task."""
-        await asyncio.sleep(self._final_sleep_time_s)
+        if exc_type is None:
+            await asyncio.sleep(self._final_sleep_time_s)
         await self.stop()
 
     async def _log_changed_status_message(self) -> None:
@@ -169,7 +170,6 @@ def stop(self) -> None:
         """Signal the logging thread to stop logging and wait for it to finish."""
         if not self._logging_thread:
             raise RuntimeError('Logging thread is not active')
-        time.sleep(self._final_sleep_time_s)
         self._stop_logging = True
         self._logging_thread.join()
         self._logging_thread = None
@@ -184,6 +184,8 @@ def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """Stop the logging thread."""
+        if exc_type is None:
+            time.sleep(self._final_sleep_time_s)
         self.stop()
 
     def _log_changed_status_message(self) -> None:

tests/unit/test_logging.py

@@ -44,6 +44,10 @@
 )
 _EXISTING_LOGS_BEFORE_REDIRECT_ATTACH = 3
 
+# Large enough that a real sleep is clearly detectable against `_FAST_EXIT_THRESHOLD_S`.
+_PATCHED_FINAL_SLEEP_S = 5
+_FAST_EXIT_THRESHOLD_S = 1.0
+
 _EXPECTED_MESSAGES_AND_LEVELS = (
     ('2025-05-13T07:24:12.588Z ACTOR: Pulling Docker image of build.', logging.INFO),
     ('2025-05-13T07:24:12.686Z ACTOR: Creating Docker container.', logging.INFO),
@@ -652,3 +656,64 @@ async def test_status_message_watcher_async_restart_after_normal_completion(http
     assert task2 is not task  # New task created
     await task2  # Let it complete (will hit terminal status again)
     assert task2.done()
+
+
+@pytest.mark.usefixtures('mock_api')
+def test_sync_watcher_manual_stop_skips_final_sleep(
+    httpserver: HTTPServer,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Manual `stop()` on the sync watcher must not pay the final sleep — only `__exit__` should."""
+    monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S)
+
+    api_url = httpserver.url_for('/').removesuffix('/')
+    run_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID)
+    watcher = run_client.get_status_message_watcher(check_period=timedelta(seconds=0))
+
+    watcher.start()
+    start = time.monotonic()
+    watcher.stop()
+    elapsed = time.monotonic() - start
+
+    assert elapsed < _FAST_EXIT_THRESHOLD_S, f'stop() should not sleep, took {elapsed:.2f}s'
+
+
+@pytest.mark.usefixtures('mock_api')
+def test_sync_watcher_exit_skips_final_sleep_on_exception(
+    httpserver: HTTPServer,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Exceptional `with`-exit must not pay the final sleep so exceptions propagate immediately."""
+    monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S)
+
+    api_url = httpserver.url_for('/').removesuffix('/')
+    run_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID)
+    watcher = run_client.get_status_message_watcher(check_period=timedelta(seconds=0))
+
+    start = time.monotonic()
+    with pytest.raises(RuntimeError, match='boom'), watcher:
+        raise RuntimeError('boom')
+    elapsed = time.monotonic() - start
+
+    assert elapsed < _FAST_EXIT_THRESHOLD_S, f'__exit__ should skip final sleep on exception, took {elapsed:.2f}s'
+
+
+@pytest.mark.usefixtures('mock_api')
+async def test_async_watcher_aexit_skips_final_sleep_on_exception(
+    httpserver: HTTPServer,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Exceptional `async with`-exit must not pay the final sleep so exceptions propagate immediately."""
+    monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S)
+
+    api_url = httpserver.url_for('/').removesuffix('/')
+    run_client = ApifyClientAsync(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID)
+    watcher = await run_client.get_status_message_watcher(check_period=timedelta(seconds=0))
+
+    start = time.monotonic()
+    with pytest.raises(RuntimeError, match='boom'):
+        async with watcher:
+            raise RuntimeError('boom')
+    elapsed = time.monotonic() - start
+
+    assert elapsed < _FAST_EXIT_THRESHOLD_S, f'__aexit__ should skip final sleep on exception, took {elapsed:.2f}s'