From 37fe1b798543d01f611d66628b46f5a8c2cde21e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 10:59:33 +0000 Subject: [PATCH 1/7] fix: persist the charged event counts across container restarts After a reboot, migration or restart the platform seeds the new container with the `APIFY_CHARGED_ACTOR_EVENT_COUNTS` recorded when the run started, so the charging manager began from zero charged events and allowed a fresh `max_total_charge_usd` budget. Reading the counts back from the run does not help either: the platform flushes charges to the run document every few seconds, and a reboot restarts the container faster than that. The charging manager now keeps its charged event counts in a Crawlee `RecoverableState` stored in the default key-value store under `__ACTOR_CHARGING_STATE`. It is written on every `PERSIST_STATE` event, which a reboot and a migration dispatch before the container stops, and once more on exit. A pay-per-event run on the platform reads it at start and takes the persisted count of each event where it is higher than the env var's. Local runs keep their counts in memory only. Closes #1209 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 44 +++++++++++++ tests/unit/actor/test_charging_manager.py | 75 ++++++++++++++++++++++- 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index fc404ffb..027ca13d 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Literal, Protocol, TypedDict from cachetools import TTLCache +from pydantic import BaseModel, ConfigDict, Field +from pydantic.alias_generators import to_camel from apify_client._models import ( FlatPricePerMonthActorPricingInfo, @@ -17,6 +19,7 @@ PayPerEventActorPricingInfo, PricePerDatasetItemActorPricingInfo, ) +from crawlee._utils.recoverable_state import RecoverableState from apify._utils import ReentrantLock, docs_group, ensure_context from apify.storages import Dataset @@ -40,6 +43,9 @@ DEFAULT_DATASET_ITEM_EVENT = 'apify-default-dataset-item' """Name of the synthetic event charged for each item pushed to the default dataset.""" +CHARGING_STATE_KEY = '__ACTOR_CHARGING_STATE' +"""Key of the default key-value store record that carries the charged event counts across a container restart.""" + PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE'] """Pricing model for an Actor.""" @@ -213,6 +219,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None + self._recoverable_state: RecoverableState[ChargingState] | None = None self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} @@ -273,6 +280,28 @@ async def __aenter__(self) -> None: total_charged_amount=count * price, ) + # The platform fills the env var with the counts recorded when the run started, so after a reboot, migration + # or restart they miss what an earlier container charged. That container persisted its counts to the default + # key-value store, so they take over where they are higher. + if self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': + self._recoverable_state = RecoverableState( + default_state=ChargingState(), + persist_state_key=CHARGING_STATE_KEY, + persistence_enabled=True, + logger=logger, + ) + persisted_state = await self._recoverable_state.initialize() + for event_name, count in persisted_state.charged_event_counts.items(): + if count > self.get_charged_event_count(event_name): + price = self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price + self._charging_state[event_name] = ChargingStateItem( + charge_count=count, + total_charged_amount=count * price, + ) + persisted_state.charged_event_counts = { + event_name: item.charge_count for event_name, item in self._charging_state.items() + } + # Set up charging log dataset for local development if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': # We are not running on the Apify platform, but PPE is enabled for testing - open a dataset that @@ -298,6 +327,9 @@ async def __aexit__( if not self.active: raise RuntimeError('Exiting an uninitialized ChargingManager') + if self._recoverable_state is not None: + await self._recoverable_state.teardown() + charging_manager_ctx.set(None) self.active = False @@ -393,6 +425,9 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str self._charging_state.setdefault(event_name, ChargingStateItem(0, Decimal())) self._charging_state[event_name].charge_count += charged_count self._charging_state[event_name].total_charged_amount += charged_count * pricing_info.price + if self._recoverable_state is not None: + counts = self._recoverable_state.current_value.charged_event_counts + counts[event_name] = self._charging_state[event_name].charge_count # Remember the key for every charge that was counted, including events the API never receives, such as # synthetic and tier-priced ones - those are counted locally and a repeat would count them twice. @@ -571,6 +606,15 @@ def _get_event_price(self, event_name: str) -> Decimal: return Decimal(0) if self._is_at_home else Decimal(1) +class ChargingState(BaseModel): + """The charged event counts of a run, as persisted in its default key-value store.""" + + model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel) + + charged_event_counts: dict[str, int] = Field(default_factory=dict) + """Number of charged events by event name, synthetic events included.""" + + @dataclass class ChargingStateItem: charge_count: int diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index ae50b041..2e228926 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -8,9 +8,12 @@ import pytest from apify_client._models import ActorChargeEvent, PayPerEventActorPricingInfo +from crawlee import service_locator +from crawlee.events._types import Event, EventPersistStateData -from apify._charging import ChargingManagerImplementation +from apify._charging import CHARGING_STATE_KEY, ChargingManagerImplementation from apify._configuration import Configuration +from apify.storages import KeyValueStore def _make_config(**kwargs: Any) -> Configuration: @@ -781,3 +784,73 @@ async def test_charge_overcharges_by_one_event_when_budget_exceeded( mock_client.run.return_value.charge.assert_awaited_once_with('search', count=2, idempotency_key=None) else: mock_client.run.return_value.charge.assert_not_called() + + +def _make_platform_ppe_config(**kwargs: Any) -> Configuration: + """Helper to create a platform PPE Configuration with charged event counts in the env var.""" + defaults: dict[str, Any] = { + 'is_at_home': True, + 'actor_run_id': 'test-run-id', + 'actor_pricing_info': _make_ppe_pricing_info(), + 'charged_event_counts': {'search': 1, 'scrape': 2}, + 'max_total_charge_usd': Decimal('10.00'), + } + defaults.update(kwargs) + return _make_config(**defaults) + + +async def test_persisted_counts_take_over_where_higher(mock_client: MagicMock) -> None: + """A platform start takes the counts persisted by an earlier container where they exceed the env var's.""" + key_value_store = await KeyValueStore.open() + await key_value_store.set_value(CHARGING_STATE_KEY, {'chargedEventCounts': {'search': 3, 'scrape': 1}}) + + cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) + async with cm: + assert cm.get_charged_event_count('search') == 3 + assert cm.get_charged_event_count('scrape') == 2 + assert cm.calculate_total_charged_amount() == Decimal('0.13') + assert cm.calculate_max_event_charge_count_within_limit('search') == 987 + + +async def test_charged_counts_are_persisted(mock_client: MagicMock) -> None: + """Charges are written to the default key-value store on persist and once more on exit.""" + key_value_store = await KeyValueStore.open() + cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) + async with cm: + await cm.charge('search', count=2) + assert cm._recoverable_state is not None + await cm._recoverable_state.persist_state() + stored = await key_value_store.get_value(CHARGING_STATE_KEY) + assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 2}} + + await cm.charge('scrape') + + stored = await key_value_store.get_value(CHARGING_STATE_KEY) + assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 3}} + + +async def test_local_run_does_not_persist_counts(mock_client: MagicMock) -> None: + """A local PPE run keeps its counts in memory only.""" + key_value_store = await KeyValueStore.open() + cm = ChargingManagerImplementation( + _make_platform_ppe_config(is_at_home=False, test_pay_per_event=True), mock_client + ) + async with cm: + await cm.charge('search') + + assert await key_value_store.get_value(CHARGING_STATE_KEY) is None + + +async def test_persist_state_event_persists_counts(mock_client: MagicMock) -> None: + """The counts are written when the `PERSIST_STATE` event fires, which a reboot and a migration trigger.""" + key_value_store = await KeyValueStore.open() + cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) + async with cm: + await cm.charge('search') + + event_manager = service_locator.get_event_manager() + async with event_manager: + event_manager.emit(event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=True)) + await event_manager.wait_for_all_listeners_to_complete() + + assert await key_value_store.get_value(CHARGING_STATE_KEY) == {'chargedEventCounts': {'search': 2, 'scrape': 2}} From 3953d9df2ffaf0354915b21f16dc5ad9e8532559 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 13:40:13 +0000 Subject: [PATCH 2/7] refactor: keep the charged event counts in the recoverable state only The recoverable state is the single owner of the counts on every run, and is persisted only on the platform. The per-event totals are derived from the counts and the prices instead of being stored alongside them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 75 ++++++++--------------- tests/unit/actor/test_charging_manager.py | 4 +- 2 files changed, 29 insertions(+), 50 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 027ca13d..11ca4ad4 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -22,7 +22,7 @@ from crawlee._utils.recoverable_state import RecoverableState from apify._utils import ReentrantLock, docs_group, ensure_context -from apify.storages import Dataset +from apify.storages import Dataset, KeyValueStore if TYPE_CHECKING: from collections.abc import AsyncGenerator @@ -219,9 +219,9 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None - self._recoverable_state: RecoverableState[ChargingState] | None = None + self._state: RecoverableState[ChargingState] | None = None + """The charged event counts, persisted to the default key-value store on the platform.""" - self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() self._idempotent_charges: TTLCache[str, IdempotentChargeItem] = TTLCache(maxsize=float('inf'), ttl=180) @@ -272,35 +272,19 @@ async def __aenter__(self) -> None: self._max_total_charge_usd = max_total_charge_usd - # Load charged event counts - for event_name, count in charged_event_counts.items(): - price = self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price - self._charging_state[event_name] = ChargingStateItem( - charge_count=count, - total_charged_amount=count * price, - ) - # The platform fills the env var with the counts recorded when the run started, so after a reboot, migration - # or restart they miss what an earlier container charged. That container persisted its counts to the default - # key-value store, so they take over where they are higher. - if self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': - self._recoverable_state = RecoverableState( - default_state=ChargingState(), - persist_state_key=CHARGING_STATE_KEY, - persistence_enabled=True, - logger=logger, - ) - persisted_state = await self._recoverable_state.initialize() - for event_name, count in persisted_state.charged_event_counts.items(): - if count > self.get_charged_event_count(event_name): - price = self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price - self._charging_state[event_name] = ChargingStateItem( - charge_count=count, - total_charged_amount=count * price, - ) - persisted_state.charged_event_counts = { - event_name: item.charge_count for event_name, item in self._charging_state.items() - } + # or restart they miss what an earlier container charged. On the platform the counts are therefore also kept + # in the default key-value store, and the persisted ones take over where they are higher. + self._state = RecoverableState( + default_state=ChargingState(), + persist_state_key=CHARGING_STATE_KEY, + persistence_enabled=self._is_at_home and self._pricing_model == 'PAY_PER_EVENT', + persist_state_kvs_factory=KeyValueStore.open, + logger=logger, + ) + counts = (await self._state.initialize()).charged_event_counts + for event_name, count in charged_event_counts.items(): + counts[event_name] = max(counts.get(event_name, 0), count) # Set up charging log dataset for local development if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': @@ -327,8 +311,8 @@ async def __aexit__( if not self.active: raise RuntimeError('Exiting an uninitialized ChargingManager') - if self._recoverable_state is not None: - await self._recoverable_state.teardown() + if self._state is not None: + await self._state.teardown() charging_manager_ctx.set(None) self.active = False @@ -422,12 +406,8 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str # Count the charge only after the API call succeeds, so a failed call retried under the same key is # counted once. - self._charging_state.setdefault(event_name, ChargingStateItem(0, Decimal())) - self._charging_state[event_name].charge_count += charged_count - self._charging_state[event_name].total_charged_amount += charged_count * pricing_info.price - if self._recoverable_state is not None: - counts = self._recoverable_state.current_value.charged_event_counts - counts[event_name] = self._charging_state[event_name].charge_count + counts = self._charged_event_counts + counts[event_name] = counts.get(event_name, 0) + charged_count # Remember the key for every charge that was counted, including events the API never receives, such as # synthetic and tier-priced ones - those are counted locally and a repeat would count them twice. @@ -467,7 +447,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str @_ensure_context def calculate_total_charged_amount(self) -> Decimal: return sum( - (item.total_charged_amount for item in self._charging_state.values()), + (count * self._get_event_price(event_name) for event_name, count in self._charged_event_counts.items()), start=Decimal(), ) @@ -496,8 +476,7 @@ def get_pricing_info(self) -> ActorPricingInfo: @_ensure_context def get_charged_event_count(self, event_name: str) -> int: - item = self._charging_state.get(event_name) - return item.charge_count if item is not None else 0 + return self._charged_event_counts.get(event_name, 0) @_ensure_context def get_max_total_charge_usd(self) -> Decimal: @@ -599,6 +578,12 @@ def _compute_charged_count(self, event_name: str, count: int) -> int: return 0 + @property + def _charged_event_counts(self) -> dict[str, int]: + if self._state is None: + raise RuntimeError('ChargingManager is not initialized') + return self._state.current_value.charged_event_counts + def _get_event_price(self, event_name: str) -> Decimal: pricing_info = self._pricing_info.get(event_name) if pricing_info is not None: @@ -615,12 +600,6 @@ class ChargingState(BaseModel): """Number of charged events by event name, synthetic events included.""" -@dataclass -class ChargingStateItem: - charge_count: int - total_charged_amount: Decimal - - @dataclass class PricingInfoItem: price: Decimal diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index 2e228926..9a32ad97 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -818,8 +818,8 @@ async def test_charged_counts_are_persisted(mock_client: MagicMock) -> None: cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) async with cm: await cm.charge('search', count=2) - assert cm._recoverable_state is not None - await cm._recoverable_state.persist_state() + assert cm._state is not None + await cm._state.persist_state() stored = await key_value_store.get_value(CHARGING_STATE_KEY) assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 2}} From 892e52fe45ab31af2e640c866bf334ffb6dea087 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 13:40:13 +0000 Subject: [PATCH 3/7] test: cover charged event counts surviving a reboot on the platform Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- tests/e2e/test_actor_charge.py | 95 ++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tests/e2e/test_actor_charge.py b/tests/e2e/test_actor_charge.py index dc09602a..162787fc 100644 --- a/tests/e2e/test_actor_charge.py +++ b/tests/e2e/test_actor_charge.py @@ -226,3 +226,98 @@ async def test_actor_push_data_combined_budget_limit( assert run.status == 'SUCCEEDED' assert run.charged_event_counts == expected_counts + + +@pytest_asyncio.fixture(scope='module', loop_scope='module') +async def ppe_reboot_actor_build(make_actor: MakeActorFunction) -> str: + async def main() -> None: + async with Actor: + key_value_store = await Actor.open_key_value_store() + charging_manager = Actor.get_charging_manager() + + if not await key_value_store.get_value('REBOOTED'): + for _ in range(3): + await Actor.charge('row-charged') + await key_value_store.set_value('REBOOTED', {'rebooted': True}) + await Actor.reboot() + return + + # The second container records what the SDK believes right at its start, before charging anything. + await key_value_store.set_value( + 'SDK_VIEW_AFTER_REBOOT', + { + 'chargedCount': charging_manager.get_charged_event_count('row-charged'), + 'totalChargedUsd': str(charging_manager.calculate_total_charged_amount()), + 'chargeableWithinLimit': charging_manager.calculate_max_event_charge_count_within_limit( + 'row-charged' + ), + }, + ) + + actor_client = await make_actor('ppe-reboot', main_func=main) + + await actor_client.update( + pricing_infos=[ + { + 'pricingModel': 'PAY_PER_EVENT', + 'apifyMarginPercentage': 0.0, + 'createdAt': '2024-01-01T00:00:00.000Z', + 'startedAt': '2024-01-01T00:00:00.000Z', + 'pricingPerEvent': { + 'actorChargeEvents': { + 'row-charged': { + 'eventTitle': 'Row charged', + 'eventPriceUsd': 0.005, + 'eventDescription': 'One charged row', + }, + }, + }, + }, + ], + ) + + actor = await actor_client.get() + + assert actor is not None + return actor.id + + +@pytest_asyncio.fixture(scope='function', loop_scope='module') +async def ppe_reboot_actor( + ppe_reboot_actor_build: str, + apify_client_async: ApifyClientAsync, +) -> ActorClientAsync: + return apify_client_async.actor(ppe_reboot_actor_build) + + +async def test_actor_charge_counts_survive_reboot( + ppe_reboot_actor: ActorClientAsync, + run_actor: RunActorFunction, + apify_client_async: ApifyClientAsync, +) -> None: + """Test that a rebooted run starts from the events charged before the reboot, instead of a fresh budget.""" + run = await run_actor(ppe_reboot_actor, max_total_charge_usd=Decimal('0.05')) + + assert run.status == 'SUCCEEDED' + assert run.stats.reboot_count == 1 + + # Refetch until the charged event counts propagate on the platform. + run = await poll_until_condition( + partial(_get_run, apify_client_async, run.id), + lambda r: r.charged_event_counts == {'row-charged': 3}, + timeout=30, + poll_interval=1, + ) + assert run.charged_event_counts == {'row-charged': 3} + + record = await apify_client_async.key_value_store(run.default_key_value_store_id).get_record( + 'SDK_VIEW_AFTER_REBOOT' + ) + assert record is not None + + # Three events at $0.005 were charged before the reboot, which leaves room for seven more under the $0.05 limit. + assert record['value'] == { + 'chargedCount': 3, + 'totalChargedUsd': '0.015', + 'chargeableWithinLimit': 7, + } From f1cd048d408691de5935121c22acda170a1a0e56 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 14:00:38 +0000 Subject: [PATCH 4/7] refactor: name the recoverable state `_charging_state` Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 14 +++++++------- tests/unit/actor/test_charging_manager.py | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 11ca4ad4..591e4dcf 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -219,7 +219,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None - self._state: RecoverableState[ChargingState] | None = None + self._charging_state: RecoverableState[ChargingState] | None = None """The charged event counts, persisted to the default key-value store on the platform.""" self._pricing_info: dict[str, PricingInfoItem] = {} @@ -275,14 +275,14 @@ async def __aenter__(self) -> None: # The platform fills the env var with the counts recorded when the run started, so after a reboot, migration # or restart they miss what an earlier container charged. On the platform the counts are therefore also kept # in the default key-value store, and the persisted ones take over where they are higher. - self._state = RecoverableState( + self._charging_state = RecoverableState( default_state=ChargingState(), persist_state_key=CHARGING_STATE_KEY, persistence_enabled=self._is_at_home and self._pricing_model == 'PAY_PER_EVENT', persist_state_kvs_factory=KeyValueStore.open, logger=logger, ) - counts = (await self._state.initialize()).charged_event_counts + counts = (await self._charging_state.initialize()).charged_event_counts for event_name, count in charged_event_counts.items(): counts[event_name] = max(counts.get(event_name, 0), count) @@ -311,8 +311,8 @@ async def __aexit__( if not self.active: raise RuntimeError('Exiting an uninitialized ChargingManager') - if self._state is not None: - await self._state.teardown() + if self._charging_state is not None: + await self._charging_state.teardown() charging_manager_ctx.set(None) self.active = False @@ -580,9 +580,9 @@ def _compute_charged_count(self, event_name: str, count: int) -> int: @property def _charged_event_counts(self) -> dict[str, int]: - if self._state is None: + if self._charging_state is None: raise RuntimeError('ChargingManager is not initialized') - return self._state.current_value.charged_event_counts + return self._charging_state.current_value.charged_event_counts def _get_event_price(self, event_name: str) -> Decimal: pricing_info = self._pricing_info.get(event_name) diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index 9a32ad97..5bba4746 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -818,8 +818,8 @@ async def test_charged_counts_are_persisted(mock_client: MagicMock) -> None: cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) async with cm: await cm.charge('search', count=2) - assert cm._state is not None - await cm._state.persist_state() + assert cm._charging_state is not None + await cm._charging_state.persist_state() stored = await key_value_store.get_value(CHARGING_STATE_KEY) assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 2}} From a73c808adaa5ea4da473d56fc7a2ec6c597f74b0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 14:06:33 +0000 Subject: [PATCH 5/7] refactor: drop the intermediary count variables Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 591e4dcf..2004fb28 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -282,9 +282,9 @@ async def __aenter__(self) -> None: persist_state_kvs_factory=KeyValueStore.open, logger=logger, ) - counts = (await self._charging_state.initialize()).charged_event_counts + await self._charging_state.initialize() for event_name, count in charged_event_counts.items(): - counts[event_name] = max(counts.get(event_name, 0), count) + self._charged_event_counts[event_name] = max(self._charged_event_counts.get(event_name, 0), count) # Set up charging log dataset for local development if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': @@ -406,8 +406,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str # Count the charge only after the API call succeeds, so a failed call retried under the same key is # counted once. - counts = self._charged_event_counts - counts[event_name] = counts.get(event_name, 0) + charged_count + self._charged_event_counts[event_name] = self._charged_event_counts.get(event_name, 0) + charged_count # Remember the key for every charge that was counted, including events the API never receives, such as # synthetic and tier-priced ones - those are counted locally and a repeat would count them twice. From 6326c3a48bbb5e18e25d62ce2294f23e87e769ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Josef=20Proch=C3=A1zka?= Date: Tue, 3 Nov 2026 16:00:11 +0100 Subject: [PATCH 6/7] Apply batched suggestions from code review Co-authored-by: Vlada Dusek --- src/apify/_charging.py | 20 +++++++++++++------- tests/e2e/test_actor_charge.py | 2 +- tests/unit/actor/test_charging_manager.py | 5 ++++- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 2004fb28..7f95c2c2 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -219,7 +219,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None - self._charging_state: RecoverableState[ChargingState] | None = None + """The charged event counts, persisted to the default key-value store in a pay-per-event run on the platform.""" """The charged event counts, persisted to the default key-value store on the platform.""" self._pricing_info: dict[str, PricingInfoItem] = {} @@ -282,7 +282,12 @@ async def __aenter__(self) -> None: persist_state_kvs_factory=KeyValueStore.open, logger=logger, ) - await self._charging_state.initialize() + try: + await self._charging_state.initialize() + except ValidationError as exc: + raise ValueError( + f'The charging state under the "{CHARGING_STATE_KEY}" key in the default key-value store is malformed.' + ) from exc for event_name, count in charged_event_counts.items(): self._charged_event_counts[event_name] = max(self._charged_event_counts.get(event_name, 0), count) @@ -311,11 +316,12 @@ async def __aexit__( if not self.active: raise RuntimeError('Exiting an uninitialized ChargingManager') - if self._charging_state is not None: - await self._charging_state.teardown() - - charging_manager_ctx.set(None) - self.active = False + try: + if self._charging_state is not None: + await self._charging_state.teardown() + finally: + charging_manager_ctx.set(None) + self.active = False @_ensure_context async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult: diff --git a/tests/e2e/test_actor_charge.py b/tests/e2e/test_actor_charge.py index 162787fc..47bd86f4 100644 --- a/tests/e2e/test_actor_charge.py +++ b/tests/e2e/test_actor_charge.py @@ -295,7 +295,7 @@ async def test_actor_charge_counts_survive_reboot( run_actor: RunActorFunction, apify_client_async: ApifyClientAsync, ) -> None: - """Test that a rebooted run starts from the events charged before the reboot, instead of a fresh budget.""" + """Test that a rebooted run starts from the events charged before the reboot.""" run = await run_actor(ppe_reboot_actor, max_total_charge_usd=Decimal('0.05')) assert run.status == 'SUCCEEDED' diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index 5bba4746..f60c25cb 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -853,4 +853,7 @@ async def test_persist_state_event_persists_counts(mock_client: MagicMock) -> No event_manager.emit(event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=True)) await event_manager.wait_for_all_listeners_to_complete() - assert await key_value_store.get_value(CHARGING_STATE_KEY) == {'chargedEventCounts': {'search': 2, 'scrape': 2}} + # Checked inside the context, since leaving it emits one more `PERSIST_STATE` event. + assert await key_value_store.get_value(CHARGING_STATE_KEY) == { + 'chargedEventCounts': {'search': 2, 'scrape': 2} + } From 8e23aec79be33e72e834791bf3e1443459110dc9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 15:04:01 +0000 Subject: [PATCH 7/7] fix: restore the charging state attribute and import `ValidationError` Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 4 ++-- tests/unit/actor/test_charging_manager.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 7f95c2c2..501bef10 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Literal, Protocol, TypedDict from cachetools import TTLCache -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, ValidationError from pydantic.alias_generators import to_camel from apify_client._models import ( @@ -219,8 +219,8 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None + self._charging_state: RecoverableState[ChargingState] | None = None """The charged event counts, persisted to the default key-value store in a pay-per-event run on the platform.""" - """The charged event counts, persisted to the default key-value store on the platform.""" self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index f60c25cb..4acf866f 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -857,3 +857,14 @@ async def test_persist_state_event_persists_counts(mock_client: MagicMock) -> No assert await key_value_store.get_value(CHARGING_STATE_KEY) == { 'chargedEventCounts': {'search': 2, 'scrape': 2} } + + +async def test_malformed_persisted_state_is_reported(mock_client: MagicMock) -> None: + """A malformed charging state record in the default key-value store fails the start with a clear error.""" + key_value_store = await KeyValueStore.open() + await key_value_store.set_value(CHARGING_STATE_KEY, {'chargedEventCounts': 'not-a-mapping'}) + + cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client) + with pytest.raises(ValueError, match='is malformed'): + async with cm: + pass