fix(crawlers): Keep Request mutations across handler runs and retries

src/crawlee/_types.py

@@ -2,7 +2,6 @@
 
 import dataclasses
 from collections.abc import Callable, Iterator, Mapping
-from copy import deepcopy
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Annotated, Any, Literal, Protocol, TypedDict, TypeVar, cast, overload
 
@@ -266,20 +265,12 @@ def __init__(
         self,
         *,
         key_value_store_getter: GetKeyValueStoreFunction,
-        request: Request,
     ) -> None:
         self._key_value_store_getter = key_value_store_getter
         self.add_requests_calls = list[AddRequestsKwargs]()
         self.push_data_calls = list[PushDataFunctionCall]()
         self.key_value_store_changes = dict[tuple[str | None, str | None, str | None], KeyValueStoreChangeRecords]()
 
-        # Isolated copies for handler execution
-        self._request = deepcopy(request)
-
-    @property
-    def request(self) -> Request:
-        return self._request
-
     async def add_requests(
         self,
         requests: Sequence[str | Request],
@@ -329,14 +320,6 @@ async def get_key_value_store(
 
         return self.key_value_store_changes[id, name, alias]
 
-    def apply_request_changes(self, target: Request) -> None:
-        """Apply tracked changes from handler copy to original request."""
-        if self.request.user_data != target.user_data:
-            target.user_data = self.request.user_data
-
-        if self.request.headers != target.headers:
-            target.headers = self.request.headers
-
 
 @docs_group('Functions')
 class AddRequestsFunction(Protocol):

src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py

@@ -293,7 +293,8 @@ async def _crawl_one(
 
         `SubCrawlerRun` contains either result of the crawl or the exception that was thrown during the crawl.
         Sub crawler pipeline call is dynamically created based on the `rendering_type`.
-        New copy-like context is created from passed `context` and `state` and is passed to sub crawler pipeline.
+        A new context is created from passed `context` and `state` and is passed to sub crawler pipeline. The original
+        `request` is shared, so its mutations persist. Only `result` and `use_state` are isolated per sub crawler.
         """
         if state is not None:
 
@@ -306,13 +307,13 @@ async def get_input_state(
         else:
             use_state_function = context.use_state
 
-        # New result is created and injected to newly created context. This is done to ensure isolation of sub crawlers.
+        # A fresh result is injected into the new context to isolate `add_requests`/`push_data`/`kvs` calls between
+        # sub crawlers. The `request` itself is shared, so mutations to it persist.
         result = RequestHandlerRunResult(
             key_value_store_getter=self.get_key_value_store,
-            request=context.request,
         )
         context_linked_to_result = BasicCrawlingContext(
-            request=result.request,
+            request=context.request,
             session=context.session,
             proxy_info=context.proxy_info,
             send_request=context.send_request,

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -72,7 +72,6 @@
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
 from ._context_pipeline import ContextPipeline
-from ._context_utils import swapped_context
 from ._logging_utils import (
     get_one_line_error_summary_if_possible,
     reduce_asyncio_timeout_error_to_relevant_traceback_parts,
@@ -1346,8 +1345,6 @@ async def _commit_request_handler_result(self, context: BasicCrawlingContext) ->
 
         await self._commit_key_value_store_changes(result, get_kvs=self.get_key_value_store)
 
-        result.apply_request_changes(target=context.request)
-
     @staticmethod
     async def _commit_key_value_store_changes(
         result: RequestHandlerRunResult, get_kvs: GetKeyValueStoreFromRequestHandlerFunction
@@ -1413,12 +1410,12 @@ async def __run_task_function(self) -> None:
         else:
             session = await self._get_session()
         proxy_info = await self._get_proxy_info(request, session)
-        result = RequestHandlerRunResult(key_value_store_getter=self.get_key_value_store, request=request)
+        result = RequestHandlerRunResult(key_value_store_getter=self.get_key_value_store)
 
         deferred_cleanup: list[Callable[[], Awaitable[None]]] = []
 
         context = BasicCrawlingContext(
-            request=result.request,
+            request=request,
             session=session,
             proxy_info=proxy_info,
             send_request=self._prepare_send_request_function(session, proxy_info),
@@ -1437,9 +1434,8 @@ async def __run_task_function(self) -> None:
             request.state = RequestState.REQUEST_HANDLER
 
             try:
-                with swapped_context(context, request):
-                    self._check_request_collision(request, session)
-                    await self._run_request_handler(context=context)
+                self._check_request_collision(request, session)
+                await self._run_request_handler(context=context)
             except asyncio.TimeoutError as e:
                 raise RequestHandlerError(e, context) from e
 

tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py

@@ -267,8 +267,9 @@ async def pre_nav_hook(context: AdaptivePlaywrightPreNavCrawlingContext) -> None
             context.request.user_data['data'] = 'bs'
 
     await crawler.run(test_urls[:1])
-    # Check that repeated pre nav hook invocations do not influence each other while probing
-    assert user_data_in_pre_nav_hook == [None, None]
+    # First probe starts from the original user data. The second probe runs on the same request, so it
+    # sees the mutation from the first one.
+    assert user_data_in_pre_nav_hook == [None, 'pw']
     # Check that the request handler sees changes to user data done by pre nav hooks
     assert user_data_in_handler == ['pw', 'bs']
 
@@ -331,8 +332,9 @@ async def post_nav_hook(context: AdaptivePlaywrightPostNavCrawlingContext) -> No
             context.request.user_data['data'] = 'bs'
 
     await crawler.run(test_urls[:1])
-    # Check that repeated post nav hook invocations do not influence each other while probing
-    assert user_data_in_post_nav_hook == [None, None]
+    # First probe starts from the original user data. The second probe runs on the same request, so it
+    # sees the mutation from the first one.
+    assert user_data_in_post_nav_hook == [None, 'pw']
     # Check that the request handler sees changes to user data done by post nav hooks
     assert user_data_in_handler == ['pw', 'bs']
 
@@ -816,7 +818,7 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
 
 
 @pytest.mark.parametrize(
-    'test_input',
+    ('test_input', 'expected_request_state'),
     [
         pytest.param(
             TestInput(
@@ -825,6 +827,7 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
                 rendering_types=cycle(['static']),
                 detection_probability_recommendation=cycle([0]),
             ),
+            ['initial', 'static'],
             id='Static only',
         ),
         pytest.param(
@@ -834,6 +837,7 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
                 rendering_types=cycle(['client only']),
                 detection_probability_recommendation=cycle([0]),
             ),
+            ['initial', 'browser'],
             id='Client only',
         ),
         pytest.param(
@@ -843,11 +847,16 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
                 rendering_types=cycle(['static', 'client only']),
                 detection_probability_recommendation=cycle([1]),
             ),
+            # Both sub crawlers run on the same request: browser first, then the static detection probe.
+            # Mutations from both handlers persist.
+            ['initial', 'browser', 'static'],
             id='Enforced rendering type detection',
         ),
     ],
 )
-async def test_change_context_state_after_handling(test_input: TestInput, server_url: URL) -> None:
+async def test_change_context_state_after_handling(
+    test_input: TestInput, expected_request_state: list[str], server_url: URL
+) -> None:
     """Test that context state is saved after handling the request."""
     predictor = _SimpleRenderingTypePredictor(
         rendering_types=test_input.rendering_types,
@@ -872,8 +881,15 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
                 used_session_id = context.session.id
                 context.session.user_data['session_state'] = True
 
+            try:
+                # `page` is only available in the browser sub crawler, static crawling raises here.
+                _ = context.page
+                handler_type = 'browser'
+            except AdaptiveContextError:
+                handler_type = 'static'
+
             if isinstance(context.request.user_data['request_state'], list):
-                context.request.user_data['request_state'].append('handler')
+                context.request.user_data['request_state'].append(handler_type)
 
         request = Request.from_url(str(server_url), user_data={'request_state': ['initial']})
 
@@ -888,8 +904,8 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
         assert check_request is not None
 
         assert session.user_data.get('session_state') is True
-        # Check that request user data was updated in the handler and only onse.
-        assert check_request.user_data.get('request_state') == ['initial', 'handler']
+
+        assert check_request.user_data.get('request_state') == expected_request_state
 
         await request_queue.drop()
 

tests/unit/crawlers/_basic/test_basic_crawler.py

@@ -2111,27 +2111,57 @@ async def handler(_: BasicCrawlingContext) -> None:
     await crawler_task
 
 
-async def test_protect_request_in_run_handlers() -> None:
-    """Test that request in crawling context are protected in run handlers."""
+async def test_request_and_session_mutations_persist() -> None:
+    """Test that request and session mutated in the request handler and error handler, and mutations are persisted."""
     request_queue = await RequestQueue.open(name='state-test')
 
-    request = Request.from_url('https://test.url/', user_data={'request_state': ['initial']})
-
-    crawler = BasicCrawler(request_manager=request_queue, max_request_retries=0)
-
-    @crawler.router.default_handler
-    async def handler(context: BasicCrawlingContext) -> None:
-        if isinstance(context.request.user_data['request_state'], list):
-            context.request.user_data['request_state'].append('modified')
-        raise ValueError('Simulated error after modifying request')
+    async with SessionPool(max_pool_size=1) as session_pool:
+        session = await session_pool.get_session()
+        session.user_data['session_state'] = ['initial']
+        session.cookies['initial'] = 'yes'
+        request = Request.from_url('https://test.url/', user_data={'request_state': ['initial']}, session_id=session.id)
+
+        crawler = BasicCrawler(
+            request_manager=request_queue,
+            max_request_retries=1,
+            concurrency_settings=ConcurrencySettings(max_concurrency=1, desired_concurrency=1),
+            session_pool=session_pool,
+        )
 
-    await crawler.run([request])
+        @crawler.error_handler
+        async def error_handler(context: BasicCrawlingContext, error: Exception) -> Request | None:
+            if isinstance(context.request.user_data['request_state'], list):
+                context.request.user_data['request_state'].append('error')
 
-    check_request = await request_queue.get_request(request.unique_key)
-    assert check_request is not None
-    assert check_request.user_data['request_state'] == ['initial']
+            if context.session and isinstance(context.session.user_data['session_state'], list):
+                context.session.user_data['session_state'].append('error')
+                context.session.cookies['error'] = 'yes'
 
-    await request_queue.drop()
+        @crawler.router.default_handler
+        async def handler(context: BasicCrawlingContext) -> None:
+            if isinstance(context.request.user_data['request_state'], list):
+                context.request.user_data['request_state'].append(f'modified_{context.request.retry_count}')
+            if context.session and isinstance(context.session.user_data['session_state'], list):
+                context.session.user_data['session_state'].append(f'modified_{context.request.retry_count}')
+                context.session.cookies[f'modified_{context.request.retry_count}'] = 'yes'
+
+            if context.request.retry_count == 0:
+                raise ValueError('Simulated error after modifying request')
+
+        await crawler.run([request])
+
+        check_request = await request_queue.get_request(request.unique_key)
+        assert check_request is not None
+        assert check_request.user_data['request_state'] == ['initial', 'modified_0', 'error', 'modified_1']
+        session = await session_pool.get_session_by_id(session.id)
+        assert session is not None
+        assert session.user_data['session_state'] == ['initial', 'modified_0', 'error', 'modified_1']
+        assert session.cookies['initial'] == 'yes'
+        assert session.cookies['error'] == 'yes'
+        assert session.cookies['modified_0'] == 'yes'
+        assert session.cookies['modified_1'] == 'yes'
+
+        await request_queue.drop()
 
 
 async def test_new_request_error_handler() -> None: