fix(idempotency): clear the in progress record when the response can't be serialized

aws_lambda_powertools/utilities/idempotency/base.py

@@ -296,7 +296,7 @@ def _get_function_response(self):
 
         else:
             try:
-                serialized_response: dict = self.output_serializer.to_dict(response) if response is not None else None
+                serialized_response: dict = self._serialize_response(response)
                 self.persistence_store.save_success(data=self.data, result=serialized_response)
             except Exception as save_exception:
                 raise IdempotencyPersistenceLayerError(
@@ -305,3 +305,14 @@ def _get_function_response(self):
                 ) from save_exception
 
         return response
+
+    def _serialize_response(self, response: Any) -> Any:
+        if response is None:
+            return None
+        try:
+            return self.output_serializer.to_dict(response)
+        except Exception as serialize_exception:
+            # Nothing reached the persistence store yet, so clear the in progress record
+            # to let the next call with the same payload run the function again
+            self.persistence_store.delete_record(data=self.data, exception=serialize_exception)
+            raise

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -299,7 +299,14 @@ def save_success(self, data: dict[str, Any], result: dict) -> None:
             # See: https://github.com/aws-powertools/powertools-lambda-python/issues/2465
             return None
 
-        response_data = json.dumps(result, cls=Encoder, sort_keys=True)
+        try:
+            response_data = json.dumps(result, cls=Encoder, sort_keys=True)
+        except Exception as exc:
+            # Nothing was written yet, so clear the in progress record to let the next call
+            # with the same payload run the function again instead of failing as in progress.
+            # A failure while writing the record below keeps it, as the write may have landed.
+            self.delete_record(data=data, exception=exc)
+            raise
 
         data_record = DataRecord(
             idempotency_key=idempotency_key,

docs/utilities/idempotency.md

@@ -409,12 +409,16 @@ To prevent extended failures, use **`register_lambda_context`** function from yo
 
 ### Handling exceptions
 
-There are two failure modes that can cause new invocations to execute your code again despite having the same payload:
+There are three failure modes that can cause new invocations to execute your code again despite having the same payload:
 
 * **Unhandled exception**. We catch them to delete the idempotency record to prevent inconsistencies, then propagate them.
 * **Persistent layer errors**. We raise **`IdempotencyPersistenceLayerError`** for any persistence layer errors _e.g., remove idempotency record_.
+* **Response that can't be serialized**. When your function returns a value the output serializer or JSON can't convert _e.g., `datetime`, `set`, `bytes`_, we try to delete the in progress record, then raise **`IdempotencyPersistenceLayerError`**.
+  The next invocation runs your code again once that record is deleted, or after it expires if the deletion fails.
 
-If an exception is handled or raised **outside** your decorated function, then idempotency will be maintained.
+If saving the response fails after it was serialized, we keep the record, as the write could have landed. New invocations return the saved response or wait until the in progress record expires.
+
+Apart from a response that can't be serialized, if an exception is handled or raised **outside** your decorated function, then idempotency will be maintained.
 
 ```python title="working_with_exceptions.py" hl_lines="21 32 38"
 --8<-- "examples/idempotency/src/working_with_exceptions.py"

examples/idempotency/src/working_with_exceptions.py

@@ -18,7 +18,7 @@
 
 @idempotent_function(data_keyword_argument="data", config=config, persistence_store=persistence_layer)
 def call_external_service(data: dict):
-    # Any exception raised will lead to idempotency record to be deleted
+    # Any exception raised, or a response that can't be serialized, will lead to idempotency record to be deleted
     result: requests.Response = requests.post(
         "https://jsonplaceholder.typicode.com/comments/",
         json=data,

tests/functional/idempotency/_boto3/test_idempotency.py

@@ -614,6 +614,137 @@ def lambda_handler(event, context):
     stubber.deactivate()
 
 
+@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
+def test_idempotent_lambda_response_not_json_serializable_clears_in_progress_record(
+    idempotency_config: IdempotencyConfig,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    hashed_idempotency_key,
+    expected_params_put_item,
+    lambda_context,
+):
+    """
+    Test idempotent decorator when lambda_handler returns a response that can't be serialized to JSON.
+    The in progress record must be cleared so a retry with the same event runs lambda_handler again.
+    """
+    stubber = stub.Stubber(persistence_store.client)
+
+    ddb_response = {}
+    expected_params_delete_item = {"TableName": TABLE_NAME, "Key": {"id": {"S": hashed_idempotency_key}}}
+
+    # Each invocation saves the in progress record and clears it after failing to serialize the response
+    for _ in range(2):
+        stubber.add_response("put_item", ddb_response, expected_params_put_item)
+        stubber.add_response("delete_item", ddb_response, expected_params_delete_item)
+    stubber.activate()
+
+    handler_calls = 0
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        nonlocal handler_calls
+        handler_calls += 1
+        return {"created": dt.datetime(2026, 1, 1)}
+
+    for _ in range(2):
+        with pytest.raises(IdempotencyPersistenceLayerError) as exc:
+            lambda_handler(lambda_apigw_event, lambda_context)
+
+        assert exc.value.args[0] == "Failed to update record state to success in idempotency store"
+        assert isinstance(exc.value.__cause__, TypeError)
+
+    assert handler_calls == 2
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+
+@pytest.mark.parametrize("use_local_cache", [False, True])
+def test_idempotent_function_output_serializer_failure_clears_in_progress_record(
+    use_local_cache: bool,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_context,
+):
+    """
+    Test idempotent_function when the output serializer fails to convert the function response.
+    The in progress record must be cleared so a retry with the same payload runs the function again.
+    """
+    mock_event = {"data": "value"}
+    function_qualified_name = "test_idempotent_function_output_serializer_failure_clears_in_progress_record.<locals>"
+    idempotency_key = (
+        f"{TESTS_MODULE_PREFIX}.{function_qualified_name}.record_handler#{hash_idempotency_key(mock_event)}"
+    )
+    stubber = stub.Stubber(persistence_store.client)
+
+    ddb_response = {}
+    expected_params_put_item = build_idempotency_put_item_stub(
+        data=mock_event,
+        function_qualified_name=function_qualified_name,
+        handler_name="record_handler",
+    )
+    expected_params_delete_item = {"TableName": TABLE_NAME, "Key": {"id": {"S": idempotency_key}}}
+
+    stubber.add_response("put_item", ddb_response, expected_params_put_item)
+    stubber.add_response("delete_item", ddb_response, expected_params_delete_item)
+    stubber.activate()
+
+    def failing_to_dict(data: Any) -> dict:
+        raise ValueError("cannot convert response")
+
+    output_serializer = CustomDictSerializer(to_dict=failing_to_dict, from_dict=lambda data: data)
+    config = IdempotencyConfig(use_local_cache=use_local_cache)
+    config.register_lambda_context(lambda_context)
+
+    @idempotent_function(
+        config=config,
+        persistence_store=persistence_store,
+        data_keyword_argument="record",
+        output_serializer=output_serializer,
+    )
+    def record_handler(record):
+        return {"message": "success!"}
+
+    with pytest.raises(IdempotencyPersistenceLayerError) as exc:
+        record_handler(record=mock_event)
+
+    assert exc.value.args[0] == "Failed to update record state to success in idempotency store"
+    assert isinstance(exc.value.__cause__, ValueError)
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+
+@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}], indirect=True)
+def test_idempotent_lambda_response_not_json_serializable_delete_failure(
+    idempotency_config: IdempotencyConfig,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    hashed_idempotency_key,
+    expected_params_put_item,
+    lambda_context,
+):
+    """
+    Test idempotent decorator when lambda_handler returns a response that can't be serialized to JSON and
+    clearing the in progress record fails too. The persistence layer error still surfaces to the caller.
+    """
+    stubber = stub.Stubber(persistence_store.client)
+
+    expected_params_delete_item = {"TableName": TABLE_NAME, "Key": {"id": {"S": hashed_idempotency_key}}}
+
+    stubber.add_response("put_item", {}, expected_params_put_item)
+    stubber.add_client_error("delete_item", "UnrecoverableError", expected_params=expected_params_delete_item)
+    stubber.activate()
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        return {"created": dt.datetime(2026, 1, 1)}
+
+    with pytest.raises(IdempotencyPersistenceLayerError) as exc:
+        lambda_handler(lambda_apigw_event, lambda_context)
+
+    assert exc.value.args[0] == "Failed to update record state to success in idempotency store"
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+
 @pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
 def test_idempotent_persistence_exception_getting(
     idempotency_config: IdempotencyConfig,

tests/functional/idempotency/_redis/test_redis_layer.py

@@ -17,6 +17,7 @@
     IdempotencyPersistenceConfigError,
     IdempotencyPersistenceConnectionError,
     IdempotencyPersistenceConsistencyError,
+    IdempotencyPersistenceLayerError,
     IdempotencyValidationError,
 )
 from aws_lambda_powertools.utilities.idempotency.idempotency import (
@@ -712,6 +713,91 @@ def lambda_handler(event, _):
     assert handler_result2 == result
 
 
+def test_idempotent_function_redis_response_not_json_serializable_runs_again(
+    persistence_store_standalone_redis: RedisCachePersistenceLayer,
+):
+    # GIVEN a function whose response can't be serialized to JSON
+    function_calls = 0
+
+    @idempotent_function(persistence_store=persistence_store_standalone_redis, data_keyword_argument="order")
+    def process_order(order):
+        nonlocal function_calls
+        function_calls += 1
+        return {"created": datetime.datetime(2026, 1, 1)}
+
+    # WHEN calling it twice with the same payload
+    # THEN both calls fail to save the response, and the second call runs the function again
+    # instead of finding the first call's record still in progress
+    for _ in range(2):
+        with pytest.raises(IdempotencyPersistenceLayerError):
+            process_order(order={"id": 1})
+
+    assert function_calls == 2
+
+
+class MockRedisFailingCompletedWrite(MockRedis):
+    # fails the write that stores the COMPLETED record, after or before it lands
+    def __init__(self, write_lands: bool, **kwargs):
+        self.write_lands = write_lands
+        super().__init__(**kwargs)
+
+    def set(self, name, value, ex: int = 0, nx: bool = False):
+        if not nx and STATUS_CONSTANTS["COMPLETED"] in str(value):
+            if self.write_lands:
+                super().set(name=name, value=value, ex=ex, nx=nx)
+            raise self.exceptions.RedisError
+        return super().set(name=name, value=value, ex=ex, nx=nx)
+
+
+@pytest.mark.parametrize("use_local_cache", [False, True])
+def test_idempotent_function_redis_completed_write_landed_then_failed_keeps_record(use_local_cache):
+    # GIVEN a store that saves the COMPLETED record but then fails the write call
+    persistence_store = CachePersistenceLayer(
+        client=MockRedisFailingCompletedWrite(write_lands=True, host="localhost", port="63005"),
+    )
+    config = IdempotencyConfig(use_local_cache=use_local_cache)
+    function_calls = 0
+
+    @idempotent_function(persistence_store=persistence_store, config=config, data_keyword_argument="order")
+    def process_order(order):
+        nonlocal function_calls
+        function_calls += 1
+        return {"order_id": order["id"]}
+
+    # WHEN the first call fails to confirm the write
+    with pytest.raises(IdempotencyPersistenceLayerError):
+        process_order(order={"id": 1})
+
+    # THEN the stored result is kept, and the next call returns it without running the function again
+    assert process_order(order={"id": 1}) == {"order_id": 1}
+    assert function_calls == 1
+
+
+@pytest.mark.parametrize("use_local_cache", [False, True])
+def test_idempotent_function_redis_completed_write_failed_keeps_in_progress_record(use_local_cache):
+    # GIVEN a store that fails the write of the COMPLETED record before saving it
+    persistence_store = CachePersistenceLayer(
+        client=MockRedisFailingCompletedWrite(write_lands=False, host="localhost", port="63005"),
+    )
+    config = IdempotencyConfig(use_local_cache=use_local_cache)
+    function_calls = 0
+
+    @idempotent_function(persistence_store=persistence_store, config=config, data_keyword_argument="order")
+    def process_order(order):
+        nonlocal function_calls
+        function_calls += 1
+        return {"order_id": order["id"]}
+
+    # WHEN the first call fails to save the response
+    with pytest.raises(IdempotencyPersistenceLayerError):
+        process_order(order={"id": 1})
+
+    # THEN the in progress record is kept, as the write could have landed
+    with pytest.raises(IdempotencyAlreadyInProgressError):
+        process_order(order={"id": 1})
+    assert function_calls == 1
+
+
 def test_redis_orphan_record_race_condition(lambda_context):
     redis_client = MockRedis(
         host="localhost",