Support replays in Idempotency utility

Durable Functions introduces the concept of function replays.
Previously, the idempotency utility would throw an
"IdempotencyItemAlreadyExistsError", as the replay has the same payload
of the initial execution. It appears as a duplicate, so is rejected.
Now, a replay is allowed

Author: ConnorKirk

aws_lambda_powertools/utilities/idempotency/base.py

@@ -123,10 +123,18 @@ def __init__(
 
         self.persistence_store = persistence_store
 
-    def handle(self) -> Any:
+    def handle(self, durable_mode: str | None = None) -> Any:
         """
         Main entry point for handling idempotent execution of a function.
 
+        Parameters
+        ----------
+        durable_mode : str | None, optional
+            Mode for handling in-progress executions. Options:
+            - "REPLAY_MODE": Allow replay of functions that are already in progress
+            - "EXECUTION_MODE": Standard durable execution mode
+            - None: Standard idempotency behavior (raises IdempotencyAlreadyInProgressError)
+
         Returns
         -------
         Any
@@ -138,12 +146,12 @@ def handle(self) -> Any:
         # In most cases we can retry successfully on this exception.
         for i in range(MAX_RETRIES + 1):  # pragma: no cover
             try:
-                return self._process_idempotency()
+                return self._process_idempotency(durable_mode)
             except IdempotencyInconsistentStateError:
                 if i == MAX_RETRIES:
                     raise  # Bubble up when exceeded max tries
 
-    def _process_idempotency(self):
+    def _process_idempotency(self, durable_mode: str | None):
         try:
             # We call save_inprogress first as an optimization for the most common case where no idempotent record
             # already exists. If it succeeds, there's no need to call get_record.
@@ -159,7 +167,8 @@ def _process_idempotency(self):
             # We give preference to ReturnValuesOnConditionCheckFailure because it is a faster and more cost-effective
             # way of retrieving the existing record after a failed conditional write operation.
             record = exc.old_data_record or self._get_idempotency_record()
-
+            if durable_mode == "REPLAY_MODE":
+                return self._get_function_response()
             # If a record is found, handle it for status
             if record:
                 return self._handle_for_status(record)

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -28,7 +28,7 @@
     from aws_lambda_powertools.utilities.idempotency.persistence.base import (
         BasePersistenceLayer,
     )
-    from aws_lambda_powertools.utilities.typing import LambdaContext
+    from aws_lambda_powertools.utilities.typing import DurableContext, LambdaContext
 
 from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
@@ -37,9 +37,9 @@
 
 @lambda_handler_decorator
 def idempotent(
-    handler: Callable[[Any, LambdaContext], Any],
+    handler: Callable[[Any, LambdaContext | DurableContext], Any],
     event: dict[str, Any],
-    context: LambdaContext,
+    context: LambdaContext | DurableContext,
     persistence_store: BasePersistenceLayer,
     config: IdempotencyConfig | None = None,
     key_prefix: str | None = None,
@@ -55,7 +55,7 @@ def idempotent(
     event: dict
         Lambda's Event
     context: dict
-        Lambda's Context
+        Lambda's Context or Durable Context
     persistence_store: BasePersistenceLayer
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
@@ -91,7 +91,15 @@ def handler(event, context):
         return handler(event, context, **kwargs)
 
     config = config or IdempotencyConfig()
-    config.register_lambda_context(context)
+
+    if hasattr(context, "state"):
+        # Extract lambda_context from DurableContext for idempotency tracking
+        config.register_lambda_context(context.lambda_context)
+        durable_mode = "REPLAY_MODE" if len(context.state.operations) > 1 else "EXECUTION_MODE"
+    else:
+        # Standard LambdaContext
+        config.register_lambda_context(context)
+        durable_mode = None
 
     args = event, context
     idempotency_handler = IdempotencyHandler(
@@ -104,7 +112,7 @@ def handler(event, context):
         function_kwargs=kwargs,
     )
 
-    return idempotency_handler.handle()
+    return idempotency_handler.handle(durable_mode=durable_mode)
 
 
 def idempotent_function(
@@ -193,6 +201,10 @@ def decorate(*args, **kwargs):
                 f" Ensure this exists in your function's signature as well as the caller used it as a keyword argument",
             )
 
+        durable_mode = None
+        if len(args) >= 2 and hasattr(args[1], "state"):
+            durable_mode = "REPLAY_MODE" if len(args[1].state.operations) > 1 else "EXECUTION_MODE"
+
         payload = kwargs.get(data_keyword_argument)
 
         idempotency_handler = IdempotencyHandler(
@@ -206,6 +218,6 @@ def decorate(*args, **kwargs):
             function_kwargs=kwargs,
         )
 
-        return idempotency_handler.handle()
+        return idempotency_handler.handle(durable_mode=durable_mode)
 
     return cast(AnyCallableT, decorate)

aws_lambda_powertools/utilities/typing/__init__.py

@@ -4,6 +4,6 @@
     [`Typing`](../utilities/typing.md)
 """
 
-from .lambda_context import LambdaContext
+from .lambda_context import DurableContext, LambdaContext
 
-__all__ = ["LambdaContext"]
+__all__ = ["DurableContext", "LambdaContext"]

aws_lambda_powertools/utilities/typing/lambda_context.py

@@ -93,3 +93,16 @@ def tenant_id(self) -> str | None:
     def get_remaining_time_in_millis() -> int:
         """Returns the number of milliseconds left before the execution times out."""
         return 0
+
+
+class DurableContext:
+    _lambda_context: LambdaContext
+    _state: object
+
+    @property
+    def lambda_context(self) -> LambdaContext:
+        return self._lambda_context
+
+    @property
+    def state(self) -> object:
+        return self._state