Repository: apify/crawlee-python

Issue:
Extract state persistence into a mixin-style base class
As of now, only the `Statistics` class saves state on `persistState` events. We should make it possible to implement this behavior in other classes by simply extending some mixin and passing it a pydantic model that would validate the serialized state from the key-value store.

EDIT: `SessionPool` also supports persistence


Discussion:

Patch:
diff --git a/src/crawlee/_utils/recoverable_state.py b/src/crawlee/_utils/recoverable_state.py
new file mode 100644
index 0000000000..2cfdcd9ec7
--- /dev/null
+++ b/src/crawlee/_utils/recoverable_state.py
@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from pydantic import BaseModel
+
+from crawlee import service_locator
+from crawlee.events._types import Event, EventPersistStateData
+from crawlee.storages._key_value_store import KeyValueStore
+
+if TYPE_CHECKING:
+    import logging
+
+TStateModel = TypeVar('TStateModel', bound=BaseModel)
+
+
+class RecoverableState(Generic[TStateModel]):
+    """A class for managing persistent recoverable state using a Pydantic model.
+
+    This class facilitates state persistence to a `KeyValueStore`, allowing data to be saved and retrieved
+    across migrations or restarts. It manages the loading, saving, and resetting of state data,
+    with optional persistence capabilities.
+
+    The state is represented by a Pydantic model that can be serialized to and deserialized from JSON.
+    The class automatically hooks into the event system to persist state when needed.
+
+    Type Parameters:
+        TStateModel: A Pydantic BaseModel type that defines the structure of the state data.
+                     Typically, it should be inferred from the `default_state` constructor parameter.
+    """
+
+    def __init__(
+        self,
+        *,
+        default_state: TStateModel,
+        persist_state_key: str,
+        persistence_enabled: bool = False,
+        persist_state_kvs_name: str | None = None,
+        persist_state_kvs_id: str | None = None,
+        logger: logging.Logger,
+    ) -> None:
+        """Initialize a new recoverable state object.
+
+        Args:
+            default_state: The default state model instance to use when no persisted state is found.
+                           A deep copy is made each time the state is used.
+            persist_state_key: The key under which the state is stored in the KeyValueStore
+            persistence_enabled: Flag to enable or disable state persistence
+            persist_state_kvs_name: The name of the KeyValueStore to use for persistence.
+                                    If neither a name nor and id are supplied, the default store will be used.
+            persist_state_kvs_id: The identifier of the KeyValueStore to use for persistence.
+                                    If neither a name nor and id are supplied, the default store will be used.
+            logger: A logger instance for logging operations related to state persistence
+        """
+        self._default_state = default_state
+        self._state_type: type[TStateModel] = self._default_state.__class__
+        self._state: TStateModel | None = None
+        self._persistence_enabled = persistence_enabled
+        self._persist_state_key = persist_state_key
+        self._persist_state_kvs_name = persist_state_kvs_name
+        self._persist_state_kvs_id = persist_state_kvs_id
+        self._key_value_store: KeyValueStore | None = None
+        self._log = logger
+
+    async def initialize(self) -> TStateModel:
+        """Initialize the recoverable state.
+
+        This method must be called before using the recoverable state. It loads the saved state
+        if persistence is enabled and registers the object to listen for PERSIST_STATE events.
+
+        Returns:
+            The loaded state model
+        """
+        if not self._persistence_enabled:
+            self._state = self._default_state.model_copy(deep=True)
+            return self.current_value
+
+        self._key_value_store = await KeyValueStore.open(
+            name=self._persist_state_kvs_name, id=self._persist_state_kvs_id
+        )
+
+        await self._load_saved_state()
+
+        event_manager = service_locator.get_event_manager()
+        event_manager.on(event=Event.PERSIST_STATE, listener=self.persist_state)
+
+        return self.current_value
+
+    async def teardown(self) -> None:
+        """Clean up resources used by the recoverable state.
+
+        If persistence is enabled, this method deregisters the object from PERSIST_STATE events
+        and persists the current state one last time.
+        """
+        if not self._persistence_enabled:
+            return
+
+        event_manager = service_locator.get_event_manager()
+        event_manager.off(event=Event.PERSIST_STATE, listener=self.persist_state)
+        await self.persist_state()
+
+    @property
+    def current_value(self) -> TStateModel:
+        """Get the current state."""
+        if self._state is None:
+            raise RuntimeError('Recoverable state has not yet been loaded')
+
+        return self._state
+
+    async def reset(self) -> None:
+        """Reset the state to the default values and clear any persisted state.
+
+        Resets the current state to the default state and, if persistence is enabled,
+        clears the persisted state from the KeyValueStore.
+        """
+        self._state = self._default_state.model_copy(deep=True)
+
+        if self._persistence_enabled:
+            if self._key_value_store is None:
+                raise RuntimeError('Recoverable state has not yet been initialized')
+
+            await self._key_value_store.set_value(self._persist_state_key, None)
+
+    async def persist_state(self, event_data: EventPersistStateData | None = None) -> None:
+        """Persist the current state to the KeyValueStore.
+
+        This method is typically called in response to a PERSIST_STATE event, but can also be called
+        directly when needed.
+
+        Args:
+            event_data: Optional data associated with a PERSIST_STATE event
+        """
+        self._log.debug(f'Persisting state of the Statistics (event_data={event_data}).')
+
+        if self._key_value_store is None or self._state is None:
+            raise RuntimeError('Recoverable state has not yet been initialized')
+
+        if self._persistence_enabled:
+            await self._key_value_store.set_value(
+                self._persist_state_key,
+                self._state.model_dump(mode='json', by_alias=True),
+                'application/json',
+            )
+
+    async def _load_saved_state(self) -> None:
+        if self._key_value_store is None:
+            raise RuntimeError('Recoverable state has not yet been initialized')
+
+        stored_state = await self._key_value_store.get_value(self._persist_state_key)
+        if stored_state is None:
+            self._state = self._default_state.model_copy(deep=True)
+        else:
+            self._state = self._state_type.model_validate(stored_state)
diff --git a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py
index 99315f8d1e..9a9970eba4 100644
--- a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py
+++ b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py
@@ -67,10 +67,11 @@ class _NonPersistentStatistics(Statistics):
 
     def __init__(self) -> None:
         super().__init__(state_model=StatisticsState)
-        self._active = True
 
     async def __aenter__(self) -> Self:
         self._active = True
+        await self._state.initialize()
+        self._after_initialize()
         return self
 
     async def __aexit__(
@@ -201,7 +202,12 @@ async def adaptive_pre_navigation_hook_pw(context: PlaywrightPreNavCrawlingConte
         static_crawler.pre_navigation_hook(adaptive_pre_navigation_hook_static)
         playwright_crawler.pre_navigation_hook(adaptive_pre_navigation_hook_pw)
 
-        self._additional_context_managers = [*self._additional_context_managers, playwright_crawler._browser_pool]  # noqa: SLF001 # Intentional access to private member.
+        self._additional_context_managers = [
+            *self._additional_context_managers,
+            static_crawler.statistics,
+            playwright_crawler.statistics,
+            playwright_crawler._browser_pool,  # noqa: SLF001 # Intentional access to private member.
+        ]
 
         # Sub crawler pipeline related
         self._pw_context_pipeline = playwright_crawler._context_pipeline  # noqa:SLF001  # Intentional access to private member.
@@ -376,6 +382,8 @@ async def _run_request_handler(self, context: BasicCrawlingContext) -> None:
 
         context.log.debug(f'Running browser request handler for {context.request.url}')
 
+        old_state_copy = None
+
         if should_detect_rendering_type:
             # Save copy of global state from `use_state` before it can be mutated by browser crawl.
             # This copy will be used in the static crawl to make sure they both run with same conditions and to
diff --git a/src/crawlee/fingerprint_suite/_types.py b/src/crawlee/fingerprint_suite/_types.py
index 674cfa4eec..276199e641 100644
--- a/src/crawlee/fingerprint_suite/_types.py
+++ b/src/crawlee/fingerprint_suite/_types.py
@@ -2,7 +2,7 @@
 
 from typing import Annotated, Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from crawlee._utils.docs import docs_group
 
@@ -14,6 +14,8 @@
 
 @docs_group('Data structures')
 class ScreenOptions(BaseModel):
+    model_config = ConfigDict(extra='forbid', populate_by_name=True)
+
     """Defines the screen constrains for the fingerprint generator."""
 
     min_width: Annotated[float | None, Field(alias='minWidth')] = None
@@ -28,15 +30,13 @@ class ScreenOptions(BaseModel):
     max_height: Annotated[float | None, Field(alias='maxHeight')] = None
     """Maximal screen height constraint for the fingerprint generator."""
 
-    class Config:
-        extra = 'forbid'
-        populate_by_name = True
-
 
 @docs_group('Data structures')
 class HeaderGeneratorOptions(BaseModel):
     """Collection of header related attributes that can be used by the fingerprint generator."""
 
+    model_config = ConfigDict(extra='forbid', populate_by_name=True)
+
     browsers: list[SupportedBrowserType] | None = None
     """List of BrowserSpecifications to generate the headers for."""
 
@@ -56,7 +56,3 @@ class HeaderGeneratorOptions(BaseModel):
 
     strict: bool | None = None
     """If true, the generator will throw an error if it cannot generate headers based on the input."""
-
-    class Config:
-        extra = 'forbid'
-        populate_by_name = True
diff --git a/src/crawlee/sessions/_models.py b/src/crawlee/sessions/_models.py
index 1e28ddd2ce..2d7ac00b05 100644
--- a/src/crawlee/sessions/_models.py
+++ b/src/crawlee/sessions/_models.py
@@ -1,11 +1,20 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
-from typing import Annotated
+from typing import Annotated, Any
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    GetPydanticSchema,
+    PlainSerializer,
+    computed_field,
+)
 
 from ._cookies import CookieParam
+from ._session import Session
 
 
 class SessionModel(BaseModel):
@@ -31,11 +40,43 @@ class SessionPoolModel(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True)
 
-    persistence_enabled: Annotated[bool, Field(alias='persistenceEnabled')]
-    persist_state_kvs_name: Annotated[str | None, Field(alias='persistStateKvsName')]
-    persist_state_key: Annotated[str, Field(alias='persistStateKey')]
     max_pool_size: Annotated[int, Field(alias='maxPoolSize')]
-    session_count: Annotated[int, Field(alias='sessionCount')]
-    usable_session_count: Annotated[int, Field(alias='usableSessionCount')]
-    retired_session_count: Annotated[int, Field(alias='retiredSessionCount')]
-    sessions: Annotated[list[SessionModel], Field(alias='sessions')]
+
+    sessions: Annotated[
+        dict[
+            str,
+            Annotated[
+                Session, GetPydanticSchema(lambda _, handler: handler(Any))
+            ],  # handler(Any) is fine - we validate manually in the BeforeValidator
+        ],
+        Field(alias='sessions'),
+        PlainSerializer(
+            lambda value: [session.get_state().model_dump(by_alias=True) for session in value.values()],
+            return_type=list,
+        ),
+        BeforeValidator(
+            lambda value: {
+                session.id: session
+                for item in value
+                if (session := Session.from_model(SessionModel.model_validate(item, by_alias=True)))
+            }
+        ),
+    ]
+
+    @computed_field(alias='sessionCount')  # type: ignore[prop-decorator]
+    @property
+    def session_count(self) -> int:
+        """Get the total number of sessions currently maintained in the pool."""
+        return len(self.sessions)
+
+    @computed_field(alias='usableSessionCount')  # type: ignore[prop-decorator]
+    @property
+    def usable_session_count(self) -> int:
+        """Get the number of sessions that are currently usable."""
+        return len([session for _, session in self.sessions.items() if session.is_usable])
+
+    @computed_field(alias='retiredSessionCount')  # type: ignore[prop-decorator]
+    @property
+    def retired_session_count(self) -> int:
+        """Get the number of sessions that are no longer usable."""
+        return self.session_count - self.usable_session_count
diff --git a/src/crawlee/sessions/_session.py b/src/crawlee/sessions/_session.py
index e92cfee030..13113cc1b4 100644
--- a/src/crawlee/sessions/_session.py
+++ b/src/crawlee/sessions/_session.py
@@ -9,11 +9,12 @@
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee._utils.docs import docs_group
 from crawlee.sessions._cookies import CookieParam, SessionCookies
-from crawlee.sessions._models import SessionModel
 
 if TYPE_CHECKING:
     from http.cookiejar import CookieJar
 
+    from crawlee.sessions._models import SessionModel
+
 logger = getLogger(__name__)
 
 
@@ -146,6 +147,8 @@ def get_state(self, *, as_dict: Literal[False]) -> SessionModel: ...
 
     def get_state(self, *, as_dict: bool = False) -> SessionModel | dict:
         """Retrieve the current state of the session either as a model or as a dictionary."""
+        from ._models import SessionModel
+
         model = SessionModel(
             id=self._id,
             max_age=self._max_age,
@@ -157,7 +160,7 @@ def get_state(self, *, as_dict: bool = False) -> SessionModel | dict:
             max_usage_count=self._max_usage_count,
             error_score=self._error_score,
             cookies=self._cookies.get_cookies_as_dicts(),
-            blocked_status_codes=self._blocked_status_codes,
+            blocked_status_codes=list(self._blocked_status_codes),
         )
         if as_dict:
             return model.model_dump()
diff --git a/src/crawlee/sessions/_session_pool.py b/src/crawlee/sessions/_session_pool.py
index 1f6409a37e..714b1c0b3e 100644
--- a/src/crawlee/sessions/_session_pool.py
+++ b/src/crawlee/sessions/_session_pool.py
@@ -9,10 +9,9 @@
 from crawlee import service_locator
 from crawlee._utils.context import ensure_context
 from crawlee._utils.docs import docs_group
-from crawlee.events._types import Event, EventPersistStateData
+from crawlee._utils.recoverable_state import RecoverableState
 from crawlee.sessions import Session
 from crawlee.sessions._models import SessionPoolModel
-from crawlee.storages import KeyValueStore
 
 if TYPE_CHECKING:
     from types import TracebackType
@@ -54,28 +53,32 @@ def __init__(
             create_session_function: A callable to create new session instances. If None, a default session settings
                 will be used. Do not set it if you are providing `create_session_settings`.
             event_manager: The event manager to handle events like persist state.
-            persistence_enabled: Flag to enable or disable state persistence of the pool. If it is enabled, make sure
-                to provide an event manager to handle the events.
+            persistence_enabled: Flag to enable or disable state persistence of the pool.
             persist_state_kvs_name: The name of the `KeyValueStore` used for state persistence.
             persist_state_key: The key under which the session pool's state is stored in the `KeyValueStore`.
         """
         if event_manager:
             service_locator.set_event_manager(event_manager)
 
+        self._state = RecoverableState(
+            default_state=SessionPoolModel(
+                max_pool_size=max_pool_size,
+                sessions={},
+            ),
+            logger=logger,
+            persistence_enabled=persistence_enabled,
+            persist_state_kvs_name=persist_state_kvs_name,
+            persist_state_key=persist_state_key or 'CRAWLEE_SESSION_POOL_STATE',
+        )
+
         self._max_pool_size = max_pool_size
         self._session_settings = create_session_settings or {}
         self._create_session_function = create_session_function
         self._persistence_enabled = persistence_enabled
-        self._persist_state_kvs_name = persist_state_kvs_name
-        self._persist_state_key = persist_state_key
 
         if self._create_session_function and self._session_settings:
             raise ValueError('Both `create_session_settings` and `create_session_function` cannot be provided.')
 
-        # Internal non-configurable attributes
-        self._kvs: KeyValueStore | None = None
-        self._sessions: dict[str, Session] = {}
-
         # Flag to indicate the context state.
         self._active = False
 
@@ -86,17 +89,17 @@ def __repr__(self) -> str:
     @property
     def session_count(self) -> int:
         """Get the total number of sessions currently maintained in the pool."""
-        return len(self._sessions)
+        return len(self._state.current_value.sessions)
 
     @property
     def usable_session_count(self) -> int:
         """Get the number of sessions that are currently usable."""
-        return len([session for _, session in self._sessions.items() if session.is_usable])
+        return self._state.current_value.usable_session_count
 
     @property
     def retired_session_count(self) -> int:
         """Get the number of sessions that are no longer usable."""
-        return self.session_count - self.usable_session_count
+        return self._state.current_value.retired_session_count
 
     @property
     def active(self) -> bool:
@@ -114,21 +117,11 @@ async def __aenter__(self) -> SessionPool:
 
         self._active = True
 
-        if self._persistence_enabled:
-            event_manager = service_locator.get_event_manager()
-            self._kvs = await KeyValueStore.open(name=self._persist_state_kvs_name)
-
-            # Attempt to restore the previously persisted state.
-            was_restored = await self._try_to_restore_previous_state()
-
-            # If the pool could not be restored, initialize it with new sessions.
-            if not was_restored:
-                await self._fill_sessions_to_max()
+        state = await self._state.initialize()
+        state.max_pool_size = self._max_pool_size
+        self._remove_retired_sessions()
 
-            # Register an event listener for persisting the session pool state.
-            event_manager.on(event=Event.PERSIST_STATE, listener=self._persist_state)
-        # If persistence is disabled, just fill the pool with sessions.
-        else:
+        if not state.sessions:
             await self._fill_sessions_to_max()
 
         return self
@@ -147,13 +140,7 @@ async def __aexit__(
         if not self._active:
             raise RuntimeError(f'The {self.__class__.__name__} is not active.')
 
-        if self._persistence_enabled:
-            event_manager = service_locator.get_event_manager()
-            # Remove the event listener for state persistence.
-            event_manager.off(event=Event.PERSIST_STATE, listener=self._persist_state)
-
-            # Persist the final state of the session pool.
-            await self._persist_state(event_data=EventPersistStateData(is_migrating=False))
+        await self._state.teardown()
 
         self._active = False
 
@@ -166,16 +153,7 @@ def get_state(self, *, as_dict: Literal[False]) -> SessionPoolModel: ...
     @ensure_context
     def get_state(self, *, as_dict: bool = False) -> SessionPoolModel | dict:
         """Retrieve the current state of the pool either as a model or as a dictionary."""
-        model = SessionPoolModel(
-            persistence_enabled=self._persistence_enabled,
-            persist_state_kvs_name=self._persist_state_kvs_name,
-            persist_state_key=self._persist_state_key,
-            max_pool_size=self._max_pool_size,
-            session_count=self.session_count,
-            usable_session_count=self.usable_session_count,
-            retired_session_count=self.retired_session_count,
-            sessions=[session.get_state(as_dict=False) for _, session in self._sessions.items()],
-        )
+        model = self._state.current_value.model_copy(deep=True)
         if as_dict:
             return model.model_dump()
         return model
@@ -190,10 +168,12 @@ def add_session(self, session: Session) -> None:
         Args:
             session: The session to add to the pool.
         """
-        if session.id in self._sessions:
+        state = self._state.current_value
+
+        if session.id in state.sessions:
             logger.warning(f'Session with ID {session.id} already exists in the pool.')
             return
-        self._sessions[session.id] = session
+        state.sessions[session.id] = session
 
     @ensure_context
     async def get_session(self) -> Session:
@@ -229,7 +209,7 @@ async def get_session_by_id(self, session_id: str) -> Session | None:
             The session object if found and usable, otherwise `None`.
         """
         await self._fill_sessions_to_max()
-        session = self._sessions.get(session_id)
+        session = self._state.current_value.sessions.get(session_id)
 
         if not session:
             logger.warning(f'Session with ID {session_id} not found.')
@@ -243,15 +223,7 @@ async def get_session_by_id(self, session_id: str) -> Session | None:
 
     async def reset_store(self) -> None:
         """Reset the KVS where the pool state is persisted."""
-        if not self._persistence_enabled:
-            logger.debug('Persistence is disabled; skipping the reset of the store.')
-            return
-
-        if not self._kvs:
-            logger.warning('SessionPool reset failed: KVS not initialized. Did you forget to call __aenter__?')
-            return
-
-        await self._kvs.set_value(key=self._persist_state_key, value=None)
+        await self._state.reset()
 
     async def _create_new_session(self) -> Session:
         """Create a new session, add it to the pool and return it."""
@@ -259,7 +231,7 @@ async def _create_new_session(self) -> Session:
             new_session = self._create_session_function()
         else:
             new_session = Session(**self._session_settings)
-        self._sessions[new_session.id] = new_session
+        self._state.current_value.sessions[new_session.id] = new_session
         return new_session
 
     async def _fill_sessions_to_max(self) -> None:
@@ -269,50 +241,12 @@ async def _fill_sessions_to_max(self) -> None:
 
     def _get_random_session(self) -> Session:
         """Get a random session from the pool."""
-        keys = list(self._sessions.keys())
-        if not keys:
+        state = self._state.current_value
+        if not state.sessions:
             raise ValueError('No sessions available in the pool.')
-        key = random.choice(keys)
-        return self._sessions[key]
+        return random.choice(list(state.sessions.values()))
 
     def _remove_retired_sessions(self) -> None:
         """Remove all sessions from the pool that are no longer usable."""
-        self._sessions = {session_id: session for session_id, session in self._sessions.items() if session.is_usable}
-
-    async def _try_to_restore_previous_state(self) -> bool:
-        """Try to restore the previous state of the pool from the KVS."""
-        if not self._persistence_enabled:
-            logger.warning('Persistence is disabled, however, the state restoration was triggered.')
-
-        if not self._kvs:
-            logger.warning('SessionPool restoration failed: KVS not initialized. Did you forget to call __aenter__?')
-            return False
-
-        previous_state = await self._kvs.get_value(key=self._persist_state_key)
-
-        if previous_state is None:
-            logger.debug('SessionPool restoration skipped: No previous state found.')
-            return False
-
-        previous_session_pool = SessionPoolModel.model_validate(previous_state)
-
-        for session_model in previous_session_pool.sessions:
-            session = Session.from_model(model=session_model)
-            if session.is_usable:
-                self._sessions[session.id] = session
-
-        return True
-
-    async def _persist_state(self, event_data: EventPersistStateData) -> None:
-        """Persist the state of the pool in the KVS."""
-        logger.debug(f'Persisting state of the SessionPool (event_data={event_data}).')
-
-        if not self._persistence_enabled:
-            logger.warning('Persistence is disabled, however, the state persistence event was triggered.')
-
-        if not self._kvs:
-            logger.warning('SessionPool persisting failed: KVS not initialized. Did you forget to call __aenter__?')
-            return
-
-        session_pool_state = self.get_state(as_dict=True)
-        await self._kvs.set_value(key=self._persist_state_key, value=session_pool_state)
+        state = self._state.current_value
+        state.sessions = {session.id: session for session in state.sessions.values() if session.is_usable}
diff --git a/src/crawlee/statistics/__init__.py b/src/crawlee/statistics/__init__.py
index 769918a50c..ba9466f872 100644
--- a/src/crawlee/statistics/__init__.py
+++ b/src/crawlee/statistics/__init__.py
@@ -1,6 +1,6 @@
 # ruff: noqa: A005
 
-from ._models import FinalStatistics, StatisticsPersistedState, StatisticsState
+from ._models import FinalStatistics, StatisticsState
 from ._statistics import Statistics
 
-__all__ = ['FinalStatistics', 'Statistics', 'StatisticsPersistedState', 'StatisticsState']
+__all__ = ['FinalStatistics', 'Statistics', 'StatisticsState']
diff --git a/src/crawlee/statistics/_error_snapshotter.py b/src/crawlee/statistics/_error_snapshotter.py
index 52fbb30b18..21dbd33d48 100644
--- a/src/crawlee/statistics/_error_snapshotter.py
+++ b/src/crawlee/statistics/_error_snapshotter.py
@@ -6,9 +6,10 @@
 import string
 from typing import TYPE_CHECKING
 
+from crawlee.storages import KeyValueStore
+
 if TYPE_CHECKING:
     from crawlee._types import BasicCrawlingContext
-    from crawlee.storages import KeyValueStore
 
 
 class ErrorSnapshotter:
@@ -19,18 +20,8 @@ class ErrorSnapshotter:
     SNAPSHOT_PREFIX = 'ERROR_SNAPSHOT'
     ALLOWED_CHARACTERS = string.ascii_letters + string.digits + '!-_.'
 
-    def __init__(self, kvs: KeyValueStore | None = None) -> None:
-        self._kvs = kvs
-
-    @property
-    def kvs(self) -> KeyValueStore:
-        if not self._kvs:
-            raise RuntimeError('The key value store was not yet set.')
-        return self._kvs
-
-    @kvs.setter
-    def kvs(self, kvs: KeyValueStore) -> None:
-        self._kvs = kvs
+    def __init__(self, *, snapshot_kvs_name: str | None = None) -> None:
+        self._kvs_name = snapshot_kvs_name
 
     async def capture_snapshot(self, error_message: str, file_and_line: str, context: BasicCrawlingContext) -> None:
         """Capture error snapshot and save it to key value store.
@@ -60,11 +51,13 @@ async def capture_snapshot(self, error_message: str, file_and_line: str, context
 
     async def _save_html(self, html: str, base_name: str) -> None:
         file_name = f'{base_name}.html'
-        await self.kvs.set_value(file_name, html, content_type='text/html')
+        kvs = await KeyValueStore.open(name=self._kvs_name)
+        await kvs.set_value(file_name, html, content_type='text/html')
 
     async def _save_screenshot(self, screenshot: bytes, base_name: str) -> None:
         file_name = f'{base_name}.jpg'
-        await self.kvs.set_value(file_name, screenshot, content_type='image/jpeg')
+        kvs = await KeyValueStore.open(name=self._kvs_name)
+        await kvs.set_value(file_name, screenshot, content_type='image/jpeg')
 
     def _sanitize_filename(self, filename: str) -> str:
         return re.sub(f'[^{re.escape(self.ALLOWED_CHARACTERS)}]', '', filename[: self.MAX_FILENAME_LENGTH])
diff --git a/src/crawlee/statistics/_error_tracker.py b/src/crawlee/statistics/_error_tracker.py
index 4e70aaa2b7..ac17c24b0c 100644
--- a/src/crawlee/statistics/_error_tracker.py
+++ b/src/crawlee/statistics/_error_tracker.py
@@ -26,13 +26,14 @@ class ErrorTracker:
     def __init__(
         self,
         *,
+        snapshot_kvs_name: str | None = None,
         show_error_name: bool = True,
         show_file_and_line_number: bool = True,
         show_error_message: bool = True,
         show_full_message: bool = False,
         save_error_snapshots: bool = False,
     ) -> None:
-        self.error_snapshotter = ErrorSnapshotter() if save_error_snapshots else None
+        self.error_snapshotter = ErrorSnapshotter(snapshot_kvs_name=snapshot_kvs_name) if save_error_snapshots else None
         self.show_error_name = show_error_name
         self.show_file_and_line_number = show_file_and_line_number
         self.show_error_message = show_error_message
diff --git a/src/crawlee/statistics/_models.py b/src/crawlee/statistics/_models.py
index 391a4a29c8..e8d91a68d2 100644
--- a/src/crawlee/statistics/_models.py
+++ b/src/crawlee/statistics/_models.py
@@ -2,10 +2,10 @@
 
 import json
 from dataclasses import asdict, dataclass
-from datetime import datetime, timedelta
-from typing import Annotated, Any
+from datetime import datetime, timedelta, timezone
+from typing import Annotated, Any, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator, computed_field
 from typing_extensions import override
 
 from crawlee._utils.console import make_table
@@ -50,6 +50,7 @@ class StatisticsState(BaseModel):
     """Statistic data about a crawler run."""
 
     model_config = ConfigDict(populate_by_name=True, ser_json_inf_nan='constants')
+    stats_id: Annotated[int | None, Field(alias='statsId')] = None
 
     requests_finished: Annotated[int, Field(alias='requestsFinished')] = 0
     requests_failed: Annotated[int, Field(alias='requestsFailed')] = 0
@@ -71,20 +72,35 @@ class StatisticsState(BaseModel):
     errors: dict[str, Any] = Field(default_factory=dict)
     retry_errors: dict[str, Any] = Field(alias='retryErrors', default_factory=dict)
     requests_with_status_code: dict[str, int] = Field(alias='requestsWithStatusCode', default_factory=dict)
-    stats_persisted_at: Annotated[datetime | None, Field(alias='statsPersistedAt')] = None
-
-
-@docs_group('Data structures')
-class StatisticsPersistedState(BaseModel):
-    """Additional statistic data to be stored in the persisted state."""
-
-    model_config = ConfigDict(populate_by_name=True)
-
-    request_retry_histogram: Annotated[list[int], Field(alias='requestRetryHistogram')]
-    stats_id: Annotated[int, Field(alias='statsId')]
-    request_avg_failed_duration: Annotated[timedelta_ms | None, Field(alias='requestAvgFailedDurationMillis')]
-    request_avg_finished_duration: Annotated[timedelta_ms | None, Field(alias='requestAvgFinishedDurationMillis')]
-    request_total_duration: Annotated[timedelta_ms, Field(alias='requestTotalDurationMillis')]
-    requests_total: Annotated[int, Field(alias='requestsTotal')]
-    crawler_last_started_at: Annotated[datetime, Field(alias='crawlerLastStartTimestamp')]
-    stats_persisted_at: Annotated[datetime, Field(alias='statsPersistedAt')]
+    stats_persisted_at: Annotated[
+        datetime | None, Field(alias='statsPersistedAt'), PlainSerializer(lambda _: datetime.now(timezone.utc))
+    ] = None
+    request_retry_histogram: Annotated[
+        dict[int, int],
+        Field(alias='requestRetryHistogram'),
+        PlainValidator(lambda value: dict(enumerate(value)), json_schema_input_type=list[int]),
+        PlainSerializer(
+            lambda value: [value.get(i, 0) for i in range(max(value.keys(), default=0) + 1)],
+            return_type=list[int],
+        ),
+    ] = {}
+
+    @computed_field(alias='requestTotalDurationMillis', return_type=timedelta_ms)  # type: ignore[prop-decorator]
+    @property
+    def request_total_duration(self) -> timedelta:
+        return self.request_total_finished_duration + self.request_total_failed_duration
+
+    @computed_field(alias='requestAvgFailedDurationMillis', return_type=Optional[timedelta_ms])  # type: ignore[prop-decorator]
+    @property
+    def request_avg_failed_duration(self) -> timedelta | None:
+        return (self.request_total_failed_duration / self.requests_failed) if self.requests_failed else None
+
+    @computed_field(alias='requestAvgFinishedDurationMillis', return_type=Optional[timedelta_ms])  # type: ignore[prop-decorator]
+    @property
+    def request_avg_finished_duration(self) -> timedelta | None:
+        return (self.request_total_finished_duration / self.requests_finished) if self.requests_finished else None
+
+    @computed_field(alias='requestsTotal')  # type: ignore[prop-decorator]
+    @property
+    def requests_total(self) -> int:
+        return self.requests_failed + self.requests_finished
diff --git a/src/crawlee/statistics/_statistics.py b/src/crawlee/statistics/_statistics.py
index d07f379d7a..3d95a889f7 100644
--- a/src/crawlee/statistics/_statistics.py
+++ b/src/crawlee/statistics/_statistics.py
@@ -4,18 +4,16 @@
 import math
 from datetime import datetime, timedelta, timezone
 from logging import Logger, getLogger
-from typing import TYPE_CHECKING, Any, Generic, Literal, cast
+from typing import TYPE_CHECKING, Generic, Literal
 
 from typing_extensions import Self, TypeVar
 
-from crawlee import service_locator
 from crawlee._utils.context import ensure_context
 from crawlee._utils.docs import docs_group
+from crawlee._utils.recoverable_state import RecoverableState
 from crawlee._utils.recurring_task import RecurringTask
-from crawlee.events._types import Event, EventPersistStateData
-from crawlee.statistics import FinalStatistics, StatisticsPersistedState, StatisticsState
+from crawlee.statistics import FinalStatistics, StatisticsState
 from crawlee.statistics._error_tracker import ErrorTracker
-from crawlee.storages import KeyValueStore
 
 if TYPE_CHECKING:
     from types import TracebackType
@@ -71,7 +69,6 @@ def __init__(
         persistence_enabled: bool = False,
         persist_state_kvs_name: str | None = None,
         persist_state_key: str | None = None,
-        key_value_store: KeyValueStore | None = None,
         log_message: str = 'Statistics',
         periodic_message_logger: Logger | None = None,
         log_interval: timedelta = timedelta(minutes=1),
@@ -82,23 +79,23 @@ def __init__(
         self._id = Statistics.__next_id
         Statistics.__next_id += 1
 
-        self._state_model = state_model
-        self.state = self._state_model()
         self._instance_start: datetime | None = None
-        self._retry_histogram = dict[int, int]()
 
-        self.error_tracker = ErrorTracker(save_error_snapshots=save_error_snapshots)
+        self.error_tracker = ErrorTracker(
+            save_error_snapshots=save_error_snapshots,
+            snapshot_kvs_name=persist_state_kvs_name,
+        )
         self.error_tracker_retry = ErrorTracker(save_error_snapshots=False)
 
         self._requests_in_progress = dict[str, RequestProcessingRecord]()
 
-        if persist_state_key is None:
-            persist_state_key = f'SDK_CRAWLER_STATISTICS_{self._id}'
-
-        self._persistence_enabled = persistence_enabled
-        self._persist_state_key = persist_state_key
-        self._persist_state_kvs_name = persist_state_kvs_name
-        self._key_value_store: KeyValueStore | None = key_value_store
+        self._state = RecoverableState(
+            default_state=state_model(stats_id=self._id),
+            persist_state_key=persist_state_key or f'SDK_CRAWLER_STATISTICS_{self._id}',
+            persistence_enabled=persistence_enabled,
+            persist_state_kvs_name=persist_state_kvs_name,
+            logger=logger,
+        )
 
         self._log_message = log_message
         self._statistics_log_format = statistics_log_format
@@ -111,10 +108,9 @@ def __init__(
     def replace_state_model(self, state_model: type[TNewStatisticsState]) -> Statistics[TNewStatisticsState]:
         """Create near copy of the `Statistics` with replaced `state_model`."""
         new_statistics: Statistics[TNewStatisticsState] = Statistics(
-            persistence_enabled=self._persistence_enabled,
-            persist_state_kvs_name=self._persist_state_kvs_name,
-            persist_state_key=self._persist_state_key,
-            key_value_store=self._key_value_store,
+            persistence_enabled=self._state._persistence_enabled,  # noqa: SLF001
+            persist_state_kvs_name=self._state._persist_state_kvs_name,  # noqa: SLF001
+            persist_state_key=self._state._persist_state_key,  # noqa: SLF001
             log_message=self._log_message,
             periodic_message_logger=self._periodic_message_logger,
             state_model=state_model,
@@ -128,7 +124,6 @@ def with_default_state(
         persistence_enabled: bool = False,
         persist_state_kvs_name: str | None = None,
         persist_state_key: str | None = None,
-        key_value_store: KeyValueStore | None = None,
         log_message: str = 'Statistics',
         periodic_message_logger: Logger | None = None,
         log_interval: timedelta = timedelta(minutes=1),
@@ -140,7 +135,6 @@ def with_default_state(
             persistence_enabled=persistence_enabled,
             persist_state_kvs_name=persist_state_kvs_name,
             persist_state_key=persist_state_key,
-            key_value_store=key_value_store,
             log_message=log_message,
             periodic_message_logger=periodic_message_logger,
             log_interval=log_interval,
@@ -166,18 +160,9 @@ async def __aenter__(self) -> Self:
         self._active = True
         self._instance_start = datetime.now(timezone.utc)
 
-        if self.state.crawler_started_at is None:
-            self.state.crawler_started_at = datetime.now(timezone.utc)
-
-        if self._key_value_store is None:
-            self._key_value_store = await KeyValueStore.open(name=self._persist_state_kvs_name)
-
-        if self.error_tracker.error_snapshotter:
-            self.error_tracker.error_snapshotter.kvs = self._key_value_store
+        await self._state.initialize()
+        self._after_initialize()
 
-        await self._maybe_load_statistics()
-        event_manager = service_locator.get_event_manager()
-        event_manager.on(event=Event.PERSIST_STATE, listener=self._persist_state)
         self._periodic_logger.start()
 
         return self
@@ -196,18 +181,24 @@ async def __aexit__(
         if not self._active:
             raise RuntimeError(f'The {self.__class__.__name__} is not active.')
 
-        self.state.crawler_finished_at = datetime.now(timezone.utc)
-        event_manager = service_locator.get_event_manager()
-        event_manager.off(event=Event.PERSIST_STATE, listener=self._persist_state)
+        self._state.current_value.crawler_finished_at = datetime.now(timezone.utc)
+
+        await self._state.teardown()
+
         await self._periodic_logger.stop()
-        await self._persist_state(event_data=EventPersistStateData(is_migrating=False))
+
         self._active = False
 
+    @property
+    def state(self) -> TStatisticsState:
+        return self._state.current_value
+
     @ensure_context
     def register_status_code(self, code: int) -> None:
         """Increment the number of times a status code has been received."""
-        self.state.requests_with_status_code.setdefault(str(code), 0)
-        self.state.requests_with_status_code[str(code)] += 1
+        state = self._state.current_value
+        state.requests_with_status_code.setdefault(str(code), 0)
+        state.requests_with_status_code[str(code)] += 1
 
     @ensure_context
     def record_request_processing_start(self, request_id_or_key: str) -> None:
@@ -223,15 +214,17 @@ def record_request_processing_finish(self, request_id_or_key: str) -> None:
         if record is None:
             return
 
+        state = self._state.current_value
         duration = record.finish()
-        self.state.requests_finished += 1
-        self.state.request_total_finished_duration += duration
+
+        state.requests_finished += 1
+        state.request_total_finished_duration += duration
         self._save_retry_count_for_request(record)
-        self.state.request_min_duration = min(
-            self.state.request_min_duration if self.state.request_min_duration is not None else timedelta.max, duration
+        state.request_min_duration = min(
+            state.request_min_duration if state.request_min_duration is not None else timedelta.max, duration
         )
-        self.state.request_max_duration = min(
-            self.state.request_max_duration if self.state.request_max_duration is not None else timedelta(), duration
+        state.request_max_duration = min(
+            state.request_max_duration if state.request_max_duration is not None else timedelta(), duration
         )
 
         del self._requests_in_progress[request_id_or_key]
@@ -243,8 +236,10 @@ def record_request_processing_failure(self, request_id_or_key: str) -> None:
         if record is None:
             return
 
-        self.state.request_total_failed_duration += record.finish()
-        self.state.requests_failed += 1
+        state = self._state.current_value
+
+        state.request_total_failed_duration += record.finish()
+        state.requests_failed += 1
         self._save_retry_count_for_request(record)
 
         del self._requests_in_progress[request_id_or_key]
@@ -256,39 +251,29 @@ def calculate(self) -> FinalStatistics:
 
         crawler_runtime = datetime.now(timezone.utc) - self._instance_start
         total_minutes = crawler_runtime.total_seconds() / 60
+        state = self._state.current_value
+        serialized_state = state.model_dump(by_alias=False)
 
         return FinalStatistics(
-            request_avg_failed_duration=(self.state.request_total_failed_duration / self.state.requests_failed)
-            if self.state.requests_failed
-            else None,
-            request_avg_finished_duration=(self.state.request_total_finished_duration / self.state.requests_finished)
-            if self.state.requests_finished
-            else None,
-            requests_finished_per_minute=round(self.state.requests_finished / total_minutes) if total_minutes else 0,
-            requests_failed_per_minute=math.floor(self.state.requests_failed / total_minutes) if total_minutes else 0,
-            request_total_duration=self.state.request_total_finished_duration
-            + self.state.request_total_failed_duration,
-            requests_total=self.state.requests_failed + self.state.requests_finished,
+            request_avg_failed_duration=state.request_avg_failed_duration,
+            request_avg_finished_duration=state.request_avg_finished_duration,
+            requests_finished_per_minute=round(state.requests_finished / total_minutes) if total_minutes else 0,
+            requests_failed_per_minute=math.floor(state.requests_failed / total_minutes) if total_minutes else 0,
+            request_total_duration=state.request_total_finished_duration + state.request_total_failed_duration,
+            requests_total=state.requests_failed + state.requests_finished,
             crawler_runtime=crawler_runtime,
-            requests_finished=self.state.requests_finished,
-            requests_failed=self.state.requests_failed,
-            retry_histogram=[
-                self._retry_histogram.get(retry_count, 0)
-                for retry_count in range(max(self._retry_histogram.keys(), default=0) + 1)
-            ],
+            requests_finished=state.requests_finished,
+            requests_failed=state.requests_failed,
+            retry_histogram=serialized_state['request_retry_histogram'],
         )
 
     async def reset(self) -> None:
         """Reset the statistics to their defaults and remove any persistent state."""
-        self.state = self._state_model()
+        await self._state.reset()
         self.error_tracker = ErrorTracker()
         self.error_tracker_retry = ErrorTracker()
-        self._retry_histogram.clear()
         self._requests_in_progress.clear()
 
-        if self._persistence_enabled and self._key_value_store:
-            await self._key_value_store.set_value(self._persist_state_key, None)
-
     def _log(self) -> None:
         stats = self.calculate()
         if self._statistics_log_format == 'table':
@@ -296,62 +281,27 @@ def _log(self) -> None:
         else:
             self._periodic_message_logger.info(self._log_message, extra=stats.to_dict())
 
-    async def _maybe_load_statistics(self) -> None:
-        if not self._persistence_enabled:
-            return
-
-        if not self._key_value_store:
-            return
-
-        stored_state = await self._key_value_store.get_value(self._persist_state_key, cast('Any', {}))
+    def _after_initialize(self) -> None:
+        state = self._state.current_value
 
-        saved_state = self.state.__class__.model_validate(stored_state)
-        self.state = saved_state
+        if state.crawler_started_at is None:
+            state.crawler_started_at = datetime.now(timezone.utc)
 
-        if saved_state.stats_persisted_at is not None and saved_state.crawler_last_started_at:
+        if state.stats_persisted_at is not None and state.crawler_last_started_at:
             self._instance_start = datetime.now(timezone.utc) - (
-                saved_state.stats_persisted_at - saved_state.crawler_last_started_at
+                state.stats_persisted_at - state.crawler_last_started_at
             )
-        elif saved_state.crawler_last_started_at:
-            self._instance_start = saved_state.crawler_last_started_at
+        elif state.crawler_last_started_at:
+            self._instance_start = state.crawler_last_started_at
 
-    async def _persist_state(self, event_data: EventPersistStateData) -> None:
-        logger.debug(f'Persisting state of the Statistics (event_data={event_data}).')
-
-        if not self._persistence_enabled:
-            return
-
-        if not self._key_value_store:
-            return
-
-        if not self._instance_start:
-            return
-
-        final_statistics = self.calculate()
-        persisted_state = StatisticsPersistedState(
-            stats_id=self._id,
-            stats_persisted_at=datetime.now(timezone.utc),
-            crawler_last_started_at=self._instance_start,
-            request_total_duration=final_statistics.request_total_duration,
-            request_avg_failed_duration=final_statistics.request_avg_failed_duration,
-            request_avg_finished_duration=final_statistics.request_avg_finished_duration,
-            requests_total=final_statistics.requests_total,
-            request_retry_histogram=final_statistics.retry_histogram,
-        )
-
-        logger.debug('Persisting state')
-
-        await self._key_value_store.set_value(
-            self._persist_state_key,
-            self.state.model_dump(mode='json', by_alias=True) | persisted_state.model_dump(mode='json', by_alias=True),
-            'application/json',
-        )
+        state.crawler_last_started_at = self._instance_start
 
     def _save_retry_count_for_request(self, record: RequestProcessingRecord) -> None:
         retry_count = record.retry_count
+        state = self._state.current_value
 
         if retry_count:
-            self.state.requests_retries += 1
+            state.requests_retries += 1
 
-        self._retry_histogram.setdefault(retry_count, 0)
-        self._retry_histogram[retry_count] += 1
+        state.request_retry_histogram.setdefault(retry_count, 0)
+        state.request_retry_histogram[retry_count] += 1
diff --git a/src/crawlee/storages/_creation_management.py b/src/crawlee/storages/_creation_management.py
index d7356a98b5..14d9b1719e 100644
--- a/src/crawlee/storages/_creation_management.py
+++ b/src/crawlee/storages/_creation_management.py
@@ -2,6 +2,7 @@
 
 import asyncio
 from typing import TYPE_CHECKING, TypeVar
+from weakref import WeakKeyDictionary
 
 from crawlee.storage_clients import MemoryStorageClient
 
@@ -16,8 +17,8 @@
 TResource = TypeVar('TResource', Dataset, KeyValueStore, RequestQueue)
 
 
-_creation_lock = asyncio.Lock()
-"""Lock for storage creation."""
+_creation_locks = WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]()
+"""Locks for storage creation (we need a separate lock for every event loop so that tests work as expected)."""
 
 _cache_dataset_by_id: dict[str, Dataset] = {}
 _cache_dataset_by_name: dict[str, Dataset] = {}
@@ -154,7 +155,11 @@ async def open_storage(
         await storage_client.purge_on_start()
 
     # Lock and create new storage
-    async with _creation_lock:
+    loop = asyncio.get_running_loop()
+    if loop not in _creation_locks:
+        _creation_locks[loop] = asyncio.Lock()
+
+    async with _creation_locks[loop]:
         if id and not is_default_on_memory:
             resource_client = _get_resource_client(storage_class, storage_client, id)
             storage_object = await resource_client.get()
diff --git a/src/crawlee/storages/_key_value_store.py b/src/crawlee/storages/_key_value_store.py
index b7d3a4b582..fc077726d1 100644
--- a/src/crawlee/storages/_key_value_store.py
+++ b/src/crawlee/storages/_key_value_store.py
@@ -3,13 +3,15 @@
 import asyncio
 from collections.abc import AsyncIterator
 from datetime import datetime, timezone
+from logging import getLogger
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, overload
 
+from pydantic import RootModel
 from typing_extensions import override
 
 from crawlee import service_locator
+from crawlee._types import JsonSerializable  # noqa: TC001
 from crawlee._utils.docs import docs_group
-from crawlee.events._types import Event, EventPersistStateData
 from crawlee.storage_clients.models import KeyValueStoreKeyInfo, KeyValueStoreMetadata, StorageMetadata
 
 from ._base import Storage
@@ -17,12 +19,18 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
-    from crawlee._types import JsonSerializable
+    from crawlee._utils.recoverable_state import RecoverableState
     from crawlee.configuration import Configuration
     from crawlee.storage_clients import StorageClient
 
 T = TypeVar('T')
 
+logger = getLogger(__name__)
+
+
+class AutosavedValue(RootModel):
+    root: dict[str, JsonSerializable]
+
 
 @docs_group('Classes')
 class KeyValueStore(Storage):
@@ -59,9 +67,13 @@ class KeyValueStore(Storage):
     ```
     """
 
-    # Cache for persistent (auto-saved) values
-    _general_cache: ClassVar[dict[str, dict[str, dict[str, JsonSerializable]]]] = {}
-    _persist_state_event_started = False
+    # Cache for recoverable (auto-saved) values
+    _autosaved_values: ClassVar[
+        dict[
+            str,
+            dict[str, RecoverableState[AutosavedValue]],
+        ]
+    ] = {}
 
     def __init__(self, id: str, name: str | None, storage_client: StorageClient) -> None:
         self._id = id
@@ -138,9 +150,9 @@ async def open(
     async def drop(self) -> None:
         from crawlee.storages._creation_management import remove_storage_from_cache
 
-        await self._resource_client.delete()
-        self._clear_cache()
         remove_storage_from_cache(storage_class=self.__class__, id=self._id, name=self._name)
+        await self._clear_cache()
+        await self._resource_client.delete()
 
     @overload
     async def get_value(self, key: str) -> Any: ...
@@ -225,58 +237,39 @@ async def get_auto_saved_value(
         Returns:
             Return the value of the key.
         """
+        from crawlee._utils.recoverable_state import RecoverableState
+
         default_value = {} if default_value is None else default_value
 
         async with self._autosave_lock:
-            if key in self._cache:
-                return self._cache[key]
+            cache = self._autosaved_values.setdefault(self._id, {})
 
-            value = await self.get_value(key, default_value)
+            if key in cache:
+                return cache[key].current_value.root
 
-            if not isinstance(value, dict):
-                raise TypeError(
-                    f'Expected dictionary for persist state value at key "{key}, but got {type(value).__name__}'
-                )
+            cache[key] = recoverable_state = RecoverableState(
+                default_state=AutosavedValue(default_value),
+                persistence_enabled=True,
+                persist_state_kvs_id=self._id,
+                persist_state_key=key,
+                logger=logger,
+            )
 
-            self._cache[key] = value
+            await recoverable_state.initialize()
 
-        self._ensure_persist_event()
+        return recoverable_state.current_value.root
 
-        return value
-
-    @property
-    def _cache(self) -> dict[str, dict[str, JsonSerializable]]:
-        """Cache dictionary for storing auto-saved values indexed by store ID."""
-        if self._id not in self._general_cache:
-            self._general_cache[self._id] = {}
-        return self._general_cache[self._id]
-
-    async def _persist_save(self, _event_data: EventPersistStateData | None = None) -> None:
-        """Save cache with persistent values. Can be used in Event Manager."""
-        for key, value in self._cache.items():
-            await self.set_value(key, value)
-
-    def _ensure_persist_event(self) -> None:
-        """Ensure persist state event handling if not already done."""
-        if self._persist_state_event_started:
-            return
-
-        event_manager = service_locator.get_event_manager()
-        event_manager.on(event=Event.PERSIST_STATE, listener=self._persist_save)
-        self._persist_state_event_started = True
-
-    def _clear_cache(self) -> None:
-        """Clear cache with persistent values."""
-        self._cache.clear()
-
-    def _drop_persist_state_event(self) -> None:
-        """Off event_manager listener and drop event status."""
-        if self._persist_state_event_started:
-            event_manager = service_locator.get_event_manager()
-            event_manager.off(event=Event.PERSIST_STATE, listener=self._persist_save)
-        self._persist_state_event_started = False
+    async def _clear_cache(self) -> None:
+        """Clear cache with autosaved values."""
+        if self.id in self._autosaved_values:
+            cache = self._autosaved_values[self.id]
+            for value in cache.values():
+                await value.teardown()
+            cache.clear()
 
     async def persist_autosaved_values(self) -> None:
-        """Force persistent values to be saved without waiting for an event in Event Manager."""
-        if self._persist_state_event_started:
-            await self._persist_save()
+        """Force autosaved values to be saved without waiting for an event in Event Manager."""
+        if self.id in self._autosaved_values:
+            cache = self._autosaved_values[self.id]
+            for value in cache.values():
+                await value.persist_state()


Test patch:
diff --git a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py
index 89b028ed81..63d078e95e 100644
--- a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py
+++ b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py
@@ -89,7 +89,7 @@ def predict(self, request: Request) -> RenderingTypePrediction:
         return RenderingTypePrediction(next(self._rendering_types), next(self._detection_probability_recommendation))
 
     @override
-    def store_result(self, request: Request, crawl_type: RenderingType) -> None:
+    def store_result(self, request: Request, rendering_type: RenderingType) -> None:
         pass
 
 
@@ -464,13 +464,13 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
         assert stored_results == [saved_data]
 
 
-def test_adaptive_playwright_crawler_statistics_in_init() -> None:
+async def test_adaptive_playwright_crawler_statistics_in_init() -> None:
     """Tests that adaptive crawler uses created AdaptivePlaywrightCrawlerStatistics from inputted Statistics."""
     persistence_enabled = True
     persist_state_kvs_name = 'some name'
     persist_state_key = 'come key'
     log_message = 'some message'
-    periodic_message_logger = logging.getLogger('some logger')  # Accessing private member to create copy like-object.
+    periodic_message_logger = logging.getLogger('some logger')
     log_interval = timedelta(minutes=2)
     statistics = Statistics.with_default_state(
         persistence_enabled=persistence_enabled,
@@ -482,11 +482,14 @@ def test_adaptive_playwright_crawler_statistics_in_init() -> None:
     )
 
     crawler = AdaptivePlaywrightCrawler.with_beautifulsoup_static_parser(statistics=statistics)
+    await crawler.run([])  # ensure that statistics get initialized
 
     assert type(crawler._statistics.state) is AdaptivePlaywrightCrawlerStatisticState
-    assert crawler._statistics._persistence_enabled == persistence_enabled
-    assert crawler._statistics._persist_state_kvs_name == persist_state_kvs_name
-    assert crawler._statistics._persist_state_key == persist_state_key
+
+    assert crawler._statistics._state._persistence_enabled == persistence_enabled
+    assert crawler._statistics._state._persist_state_kvs_name == persist_state_kvs_name
+    assert crawler._statistics._state._persist_state_key == persist_state_key
+
     assert crawler._statistics._log_message == log_message
     assert crawler._statistics._periodic_message_logger == periodic_message_logger
 
diff --git a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler_statistics.py b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler_statistics.py
index 435e374ae3..2bf77d42c9 100644
--- a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler_statistics.py
+++ b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler_statistics.py
@@ -14,7 +14,7 @@ async def test_predictor_state_persistence() -> None:
         adaptive_statistics.state.rendering_type_mispredictions = 2
         adaptive_statistics.state.http_only_request_handler_runs = 3
 
-        persistence_state_key = adaptive_statistics._persist_state_key
+        persistence_state_key = adaptive_statistics._state._persist_state_key
     # adaptive_statistics are persisted after leaving the context
 
     # new_adaptive_statistics are initialized from the persisted values.
diff --git a/tests/unit/crawlers/_playwright/test_playwright_crawler.py b/tests/unit/crawlers/_playwright/test_playwright_crawler.py
index 35b663a567..8a339dc212 100644
--- a/tests/unit/crawlers/_playwright/test_playwright_crawler.py
+++ b/tests/unit/crawlers/_playwright/test_playwright_crawler.py
@@ -597,7 +597,7 @@ async def request_handler(context: PlaywrightCrawlingContext) -> None:
     # Three errors twice retried errors, but only 2 unique -> 4 (2 x (html and jpg)) artifacts expected.
     assert crawler.statistics.error_tracker.total == 3 * max_retries
     assert crawler.statistics.error_tracker.unique_error_count == 2
-    assert len(kvs_content) == 4
+    assert len(list(kvs_content.keys())) == 4
 
 
 async def test_respect_robots_txt(server_url: URL) -> None:
diff --git a/tests/unit/sessions/test_models.py b/tests/unit/sessions/test_models.py
index 450777f226..fee469c475 100644
--- a/tests/unit/sessions/test_models.py
+++ b/tests/unit/sessions/test_models.py
@@ -5,7 +5,7 @@
 import pytest
 
 from crawlee.sessions._cookies import CookieParam
-from crawlee.sessions._models import SessionModel, SessionPoolModel
+from crawlee.sessions._models import SessionModel
 
 SESSION_CREATED_AT = datetime.now(timezone.utc)
 
@@ -78,33 +78,3 @@ def test_session_model(
 
     # Check that max_age is correctly parsed into a timedelta object
     assert session_direct.max_age == session_camel.max_age == session_snake.max_age == timedelta(minutes=30)
-
-
-def test_create_session_pool_with_direct_sessions(session_direct: SessionModel) -> None:
-    """Test creating a SessionPoolModel with direct session model instances."""
-    session_pool = SessionPoolModel(
-        persistence_enabled=False,
-        persist_state_kvs_name='test_session_pool',
-        persist_state_key='crawlee_session_pool_state',
-        max_pool_size=3,
-        session_count=0,
-        usable_session_count=0,
-        retired_session_count=0,
-        sessions=[session_direct],
-    )
-    assert session_pool.sessions == [session_direct]
-
-
-def test_create_session_pool_with_args_sessions(session_args_camel: dict, session_args_snake: dict) -> None:
-    """Test creating a SessionPoolModel using sessions initialized from camelCase and snake_case dicts."""
-    session_pool_camel = SessionPoolModel(
-        persistence_enabled=False,
-        persist_state_kvs_name='test_session_pool',
-        persist_state_key='crawlee_session_pool_state',
-        max_pool_size=3,
-        session_count=0,
-        usable_session_count=0,
-        retired_session_count=0,
-        sessions=[session_args_camel, session_args_snake],
-    )
-    assert session_pool_camel.sessions == [SessionModel(**session_args_camel), SessionModel(**session_args_snake)]
diff --git a/tests/unit/sessions/test_session_pool.py b/tests/unit/sessions/test_session_pool.py
index 6bdce8bd46..e422628124 100644
--- a/tests/unit/sessions/test_session_pool.py
+++ b/tests/unit/sessions/test_session_pool.py
@@ -139,8 +139,7 @@ async def test_session_pool_persist(event_manager: EventManager, kvs_name: str |
         assert sp_model.retired_session_count == sp.retired_session_count
 
         # Check if all the sessions are correctly persisted
-        for session_model in sp_model.sessions:
-            kvs_session = Session.from_model(model=session_model)
+        for kvs_session in sp_model.sessions.values():
             session = await sp.get_session_by_id(kvs_session.id)
             assert kvs_session == session
 
diff --git a/tests/unit/storages/test_key_value_store.py b/tests/unit/storages/test_key_value_store.py
index 955d483546..ea3f4e5f7d 100644
--- a/tests/unit/storages/test_key_value_store.py
+++ b/tests/unit/storages/test_key_value_store.py
@@ -194,7 +194,7 @@ async def test_get_auto_saved_value_auto_save_race_conditions(key_value_store: K
 
     sleep_time_iterator = chain(iter([0.5]), repeat(0))
 
-    async def delayed_get_value(key: str, default_value: None) -> None:
+    async def delayed_get_value(key: str, default_value: None = None) -> None:
         await asyncio.sleep(next(sleep_time_iterator))
         return await KeyValueStore.get_value(key_value_store, key=key, default_value=default_value)