fix: honor error_handler replacement requests on SessionError and preserve forefront on retries

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -1148,6 +1148,56 @@ def _check_url_patterns(
         # The URL does not match any `include` pattern - reject it
         return False
 
+    async def _handle_error_handler_replacement(
+        self,
+        context: TCrawlingContext | BasicCrawlingContext,
+        error: Exception,
+        *,
+        request_manager: RequestManager,
+        retire_session: bool = False,
+    ) -> bool:
+        """Invoke the user `error_handler` and enqueue a replacement request when appropriate.
+
+        A returned request with the same `unique_key` as the original is treated as no
+        replacement: `add_request` would be a no-op while the original is still in progress,
+        and marking the original handled would drop the work.
+
+        Args:
+            context: Current crawling context.
+            error: Exception that triggered the error handler.
+            request_manager: Request manager used to enqueue the replacement.
+            retire_session: When True, retire `context.session` before awaiting I/O so a
+                known-blocked session cannot stay in the pool if a later await fails.
+
+        Returns:
+            True if the original request was replaced and marked handled; False to continue
+            with the normal retry or failure path.
+        """
+        if not self._error_handler:
+            return False
+
+        try:
+            new_request = await self._error_handler(context, error)
+        except Exception as e:
+            context.request.state = RequestState.ERROR
+            raise UserDefinedErrorHandlerError('Exception thrown in user-defined request error handler') from e
+
+        if new_request is None or new_request.unique_key == context.request.unique_key:
+            return False
+
+        # Preserve counters so a handler minting a fresh unique_key each time cannot loop forever.
+        new_request.retry_count = context.request.retry_count
+        if context.request.session_rotation_count is not None:
+            new_request.session_rotation_count = context.request.session_rotation_count
+
+        if retire_session and context.session:
+            context.session.retire()
+
+        await request_manager.add_request(new_request, forefront=new_request.forefront)
+        await self._mark_request_as_handled(context.request)
+        self._statistics.record_request_processing_finish(context.request.unique_key)
+        return True
+
     async def _handle_request_retries(
         self,
         context: TCrawlingContext | BasicCrawlingContext,
@@ -1169,18 +1219,10 @@ async def _handle_request_retries(
             )
             await self._statistics.error_tracker.add(error=error, context=context)
 
-            if self._error_handler:
-                try:
-                    new_request = await self._error_handler(context, error)
-                except Exception as e:
-                    raise UserDefinedErrorHandlerError('Exception thrown in user-defined request error handler') from e
-                else:
-                    if new_request is not None and new_request != request:
-                        await request_manager.add_request(new_request)
-                        await self._mark_request_as_handled(request)
-                        return
-
-            await request_manager.reclaim_request(request)
+            if await self._handle_error_handler_replacement(context, error, request_manager=request_manager):
+                return
+
+            await request_manager.reclaim_request(request, forefront=request.forefront)
         else:
             request.state = RequestState.ERROR
             await self._mark_request_as_handled(request)
@@ -1474,22 +1516,42 @@ async def __run_task_function(self) -> None:
             if not session:
                 raise RuntimeError('SessionError raised in a crawling context without a session') from session_error
 
-            if self._error_handler:
-                await self._error_handler(context, session_error)
+            request.state = RequestState.ERROR_HANDLER
 
             if self._should_retry_request(context, session_error):
+                await self._statistics.error_tracker_retry.add(error=session_error, context=context)
+
+                # Replacement requests are only honored while rotations remain, so exhausted rotations
+                # still go through failed_request_handler instead of being silently replaced.
+                if await self._handle_error_handler_replacement(
+                    context,
+                    session_error,
+                    request_manager=request_manager,
+                    retire_session=True,
+                ):
+                    return
+
                 exc_only = ''.join(traceback.format_exception_only(session_error)).strip()
                 self._logger.warning('Encountered "%s", rotating session and retrying...', exc_only)
 
-                if session:
-                    session.retire()
+                session.retire()
 
                 # Increment session rotation count.
                 request.session_rotation_count = (request.session_rotation_count or 0) + 1
 
-                await request_manager.reclaim_request(request)
-                await self._statistics.error_tracker_retry.add(error=session_error, context=context)
+                await request_manager.reclaim_request(request, forefront=request.forefront)
             else:
+                # Still invoke the error_handler for side effects, but never replace once rotations are exhausted.
+                if self._error_handler:
+                    try:
+                        await self._error_handler(context, session_error)
+                    except Exception as e:
+                        request.state = RequestState.ERROR
+                        raise UserDefinedErrorHandlerError(
+                            'Exception thrown in user-defined request error handler'
+                        ) from e
+
+                request.state = RequestState.ERROR
                 await self._mark_request_as_handled(request)
 
                 await self._handle_failed_request(context, session_error)

tests/unit/crawlers/_basic/test_basic_crawler.py

@@ -15,7 +15,7 @@
 from datetime import timedelta
 from itertools import product
 from typing import TYPE_CHECKING, Any, Literal, cast
-from unittest.mock import AsyncMock, Mock, call, patch
+from unittest.mock import ANY, AsyncMock, Mock, call, patch
 
 import pytest
 
@@ -274,6 +274,131 @@ async def error_handler(context: BasicCrawlingContext, error: Exception) -> None
     assert error_handler_mock.call_count == 1
 
 
+async def test_session_error_handler_can_replace_request() -> None:
+    """`error_handler` return value must be honored for SessionError while rotations remain."""
+    queue = await RequestQueue.open()
+    crawler = BasicCrawler(request_manager=queue, max_session_rotations=3)
+
+    request = Request.from_url('https://a.placeholder.com')
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        if '|recovered' in context.request.unique_key:
+            return
+        raise SessionError('blocked')
+
+    @crawler.error_handler
+    async def error_handler(context: BasicCrawlingContext, error: Exception) -> Request | None:
+        return Request.from_url(
+            context.request.url,
+            unique_key=f'{context.request.unique_key}|recovered',
+        )
+
+    await crawler.run([request])
+
+    original_request = await queue.get_request(request.unique_key)
+    recovered_request = await queue.get_request(f'{request.unique_key}|recovered')
+
+    assert original_request is not None
+    assert original_request.was_already_handled
+    assert recovered_request is not None
+    assert recovered_request.state == RequestState.DONE
+    assert recovered_request.was_already_handled
+
+
+async def test_session_error_handler_same_unique_key_rotates() -> None:
+    """Returning a request with the same unique_key must rotate, not drop the original."""
+    handler_calls = 0
+    failed_calls = 0
+    crawler = BasicCrawler(max_session_rotations=3)
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        nonlocal handler_calls
+        handler_calls += 1
+        raise SessionError('blocked')
+
+    @crawler.error_handler
+    async def error_handler(context: BasicCrawlingContext, error: Exception) -> Request | None:
+        return Request.from_url(context.request.url)
+
+    @crawler.failed_request_handler
+    async def failed_request_handler(context: BasicCrawlingContext, error: Exception) -> None:
+        nonlocal failed_calls
+        failed_calls += 1
+
+    stats = await crawler.run(['https://a.placeholder.com'])
+
+    assert handler_calls == 3
+    assert failed_calls == 1
+    assert stats.requests_failed == 1
+
+
+async def test_session_error_handler_replacement_ignored_when_rotations_exhausted() -> None:
+    """When rotations are exhausted, replacement requests must not skip failed_request_handler."""
+    queue = await RequestQueue.open()
+    failed_handler_mock = AsyncMock()
+    crawler = BasicCrawler(request_manager=queue, max_session_rotations=1)
+
+    request = Request.from_url('https://a.placeholder.com')
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        raise SessionError('blocked')
+
+    @crawler.error_handler
+    async def error_handler(context: BasicCrawlingContext, error: Exception) -> Request | None:
+        return Request.from_url(
+            context.request.url,
+            unique_key=f'{context.request.unique_key}|should-not-run',
+        )
+
+    @crawler.failed_request_handler
+    async def failed_request_handler(context: BasicCrawlingContext, error: Exception) -> None:
+        await failed_handler_mock(context, error)
+
+    await crawler.run([request])
+
+    failed_handler_mock.assert_awaited_once()
+    assert await queue.get_request(f'{request.unique_key}|should-not-run') is None
+    original_request = await queue.get_request(request.unique_key)
+    assert original_request is not None
+    assert original_request.state == RequestState.ERROR
+    assert original_request.was_already_handled
+
+
+async def test_session_error_handler_exception_is_wrapped() -> None:
+    """Exceptions from error_handler on SessionError path become UserDefinedErrorHandlerError."""
+    crawler = BasicCrawler(max_session_rotations=2)
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        raise SessionError('blocked')
+
+    @crawler.error_handler
+    async def error_handler(context: BasicCrawlingContext, error: Exception) -> None:
+        raise RuntimeError('Crash in session error handler')
+
+    with pytest.raises(UserDefinedErrorHandlerError):
+        await crawler.run(['https://a.placeholder.com'])
+
+
+async def test_reclaim_uses_request_forefront_flag() -> None:
+    """Retries must reclaim with `request.forefront` so tiered-proxy priority retries stay at the front."""
+    queue = await RequestQueue.open()
+    crawler = BasicCrawler(request_manager=queue, max_request_retries=1)
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        context.request.forefront = True
+        raise RuntimeError('Arbitrary crash for testing purposes')
+
+    with patch.object(queue, 'reclaim_request', wraps=queue.reclaim_request) as reclaim_mock:
+        await crawler.run(['https://a.placeholder.com'])
+
+    reclaim_mock.assert_awaited_once_with(ANY, forefront=True)
+
+
 async def test_handles_error_in_error_handler() -> None:
     crawler = BasicCrawler(max_request_retries=3)