feat(charging): Support the idempotency_key param in Actor.charge()

src/apify/_actor.py

@@ -788,18 +788,21 @@ def get_charging_manager(self) -> ChargingManager:
         return self._charging_manager_implementation
 
     @_ensure_context
-    async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
+    async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult:
         """Charge for a specified number of events - sub-operations of the Actor.
 
         This is relevant only for the pay-per-event pricing model.
 
         Args:
             event_name: Name of the event to be charged for.
             count: Number of events to charge for.
+            idempotency_key: A unique key preventing a retried operation from being charged for twice. A repeat
+                under the same key is not sent to the API and reports the `charged_count` of the original call.
+                A key belongs to a single event and is only remembered for the lifetime of the run.
         """
         # charging_manager.charge() acquires charge_lock internally.
         charging_manager = self.get_charging_manager()
-        return await charging_manager.charge(event_name, count=count)
+        return await charging_manager.charge(event_name, count=count, idempotency_key=idempotency_key)
 
     @overload
     def on(

src/apify/_charging.py

@@ -210,14 +210,17 @@ class ChargingManager(Protocol):
     charge_lock: ReentrantLock
     """Lock to synchronize charge operations. Prevents race conditions between `charge` and `push_data` calls."""
 
-    async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
+    async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult:
         """Charge for a specified number of events - sub-operations of the Actor.
 
         This is relevant only for the pay-per-event pricing model.
 
         Args:
             event_name: Name of the event to be charged for.
             count: Number of events to charge for.
+            idempotency_key: A unique key preventing a retried operation from being charged for twice. A repeat
+                under the same key is not sent to the API and reports the `charged_count` of the original call.
+                A key belongs to a single event and is only remembered for the lifetime of the run.
         """
 
     def calculate_total_charged_amount(self) -> Decimal:
@@ -329,6 +332,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No
         self._charging_state: dict[str, ChargingStateItem] = {}
         self._pricing_info: dict[str, PricingInfoItem] = {}
         self._tier_priced_events: set[str] = set()
+        self._idempotent_charges: dict[str, IdempotentChargeItem] = {}
 
         self._not_ppe_warning_printed = False
         self.active = False
@@ -412,7 +416,10 @@ async def __aexit__(
         self.active = False
 
     @_ensure_context
-    async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
+    async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult:
+        if idempotency_key == '':
+            raise ValueError('idempotency_key must not be an empty string')
+
         # For runs that do not use the pay-per-event pricing model, just print a warning and return
         if self._pricing_model != 'PAY_PER_EVENT':
             if not self._not_ppe_warning_printed:
@@ -435,6 +442,21 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
             )
 
         async with self.charge_lock():
+            # The platform discards a charge repeated under a key it has already seen, so repeating it here would
+            # inflate the local charging state and make the run hit `max_total_charge_usd` sooner than it should.
+            if idempotency_key is not None and (previous := self._idempotent_charges.get(idempotency_key)):
+                if previous.event_name != event_name:
+                    raise ValueError(
+                        f"Idempotency key '{idempotency_key}' was already used to charge for event "
+                        f"'{previous.event_name}', so it cannot be reused for event '{event_name}'."
+                    )
+
+                return ChargeResult(
+                    event_charge_limit_reached=self.is_event_charge_limit_reached(event_name),
+                    charged_count=previous.charged_count,
+                    chargeable_within_limit=self.compute_chargeable(),
+                )
+
             # Determine the maximum amount of events that can be charged within the budget
             max_chargeable = self.calculate_max_event_charge_count_within_limit(event_name)
             charged_count = min(count, max_chargeable if max_chargeable is not None else count)
@@ -455,10 +477,7 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
                 ),
             )
 
-            # Update the charging state
-            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
+            charge_sent = False
 
             # If running on the platform, call the charge endpoint
             if self._is_at_home:
@@ -470,7 +489,12 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
                     # the platform handles them automatically based on dataset writes.
                     pass
                 elif event_name in self._pricing_info:
-                    await self._client.run(self._actor_run_id).charge(event_name, count=charged_count)
+                    await self._client.run(self._actor_run_id).charge(
+                        event_name,
+                        count=charged_count,
+                        idempotency_key=idempotency_key,
+                    )
+                    charge_sent = True
                     logger.debug(f"Charged {charged_count} occurrence(s) of event '{event_name}'.")
                 elif event_name in self._tier_priced_events:
                     logger.warning(
@@ -479,6 +503,19 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
                 else:
                     logger.warning(f"Attempting to charge for an unknown event '{event_name}'")
 
+            # Update the charging state
+            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
+
+            # Only remember a key that stands for a charge the platform actually received. Off the platform there
+            # is no request at all and the registry is the only thing providing deduplication.
+            if idempotency_key is not None and (not self._is_at_home or charge_sent):
+                self._idempotent_charges[idempotency_key] = IdempotentChargeItem(
+                    event_name=event_name,
+                    charged_count=charged_count,
+                )
+
             # Log the charged operation (if enabled)
             if self._charging_log_dataset:
                 await self._charging_log_dataset.push_data(
@@ -487,6 +524,7 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult:
                         'event_title': pricing_info.title,
                         'event_price_usd': float(round(pricing_info.price, 3)),
                         'charged_count': charged_count,
+                        'idempotency_key': idempotency_key,
                         'timestamp': datetime.now(UTC).isoformat(),
                     }
                 )
@@ -636,6 +674,12 @@ class PricingInfoItem:
     title: str
 
 
+@dataclass(frozen=True)
+class IdempotentChargeItem:
+    event_name: str
+    charged_count: int
+
+
 class _FetchedPricingInfoDict(TypedDict):
     pricing_info: ActorPricingInfoModel | None
     charged_event_counts: dict[str, int]

tests/unit/actor/test_actor_charge.py

@@ -37,7 +37,7 @@ async def setup_mocked_charging(
             setup.charging_mgr._pricing_info['event'] = PricingInfoItem(Decimal('1.0'), 'Event')
 
             result = await Actor.charge('event', count=1)
-            setup.mock_charge.assert_called_once_with('event', count=1)
+            setup.mock_charge.assert_called_once_with('event', count=1, idempotency_key=None)
     """
     # Mock the ApifyClientAsync
     mock_client = Mock()
@@ -82,7 +82,7 @@ async def test_actor_charge_push_data_with_no_remaining_budget() -> None:
         result1 = await Actor.charge('some-event', count=1)  # Costs $1, leaving $0.5
 
         # Verify the first charge call was made correctly
-        setup.mock_charge.assert_called_once_with('some-event', count=1)
+        setup.mock_charge.assert_called_once_with('some-event', count=1, idempotency_key=None)
         setup.mock_charge.reset_mock()
 
         assert result1.charged_count == 1
@@ -117,10 +117,27 @@ async def test_actor_charge_api_call_verification() -> None:
 
         # Call charge with count=1 - this SHOULD call the API
         result2 = await Actor.charge('test-event', count=1)
-        setup.mock_charge.assert_called_once_with('test-event', count=1)
+        setup.mock_charge.assert_called_once_with('test-event', count=1, idempotency_key=None)
         assert result2.charged_count == 1
 
 
+async def test_actor_charge_forwards_idempotency_key() -> None:
+    """Verify that Actor.charge passes the idempotency key down to the API and deduplicates repeats."""
+    async with setup_mocked_charging(
+        Configuration(max_total_charge_usd=Decimal('10.0'), test_pay_per_event=True), {'test-event': Decimal('1.0')}
+    ) as setup:
+        result1 = await Actor.charge('test-event', count=1, idempotency_key='key-1')
+        setup.mock_charge.assert_called_once_with('test-event', count=1, idempotency_key='key-1')
+        assert result1.charged_count == 1
+
+        setup.mock_charge.reset_mock()
+
+        result2 = await Actor.charge('test-event', count=1, idempotency_key='key-1')
+        setup.mock_charge.assert_not_called()
+        assert result2.charged_count == 1
+        assert setup.charging_mgr.get_charged_event_count('test-event') == 1
+
+
 async def test_max_event_charge_count_within_limit_tolerates_overdraw() -> None:
     """Test that calculate_max_event_charge_count_within_limit does not return nonsensical (e.g., negative) values when
     the total number of charged events overdraws the max_total_charge_usd limit."""

tests/unit/actor/test_charging_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from decimal import Decimal
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
@@ -466,3 +467,132 @@ async def test_compute_chargeable_updates_after_charge(mock_client: MagicMock) -
         # $6.00 remaining: search=$1.00 → 6, scrape=$2.00 → 3
         assert chargeable['search'] == 6
         assert chargeable['scrape'] == 3
+
+
+async def test_charge_forwards_idempotency_key_to_client(mock_client: MagicMock) -> None:
+    """Test that the idempotency key is passed through to the API client."""
+    pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')})
+    config = _make_config(
+        is_at_home=True,
+        actor_run_id='test-run-id',
+        actor_pricing_info=pricing_info,
+        charged_event_counts={},
+        max_total_charge_usd=Decimal('10.00'),
+    )
+    cm = ChargingManagerImplementation(config, mock_client)
+    async with cm:
+        await cm.charge('search', count=2, idempotency_key='key-1')
+        mock_client.run.return_value.charge.assert_awaited_once_with('search', count=2, idempotency_key='key-1')
+
+        # Without a key the client gets None and generates a unique one per call, so nothing is deduplicated.
+        mock_client.run.return_value.charge.reset_mock()
+        await cm.charge('search', count=1)
+        mock_client.run.return_value.charge.assert_awaited_once_with('search', count=1, idempotency_key=None)
+
+
+async def test_charge_deduplicates_repeated_idempotency_key(mock_client: MagicMock) -> None:
+    """Test that a repeated key skips both the API call and the local charging state update."""
+    pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')})
+    config = _make_config(
+        is_at_home=True,
+        actor_run_id='test-run-id',
+        actor_pricing_info=pricing_info,
+        charged_event_counts={},
+        max_total_charge_usd=Decimal('10.00'),
+    )
+    cm = ChargingManagerImplementation(config, mock_client)
+    async with cm:
+        first = await cm.charge('search', count=2, idempotency_key='key-1')
+        assert first.charged_count == 2
+
+        # The repeat reports what was charged under this key originally, whatever count it asks for.
+        repeated = await cm.charge('search', count=1, idempotency_key='key-1')
+        assert repeated.charged_count == 2
+
+        assert cm.get_charged_event_count('search') == 2
+        assert mock_client.run.return_value.charge.await_count == 1
+
+        # A distinct key is charged as usual.
+        await cm.charge('search', count=1, idempotency_key='key-2')
+        assert cm.get_charged_event_count('search') == 3
+        assert mock_client.run.return_value.charge.await_count == 2
+
+
+async def test_charge_does_not_register_key_for_uncharged_event(mock_client: MagicMock) -> None:
+    """Test that an event that never reached the API does not consume the idempotency key."""
+    pricing_info = PayPerEventActorPricingInfo.model_validate(
+        {
+            'pricingModel': 'PAY_PER_EVENT',
+            'pricingPerEvent': {
+                'actorChargeEvents': {
+                    'search': {'eventPriceUsd': 1.00, 'eventTitle': 'Search event'},
+                    'tiered': {'eventTieredPricingUsd': {}, 'eventTitle': 'Tiered event'},
+                }
+            },
+        }
+    )
+    config = _make_config(
+        is_at_home=True,
+        actor_run_id='test-run-id',
+        actor_pricing_info=pricing_info,
+        charged_event_counts={},
+        max_total_charge_usd=Decimal('10.00'),
+    )
+    cm = ChargingManagerImplementation(config, mock_client)
+    async with cm:
+        # Neither a tier-priced nor an unknown event is chargeable via the API.
+        await cm.charge('tiered', count=1, idempotency_key='key-1')
+        await cm.charge('typo-event', count=1, idempotency_key='key-2')
+        mock_client.run.return_value.charge.assert_not_awaited()
+
+        # Both keys are still free, so a corrected charge under either one reaches the platform.
+        assert (await cm.charge('search', count=1, idempotency_key='key-1')).charged_count == 1
+        assert (await cm.charge('search', count=1, idempotency_key='key-2')).charged_count == 1
+        assert mock_client.run.return_value.charge.await_count == 2
+
+
+async def test_charge_rejects_invalid_idempotency_key(mock_client: MagicMock) -> None:
+    """Test that an empty key and a key reused for another event are both refused."""
+    pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00'), 'scrape': Decimal('2.00')})
+    config = _make_config(
+        is_at_home=True,
+        actor_run_id='test-run-id',
+        actor_pricing_info=pricing_info,
+        charged_event_counts={},
+        max_total_charge_usd=Decimal('10.00'),
+    )
+    cm = ChargingManagerImplementation(config, mock_client)
+    async with cm:
+        with pytest.raises(ValueError, match='must not be an empty string'):
+            await cm.charge('search', count=1, idempotency_key='')
+
+        await cm.charge('search', count=1, idempotency_key='key-1')
+
+        with pytest.raises(ValueError, match='cannot be reused for event'):
+            await cm.charge('scrape', count=1, idempotency_key='key-1')
+
+        # The refused charge must leave no trace behind.
+        assert cm.get_charged_event_count('scrape') == 0
+        assert cm.calculate_total_charged_amount() == Decimal('1.00')
+        assert mock_client.run.return_value.charge.await_count == 1
+
+
+async def test_concurrent_charges_under_one_key_charge_once(mock_client: MagicMock) -> None:
+    """Test that concurrent charges under one key result in a single charge."""
+    pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')})
+    config = _make_config(
+        is_at_home=True,
+        actor_run_id='test-run-id',
+        actor_pricing_info=pricing_info,
+        charged_event_counts={},
+        max_total_charge_usd=Decimal('10.00'),
+    )
+    cm = ChargingManagerImplementation(config, mock_client)
+    async with cm:
+        results = await asyncio.gather(
+            *(cm.charge('search', count=1, idempotency_key='key-1') for _ in range(5)),
+        )
+
+        assert [result.charged_count for result in results] == [1] * 5
+        assert cm.get_charged_event_count('search') == 1
+        assert mock_client.run.return_value.charge.await_count == 1