From dc316fdbbe2260c429249c6fd73704160e254f0a Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Fri, 29 Mar 2024 01:21:10 +1100
Subject: [PATCH 01/10] feat(idempotent-response-manipulation): Added
 capability of providing an IdempotentHook functiont to be called when an
 idempotent response is being returned.

---
 .../utilities/idempotency/__init__.py         |   3 +
 .../utilities/idempotency/base.py             |   8 +-
 .../utilities/idempotency/config.py           |  43 ++++-
 .../idempotency/test_idempotency.py           | 156 ++++++++++++++----
 4 files changed, 174 insertions(+), 36 deletions(-)

diff --git a/aws_lambda_powertools/utilities/idempotency/__init__.py b/aws_lambda_powertools/utilities/idempotency/__init__.py
index ae27330cc1f..77ab4862322 100644
--- a/aws_lambda_powertools/utilities/idempotency/__init__.py
+++ b/aws_lambda_powertools/utilities/idempotency/__init__.py
@@ -2,6 +2,7 @@
 Utility for adding idempotency to lambda functions
 """
 
+from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData, IdempotentHookFunction
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
@@ -17,4 +18,6 @@
     "idempotent",
     "idempotent_function",
     "IdempotencyConfig",
+    "IdempotentHookData",
+    "IdempotentHookFunction",
 )
diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py
index 771547fe33c..86d9e3cd2ea 100644
--- a/aws_lambda_powertools/utilities/idempotency/base.py
+++ b/aws_lambda_powertools/utilities/idempotency/base.py
@@ -3,7 +3,10 @@
 from copy import deepcopy
 from typing import Any, Callable, Dict, Optional, Tuple
 
-from aws_lambda_powertools.utilities.idempotency.config import IdempotencyConfig
+from aws_lambda_powertools.utilities.idempotency.config import (
+    IdempotencyConfig,
+    IdempotentHookData,
+)
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
@@ -227,6 +230,9 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
             )
         response_dict: Optional[dict] = data_record.response_json_as_dict()
         if response_dict is not None:
+            if self.config.response_hook is not None:
+                idempotent_data = IdempotentHookData(data_record)
+                return self.output_serializer.from_dict(self.config.response_hook(response_dict, idempotent_data))
             return self.output_serializer.from_dict(response_dict)
         return None
 
diff --git a/aws_lambda_powertools/utilities/idempotency/config.py b/aws_lambda_powertools/utilities/idempotency/config.py
index e78f339fdc9..1dcfda38c04 100644
--- a/aws_lambda_powertools/utilities/idempotency/config.py
+++ b/aws_lambda_powertools/utilities/idempotency/config.py
@@ -1,8 +1,45 @@
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
+from aws_lambda_powertools.shared.types import Protocol
+from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
+class IdempotentHookData:
+    """
+    Idempotent Hook Data
+
+    Contains data relevant to the current Idempotent record which matches the current request.
+    All IdempotentHook functions will be passed this data as well as the current Response.
+    """
+
+    def __init__(self, data_record: DataRecord) -> None:
+        self._idempotency_key = data_record.idempotency_key
+        self._status = data_record.status
+        self._expiry_timestamp = data_record.expiry_timestamp
+
+    @property
+    def idempotency_key(self) -> str:
+        return self._idempotency_key
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    @property
+    def expiry_timestamp(self) -> int:
+        return self._expiry_timestamp
+
+
+class IdempotentHookFunction(Protocol):
+    """
+    The IdempotentHookFunction.
+    This class defines the calling signature for IdempotentHookFunction callbacks.
+    """
+
+    def __call__(self, response: Any, idempotent_data: IdempotentHookData): ...
+
+
 class IdempotencyConfig:
     def __init__(
         self,
@@ -15,6 +52,7 @@ def __init__(
         local_cache_max_items: int = 256,
         hash_function: str = "md5",
         lambda_context: Optional[LambdaContext] = None,
+        response_hook: IdempotentHookFunction = None,
     ):
         """
         Initialize the base persistence layer
@@ -37,6 +75,8 @@ def __init__(
             Function to use for calculating hashes, by default md5.
         lambda_context: LambdaContext, optional
             Lambda Context containing information about the invocation, function and execution environment.
+        response_hook: Callable, optional
+            Hook function to be called when a response is idempotent.
         """
         self.event_key_jmespath = event_key_jmespath
         self.payload_validation_jmespath = payload_validation_jmespath
@@ -47,6 +87,7 @@ def __init__(
         self.local_cache_max_items = local_cache_max_items
         self.hash_function = hash_function
         self.lambda_context: Optional[LambdaContext] = lambda_context
+        self.response_hook: IdempotentHookFunction = response_hook
 
     def register_lambda_context(self, lambda_context: LambdaContext):
         """Captures the Lambda context, to calculate the remaining time before the invocation times out"""
diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py
index 2591cf8e043..abc45bb4104 100644
--- a/tests/functional/idempotency/test_idempotency.py
+++ b/tests/functional/idempotency/test_idempotency.py
@@ -1,7 +1,8 @@
 import copy
 import datetime
 import warnings
-from unittest.mock import MagicMock
+from typing import Any
+from unittest.mock import MagicMock, Mock
 
 import jmespath
 import pytest
@@ -26,6 +27,7 @@
     IdempotencyHandler,
     _prepare_data,
 )
+from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
@@ -240,6 +242,39 @@ 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_expired(
+    idempotency_config: IdempotencyConfig,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    lambda_response,
+    expected_params_update_item,
+    expected_params_put_item,
+    lambda_context,
+):
+    """
+    Test idempotent decorator when lambda is called with an event it successfully handled already, but outside of the
+    expiry window
+    """
+
+    stubber = stub.Stubber(persistence_store.client)
+
+    ddb_response = {}
+
+    stubber.add_response("put_item", ddb_response, expected_params_put_item)
+    stubber.add_response("update_item", ddb_response, expected_params_update_item)
+    stubber.activate()
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        return lambda_response
+
+    lambda_handler(lambda_apigw_event, lambda_context)
+
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+
 @pytest.mark.parametrize("idempotency_config", [{"use_local_cache": True}], indirect=True)
 def test_idempotent_lambda_first_execution_cached(
     idempotency_config: IdempotencyConfig,
@@ -324,39 +359,6 @@ 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_expired(
-    idempotency_config: IdempotencyConfig,
-    persistence_store: DynamoDBPersistenceLayer,
-    lambda_apigw_event,
-    lambda_response,
-    expected_params_update_item,
-    expected_params_put_item,
-    lambda_context,
-):
-    """
-    Test idempotent decorator when lambda is called with an event it successfully handled already, but outside of the
-    expiry window
-    """
-
-    stubber = stub.Stubber(persistence_store.client)
-
-    ddb_response = {}
-
-    stubber.add_response("put_item", ddb_response, expected_params_put_item)
-    stubber.add_response("update_item", ddb_response, expected_params_update_item)
-    stubber.activate()
-
-    @idempotent(config=idempotency_config, persistence_store=persistence_store)
-    def lambda_handler(event, context):
-        return lambda_response
-
-    lambda_handler(lambda_apigw_event, lambda_context)
-
-    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_lambda_exception(
     idempotency_config: IdempotencyConfig,
@@ -1986,3 +1988,89 @@ def lambda_handler(event, context):
 
     # THEN we should not cache a transaction that failed validation
     assert cache_spy.call_count == 0
+
+
+@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
+def test_responsehook_lambda_first_execution(
+    idempotency_config: IdempotencyConfig,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    expected_params_update_item,
+    expected_params_put_item,
+    lambda_response,
+    lambda_context,
+):
+    """
+    Test response_hook is not called for the idempotent decorator when lambda is executed
+    with an event with a previously unknown event key
+    """
+
+    idempotent_response_hook = Mock()
+
+    stubber = stub.Stubber(persistence_store.client)
+    ddb_response = {}
+
+    stubber.add_response("put_item", ddb_response, expected_params_put_item)
+    stubber.add_response("update_item", ddb_response, expected_params_update_item)
+    stubber.activate()
+
+    idempotency_config.response_hook = idempotent_response_hook
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        return lambda_response
+
+    lambda_handler(lambda_apigw_event, lambda_context)
+
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+    assert not idempotent_response_hook.called
+
+
+@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
+def test_idempotent_lambda_already_completed_response_hook_is_called(
+    idempotency_config: IdempotencyConfig,
+    persistence_store: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    timestamp_future,
+    hashed_idempotency_key,
+    serialized_lambda_response,
+    deserialized_lambda_response,
+    lambda_context,
+):
+    """
+    Test idempotent decorator where event with matching event key has already been successfully processed
+    """
+
+    def idempotent_response_hook(response: Any, idempotent_data: IdempotentHookData):
+        """Modify the response provided by adding a new key"""
+        response["idempotent_response"] = True
+
+        return response
+
+    idempotency_config.response_hook = idempotent_response_hook
+
+    stubber = stub.Stubber(persistence_store.client)
+    ddb_response = {
+        "Item": {
+            "id": {"S": hashed_idempotency_key},
+            "expiration": {"N": timestamp_future},
+            "data": {"S": serialized_lambda_response},
+            "status": {"S": "COMPLETED"},
+        },
+    }
+    stubber.add_client_error("put_item", "ConditionalCheckFailedException", modeled_fields=ddb_response)
+    stubber.activate()
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        raise Exception
+
+    lambda_resp = lambda_handler(lambda_apigw_event, lambda_context)
+
+    # Then idempotent_response value will be added to the response
+    assert lambda_resp["idempotent_response"]
+
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()

From df633f668a315c6af79ac8151cbcbf6c6ce03a0f Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Fri, 29 Mar 2024 22:48:45 +1100
Subject: [PATCH 02/10] chore(mypy): resolve myopy static typing issues, make
 response+hook properly optional

---
 aws_lambda_powertools/utilities/idempotency/config.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/aws_lambda_powertools/utilities/idempotency/config.py b/aws_lambda_powertools/utilities/idempotency/config.py
index 1dcfda38c04..6e1d6c7c423 100644
--- a/aws_lambda_powertools/utilities/idempotency/config.py
+++ b/aws_lambda_powertools/utilities/idempotency/config.py
@@ -27,7 +27,7 @@ def status(self) -> str:
         return self._status
 
     @property
-    def expiry_timestamp(self) -> int:
+    def expiry_timestamp(self) -> Optional[int]:
         return self._expiry_timestamp
 
 
@@ -52,7 +52,7 @@ def __init__(
         local_cache_max_items: int = 256,
         hash_function: str = "md5",
         lambda_context: Optional[LambdaContext] = None,
-        response_hook: IdempotentHookFunction = None,
+        response_hook: Optional[IdempotentHookFunction] = None,
     ):
         """
         Initialize the base persistence layer
@@ -87,7 +87,7 @@ def __init__(
         self.local_cache_max_items = local_cache_max_items
         self.hash_function = hash_function
         self.lambda_context: Optional[LambdaContext] = lambda_context
-        self.response_hook: IdempotentHookFunction = response_hook
+        self.response_hook: Optional[IdempotentHookFunction] = response_hook
 
     def register_lambda_context(self, lambda_context: LambdaContext):
         """Captures the Lambda context, to calculate the remaining time before the invocation times out"""

From 736fdde01b9fee0a57b8bb8ea0fce71d3fc7346c Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Sun, 31 Mar 2024 13:51:52 +1100
Subject: [PATCH 03/10] feat(response_hook): added some documentation, call
 response_hook after custom de-serialization

---
 .../utilities/idempotency/base.py             |  6 ++-
 .../utilities/idempotency/config.py           |  6 +--
 docs/utilities/idempotency.md                 | 39 +++++++++++++------
 .../src/working_with_response_hook.py         | 27 +++++++++++++
 .../idempotency/test_idempotency.py           |  2 +-
 5 files changed, 63 insertions(+), 17 deletions(-)
 create mode 100644 examples/idempotency/src/working_with_response_hook.py

diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py
index 86d9e3cd2ea..9ab5def6f11 100644
--- a/aws_lambda_powertools/utilities/idempotency/base.py
+++ b/aws_lambda_powertools/utilities/idempotency/base.py
@@ -231,8 +231,10 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
         response_dict: Optional[dict] = data_record.response_json_as_dict()
         if response_dict is not None:
             if self.config.response_hook is not None:
-                idempotent_data = IdempotentHookData(data_record)
-                return self.output_serializer.from_dict(self.config.response_hook(response_dict, idempotent_data))
+                return self.config.response_hook(
+                    self.output_serializer.from_dict(response_dict),
+                    IdempotentHookData(data_record),
+                )
             return self.output_serializer.from_dict(response_dict)
         return None
 
diff --git a/aws_lambda_powertools/utilities/idempotency/config.py b/aws_lambda_powertools/utilities/idempotency/config.py
index 6e1d6c7c423..aa11793bd8a 100644
--- a/aws_lambda_powertools/utilities/idempotency/config.py
+++ b/aws_lambda_powertools/utilities/idempotency/config.py
@@ -37,7 +37,7 @@ class IdempotentHookFunction(Protocol):
     This class defines the calling signature for IdempotentHookFunction callbacks.
     """
 
-    def __call__(self, response: Any, idempotent_data: IdempotentHookData): ...
+    def __call__(self, response: Any, idempotent_data: IdempotentHookData) -> Any: ...
 
 
 class IdempotencyConfig:
@@ -75,8 +75,8 @@ def __init__(
             Function to use for calculating hashes, by default md5.
         lambda_context: LambdaContext, optional
             Lambda Context containing information about the invocation, function and execution environment.
-        response_hook: Callable, optional
-            Hook function to be called when a response is idempotent.
+        response_hook: IdempotentHookFunction, optional
+            Hook function to be called when an idempotent response is returned from the idempotent store.
         """
         self.event_key_jmespath = event_key_jmespath
         self.payload_validation_jmespath = payload_validation_jmespath
diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 17848a7828b..d5d089cb1c1 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -73,8 +73,8 @@ We currently support Amazon DynamoDB and Redis as a storage layer. The following
 If you're not [changing the default configuration for the DynamoDB persistence layer](#dynamodbpersistencelayer), this is the expected default configuration:
 
 | Configuration      | Value        | Notes                                                                               |
-| ------------------ | ------------ | ----------------------------------------------------------------------------------- |
-| Partition key      | `id`         |
+| ------------------ | ------------ |-------------------------------------------------------------------------------------|
+| Partition key      | `id`         |                                                                                     |
 | TTL attribute name | `expiration` | This can only be configured after your table is created if you're using AWS Console |
 
 ???+ tip "Tip: You can share a single state table for all functions"
@@ -699,15 +699,16 @@ For advanced configurations, such as setting up SSL certificates or customizing
 
 Idempotent decorator can be further configured with **`IdempotencyConfig`** as seen in the previous example. These are the available options for further configuration
 
-| Parameter                       | Default | Description                                                                                                                                                               |
-| ------------------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"} |
-| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                              |
-| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                            |
-| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                  |
-| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                              |
-| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                               |
-| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library. |
+| Parameter                       | Default | Description                                                                                                                                                                                                                         |
+|---------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"}                                                           |
+| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                                                                                        |
+| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                                                                                      |
+| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                                                                            |
+| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                        |
+| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                         |
+| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                           |
+| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response.  This function hook is called when an already returned response is found.  See [Modifying The Idempotent Response](idempotency.md#modifying-the-idempotent-repsonse) |
 
 ### Handling concurrent executions with the same payload
 
@@ -909,6 +910,22 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
     For example, the `_put_record` method needs to raise an exception if a non-expired record already exists in the data store with a matching key.
 
+### Modifying the Idempotent Repsonse
+
+The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an idempotent response is loaded from the PersistenceStore.
+
+You can provide the response_hook using _**IdempotentConfig**_.
+
+=== "Using an Idempotent Response Hook"
+
+```python hl_lines="10-15 19"
+--8<-- "examples/idempotency/src/working_with_response_hook.py"
+```
+
+???+ info "Info: Using custom de-serialization?"
+
+    The response_hook is called after the custom de-serialization so the payload you process will be the de-serialized version.
+
 ## Compatibility with other utilities
 
 ### Batch
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
new file mode 100644
index 00000000000..61247959177
--- /dev/null
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -0,0 +1,27 @@
+from typing import Dict
+
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    IdempotencyConfig,
+    IdempotentHookData,
+    idempotent,
+)
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+
+def my_response_hook(response: Dict, idempotent_data: IdempotentHookData) -> Dict:
+    # How to add a field to the response
+    response["is_idempotent_response"] = True
+
+    # Must return the response here
+    return response
+
+
+persistence_layer = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
+
+config = IdempotencyConfig(event_key_jmespath="body", response_hook=my_response_hook)
+
+
+@idempotent(persistence_store=persistence_layer, config=config)
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return event
diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py
index abc45bb4104..4be0897eed1 100644
--- a/tests/functional/idempotency/test_idempotency.py
+++ b/tests/functional/idempotency/test_idempotency.py
@@ -2043,7 +2043,7 @@ def test_idempotent_lambda_already_completed_response_hook_is_called(
     Test idempotent decorator where event with matching event key has already been successfully processed
     """
 
-    def idempotent_response_hook(response: Any, idempotent_data: IdempotentHookData):
+    def idempotent_response_hook(response: Any, idempotent_data: IdempotentHookData) -> Any:
         """Modify the response provided by adding a new key"""
         response["idempotent_response"] = True
 

From 8ea0c4fc034f670603da66413aaf74e2480df448 Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Wed, 3 Apr 2024 22:14:04 +1100
Subject: [PATCH 04/10] feat(response_hook): review items

---
 .../utilities/idempotency/__init__.py         |  5 ++-
 .../utilities/idempotency/base.py             |  9 +++--
 .../utilities/idempotency/config.py           | 40 +------------------
 .../utilities/idempotency/hook.py             | 13 ++++++
 docs/utilities/idempotency.md                 | 28 ++++++-------
 .../src/working_with_response_hook.py         | 14 +++++--
 .../idempotency/test_idempotency.py           |  3 +-
 7 files changed, 48 insertions(+), 64 deletions(-)
 create mode 100644 aws_lambda_powertools/utilities/idempotency/hook.py

diff --git a/aws_lambda_powertools/utilities/idempotency/__init__.py b/aws_lambda_powertools/utilities/idempotency/__init__.py
index 77ab4862322..0c46553cc59 100644
--- a/aws_lambda_powertools/utilities/idempotency/__init__.py
+++ b/aws_lambda_powertools/utilities/idempotency/__init__.py
@@ -2,7 +2,9 @@
 Utility for adding idempotency to lambda functions
 """
 
-from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData, IdempotentHookFunction
+from aws_lambda_powertools.utilities.idempotency.hook import (
+    IdempotentHookFunction,
+)
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
@@ -18,6 +20,5 @@
     "idempotent",
     "idempotent_function",
     "IdempotencyConfig",
-    "IdempotentHookData",
     "IdempotentHookFunction",
 )
diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py
index 9ab5def6f11..4a249f20521 100644
--- a/aws_lambda_powertools/utilities/idempotency/base.py
+++ b/aws_lambda_powertools/utilities/idempotency/base.py
@@ -5,7 +5,6 @@
 
 from aws_lambda_powertools.utilities.idempotency.config import (
     IdempotencyConfig,
-    IdempotentHookData,
 )
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
@@ -230,12 +229,14 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
             )
         response_dict: Optional[dict] = data_record.response_json_as_dict()
         if response_dict is not None:
+            serialized_response = self.output_serializer.from_dict(response_dict)
             if self.config.response_hook is not None:
                 return self.config.response_hook(
-                    self.output_serializer.from_dict(response_dict),
-                    IdempotentHookData(data_record),
+                    serialized_response,
+                    data_record,
                 )
-            return self.output_serializer.from_dict(response_dict)
+            return serialized_response
+
         return None
 
     def _get_function_response(self):
diff --git a/aws_lambda_powertools/utilities/idempotency/config.py b/aws_lambda_powertools/utilities/idempotency/config.py
index aa11793bd8a..826dbbe4089 100644
--- a/aws_lambda_powertools/utilities/idempotency/config.py
+++ b/aws_lambda_powertools/utilities/idempotency/config.py
@@ -1,45 +1,9 @@
-from typing import Any, Dict, Optional
+from typing import Dict, Optional
 
-from aws_lambda_powertools.shared.types import Protocol
-from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
+from aws_lambda_powertools.utilities.idempotency import IdempotentHookFunction
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
-class IdempotentHookData:
-    """
-    Idempotent Hook Data
-
-    Contains data relevant to the current Idempotent record which matches the current request.
-    All IdempotentHook functions will be passed this data as well as the current Response.
-    """
-
-    def __init__(self, data_record: DataRecord) -> None:
-        self._idempotency_key = data_record.idempotency_key
-        self._status = data_record.status
-        self._expiry_timestamp = data_record.expiry_timestamp
-
-    @property
-    def idempotency_key(self) -> str:
-        return self._idempotency_key
-
-    @property
-    def status(self) -> str:
-        return self._status
-
-    @property
-    def expiry_timestamp(self) -> Optional[int]:
-        return self._expiry_timestamp
-
-
-class IdempotentHookFunction(Protocol):
-    """
-    The IdempotentHookFunction.
-    This class defines the calling signature for IdempotentHookFunction callbacks.
-    """
-
-    def __call__(self, response: Any, idempotent_data: IdempotentHookData) -> Any: ...
-
-
 class IdempotencyConfig:
     def __init__(
         self,
diff --git a/aws_lambda_powertools/utilities/idempotency/hook.py b/aws_lambda_powertools/utilities/idempotency/hook.py
new file mode 100644
index 00000000000..0027399b937
--- /dev/null
+++ b/aws_lambda_powertools/utilities/idempotency/hook.py
@@ -0,0 +1,13 @@
+from typing import Any
+
+from aws_lambda_powertools.shared.types import Protocol
+from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
+
+
+class IdempotentHookFunction(Protocol):
+    """
+    The IdempotentHookFunction.
+    This class defines the calling signature for IdempotentHookFunction callbacks.
+    """
+
+    def __call__(self, response: Any, idempotent_data: DataRecord) -> Any: ...
diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index d5d089cb1c1..54b9685fdce 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -699,16 +699,16 @@ For advanced configurations, such as setting up SSL certificates or customizing
 
 Idempotent decorator can be further configured with **`IdempotencyConfig`** as seen in the previous example. These are the available options for further configuration
 
-| Parameter                       | Default | Description                                                                                                                                                                                                                         |
-|---------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"}                                                           |
-| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                                                                                        |
-| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                                                                                      |
-| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                                                                            |
-| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                        |
-| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                         |
-| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                           |
-| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response.  This function hook is called when an already returned response is found.  See [Modifying The Idempotent Response](idempotency.md#modifying-the-idempotent-repsonse) |
+| Parameter                       | Default | Description                                                                                                                                                                                                                                  |
+|---------------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"}                                                                    |
+| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                                                                                                 |
+| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                                                                                               |
+| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                                                                                     |
+| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                                 |
+| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                                  |
+| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                                    |
+| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response.  This function hook is called when an existing idempotent response is found.  See [Manipulating The Idempotent Response](idempotency.md#manipulating-the-idempotent-response) |
 
 ### Handling concurrent executions with the same payload
 
@@ -910,15 +910,15 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
     For example, the `_put_record` method needs to raise an exception if a non-expired record already exists in the data store with a matching key.
 
-### Modifying the Idempotent Repsonse
+### Manipulating the Idempotent Response
 
-The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an idempotent response is loaded from the PersistenceStore.
+The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an already returned response is loaded from the PersistenceStore.  The Hook function will be called with the current de-serialized response object and the Idempotent DataRecord.
 
-You can provide the response_hook using _**IdempotentConfig**_.
+You can provide the response_hook using_**IdempotentConfig**_.
 
 === "Using an Idempotent Response Hook"
 
-```python hl_lines="10-15 19"
+```python hl_lines="15-23 28"
 --8<-- "examples/idempotency/src/working_with_response_hook.py"
 ```
 
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
index 61247959177..63419cbc4b9 100644
--- a/examples/idempotency/src/working_with_response_hook.py
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -1,17 +1,23 @@
+from datetime import datetime
 from typing import Dict
 
 from aws_lambda_powertools.utilities.idempotency import (
     DynamoDBPersistenceLayer,
     IdempotencyConfig,
-    IdempotentHookData,
     idempotent,
 )
+from aws_lambda_powertools.utilities.idempotency.persistence.base import (
+    DataRecord,
+)
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
-def my_response_hook(response: Dict, idempotent_data: IdempotentHookData) -> Dict:
-    # How to add a field to the response
-    response["is_idempotent_response"] = True
+def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
+    # Return inserted Header data into the Idempotent Response
+    expiry_time = datetime.fromtimestamp(idempotent_data.expiry_timestamp)
+
+    response["headers"]["x-idempotent-key"] = idempotent_data.idempotency_key
+    response["headers"]["x-idempotent-expiration"] = expiry_time.isoformat()
 
     # Must return the response here
     return response
diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py
index 4be0897eed1..d33469d680f 100644
--- a/tests/functional/idempotency/test_idempotency.py
+++ b/tests/functional/idempotency/test_idempotency.py
@@ -27,7 +27,6 @@
     IdempotencyHandler,
     _prepare_data,
 )
-from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
@@ -2043,7 +2042,7 @@ def test_idempotent_lambda_already_completed_response_hook_is_called(
     Test idempotent decorator where event with matching event key has already been successfully processed
     """
 
-    def idempotent_response_hook(response: Any, idempotent_data: IdempotentHookData) -> Any:
+    def idempotent_response_hook(response: Any, idempotent_data: DataRecord) -> Any:
         """Modify the response provided by adding a new key"""
         response["idempotent_response"] = True
 

From b83e3a67c0b9090f0a997fc2313f1cdd7e27590a Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Wed, 3 Apr 2024 23:34:16 +1100
Subject: [PATCH 05/10] chore(mypy): resolve type erro r in example code -
 expiry_timestamp can be None

---
 docs/utilities/idempotency.md                          | 4 ++--
 examples/idempotency/src/working_with_response_hook.py | 8 +++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 54b9685fdce..a0f6e3aacc4 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -914,11 +914,11 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
 The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an already returned response is loaded from the PersistenceStore.  The Hook function will be called with the current de-serialized response object and the Idempotent DataRecord.
 
-You can provide the response_hook using_**IdempotentConfig**_.
+You can provide the response_hook using _**IdempotentConfig**_.
 
 === "Using an Idempotent Response Hook"
 
-```python hl_lines="15-23 28"
+```python hl_lines="15-26 31"
 --8<-- "examples/idempotency/src/working_with_response_hook.py"
 ```
 
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
index 63419cbc4b9..5903935e61d 100644
--- a/examples/idempotency/src/working_with_response_hook.py
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -14,10 +14,12 @@
 
 def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
     # Return inserted Header data into the Idempotent Response
-    expiry_time = datetime.fromtimestamp(idempotent_data.expiry_timestamp)
-
     response["headers"]["x-idempotent-key"] = idempotent_data.idempotency_key
-    response["headers"]["x-idempotent-expiration"] = expiry_time.isoformat()
+
+    # expiry_timestamp could be None so include if set
+    if idempotent_data.expiry_timestamp:
+        expiry_time = datetime.fromtimestamp(idempotent_data.expiry_timestamp)
+        response["headers"]["x-idempotent-expiration"] = expiry_time.isoformat()
 
     # Must return the response here
     return response

From ec79ded3e94b93057662ef4d2772a287ee6d1f38 Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Thu, 4 Apr 2024 00:04:08 +1100
Subject: [PATCH 06/10] chore(docs): fix formatting error in markdown

---
 docs/utilities/idempotency.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index a0f6e3aacc4..cccb8b8f1ea 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -914,7 +914,7 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
 The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an already returned response is loaded from the PersistenceStore.  The Hook function will be called with the current de-serialized response object and the Idempotent DataRecord.
 
-You can provide the response_hook using _**IdempotentConfig**_.
+You can provide the _**response_hook**_ using _**IdempotentConfig**_.
 
 === "Using an Idempotent Response Hook"
 

From be0f3f72e28da58e3c19b58fb35d71090182ca41 Mon Sep 17 00:00:00 2001
From: walmsles <2704782+walmsles@users.noreply.github.com>
Date: Thu, 4 Apr 2024 00:06:11 +1100
Subject: [PATCH 07/10] chore(docs): fix highlighting of example code - lines
 moved

---
 docs/utilities/idempotency.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index cccb8b8f1ea..749c44fa15b 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -918,7 +918,7 @@ You can provide the _**response_hook**_ using _**IdempotentConfig**_.
 
 === "Using an Idempotent Response Hook"
 
-```python hl_lines="15-26 31"
+```python hl_lines="15-25 30"
 --8<-- "examples/idempotency/src/working_with_response_hook.py"
 ```
 

From f0596a4071d5227468451a6bd1c442a3f4e140f7 Mon Sep 17 00:00:00 2001
From: Leandro Damascena <lcdama@amazon.pt>
Date: Wed, 3 Apr 2024 16:10:28 +0100
Subject: [PATCH 08/10] Improving doc

---
 .../utilities/idempotency/base.py             |  1 +
 docs/utilities/idempotency.md                 | 52 ++++++++++++++++---
 .../src/working_with_response_hook.py         | 11 ++--
 3 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py
index 4a249f20521..f5ed9e2e476 100644
--- a/aws_lambda_powertools/utilities/idempotency/base.py
+++ b/aws_lambda_powertools/utilities/idempotency/base.py
@@ -231,6 +231,7 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
         if response_dict is not None:
             serialized_response = self.output_serializer.from_dict(response_dict)
             if self.config.response_hook is not None:
+                logger.debug("Response hook configured, invoking function")
                 return self.config.response_hook(
                     serialized_response,
                     data_record,
diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 749c44fa15b..03cd40eeb2d 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -454,6 +454,40 @@ sequenceDiagram
 <i>Idempotent successful request cached</i>
 </center>
 
+#### Successful request with response_hook configured
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Response hook
+    participant Persistence Layer
+    alt initial request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Note over Lambda,Persistence Layer: Set record status to INPROGRESS. <br> Prevents concurrent invocations <br> with the same payload
+        Lambda-->>Lambda: Call your function
+        Lambda->>Persistence Layer: Update record with result
+        deactivate Persistence Layer
+        Persistence Layer-->>Persistence Layer: Update record
+        Note over Lambda,Persistence Layer: Set record status to COMPLETE. <br> New invocations with the same payload <br> now return the same result
+        Lambda-->>Client: Response sent to client
+    else retried request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Persistence Layer-->>Response hook: Already exists in persistence layer.
+        deactivate Persistence Layer
+        Note over Response hook,Persistence Layer: Record status is COMPLETE and not expired
+        Response hook->>Lambda: Response hook invoked
+        Lambda-->>Client: Same response sent to client
+    end
+```
+<i>Idempotent successful request with response hook</i>
+</center>
+
 #### Expired idempotency records
 
 <center>
@@ -708,7 +742,7 @@ Idempotent decorator can be further configured with **`IdempotencyConfig`** as s
 | **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                                 |
 | **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                                  |
 | **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                                    |
-| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response.  This function hook is called when an existing idempotent response is found.  See [Manipulating The Idempotent Response](idempotency.md#manipulating-the-idempotent-response) |
+| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response. This function hook is called when an existing idempotent response is found. See [Manipulating The Idempotent Response](idempotency.md#manipulating-the-idempotent-response) |
 
 ### Handling concurrent executions with the same payload
 
@@ -912,15 +946,19 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
 ### Manipulating the Idempotent Response
 
-The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an already returned response is loaded from the PersistenceStore.  The Hook function will be called with the current de-serialized response object and the Idempotent DataRecord.
-
-You can provide the _**response_hook**_ using _**IdempotentConfig**_.
+You can set up a `response_hook` in the `IdempotentConfig` class to access the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record.
 
 === "Using an Idempotent Response Hook"
 
-```python hl_lines="15-25 30"
---8<-- "examples/idempotency/src/working_with_response_hook.py"
-```
+    ```python hl_lines="15 17 20 31"
+    --8<-- "examples/idempotency/src/working_with_response_hook.py"
+    ```
+
+=== "Sample event"
+
+    ```json
+    --8<-- "examples/idempotency/src/getting_started_with_idempotency_payload.json"
+    ```
 
 ???+ info "Info: Using custom de-serialization?"
 
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
index 5903935e61d..1bb203649a8 100644
--- a/examples/idempotency/src/working_with_response_hook.py
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -14,12 +14,13 @@
 
 def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
     # Return inserted Header data into the Idempotent Response
-    response["headers"]["x-idempotent-key"] = idempotent_data.idempotency_key
+    response["x-idempotent-key"] = idempotent_data.idempotency_key
 
     # expiry_timestamp could be None so include if set
-    if idempotent_data.expiry_timestamp:
-        expiry_time = datetime.fromtimestamp(idempotent_data.expiry_timestamp)
-        response["headers"]["x-idempotent-expiration"] = expiry_time.isoformat()
+    expiry_timestamp = idempotent_data.expiry_timestamp
+    if expiry_timestamp:
+        expiry_time = datetime.fromtimestamp(int(expiry_timestamp))
+        response["x-idempotent-expiration"] = expiry_time.isoformat()
 
     # Must return the response here
     return response
@@ -27,7 +28,7 @@ def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
 
 persistence_layer = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
 
-config = IdempotencyConfig(event_key_jmespath="body", response_hook=my_response_hook)
+config = IdempotencyConfig(response_hook=my_response_hook)
 
 
 @idempotent(persistence_store=persistence_layer, config=config)

From 5f50be5a63e14a6162d79f838c54973f9e68b5a7 Mon Sep 17 00:00:00 2001
From: Leandro Damascena <lcdama@amazon.pt>
Date: Wed, 3 Apr 2024 22:15:27 +0100
Subject: [PATCH 09/10] Improving doc

---
 docs/utilities/idempotency.md                 | 16 +++++++--
 .../src/working_with_response_hook.py         | 36 ++++++++++++++-----
 .../working_with_response_hook_payload.json   |  8 +++++
 3 files changed, 50 insertions(+), 10 deletions(-)
 create mode 100644 examples/idempotency/src/working_with_response_hook_payload.json

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 03cd40eeb2d..288dd738695 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -950,20 +950,32 @@ You can set up a `response_hook` in the `IdempotentConfig` class to access the r
 
 === "Using an Idempotent Response Hook"
 
-    ```python hl_lines="15 17 20 31"
+    ```python hl_lines="18 20 23 32"
     --8<-- "examples/idempotency/src/working_with_response_hook.py"
     ```
 
 === "Sample event"
 
     ```json
-    --8<-- "examples/idempotency/src/getting_started_with_idempotency_payload.json"
+    --8<-- "examples/idempotency/src/working_with_response_hook_payload.json"
     ```
 
 ???+ info "Info: Using custom de-serialization?"
 
     The response_hook is called after the custom de-serialization so the payload you process will be the de-serialized version.
 
+#### Being a good citizen
+
+Using Response hooks can add subtle improvements to manipulating returned data from idempotent operations, but also add significant complexity if you're not careful.
+
+Keep the following in mind when authoring hooks for Idempotency utility:
+
+1. **Response hook works exclusively when operations are idempotent.** Carefully consider the logic within the `Response hook` and prevent any attempt to access the key from relying exclusively on idempotent operations.
+
+2. **Catch your own exceptions.** Catch and handle known exceptions to your logic.
+
+3. **Watch out when you are decorating the Lambda Handler and using the Response hook.** If you don't catch and handle exceptions in your `Response hook`, your function might not run properly.
+
 ## Compatibility with other utilities
 
 ### Batch
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
index 1bb203649a8..a5509eecd81 100644
--- a/examples/idempotency/src/working_with_response_hook.py
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -1,16 +1,20 @@
-from datetime import datetime
+import datetime
+import uuid
 from typing import Dict
 
+from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities.idempotency import (
     DynamoDBPersistenceLayer,
     IdempotencyConfig,
-    idempotent,
+    idempotent_function,
 )
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     DataRecord,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
+logger = Logger()
+
 
 def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
     # Return inserted Header data into the Idempotent Response
@@ -19,18 +23,34 @@ def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
     # expiry_timestamp could be None so include if set
     expiry_timestamp = idempotent_data.expiry_timestamp
     if expiry_timestamp:
-        expiry_time = datetime.fromtimestamp(int(expiry_timestamp))
+        expiry_time = datetime.datetime.fromtimestamp(int(expiry_timestamp))
         response["x-idempotent-expiration"] = expiry_time.isoformat()
 
     # Must return the response here
     return response
 
 
-persistence_layer = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
-
+dynamodb = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
 config = IdempotencyConfig(response_hook=my_response_hook)
 
 
-@idempotent(persistence_store=persistence_layer, config=config)
-def lambda_handler(event: dict, context: LambdaContext) -> dict:
-    return event
+@idempotent_function(data_keyword_argument="order", config=config, persistence_store=dynamodb)
+def process_order(order: dict) -> dict:
+    # create the order_id
+    order_id = str(uuid.uuid4())
+
+    # create your logic to save the order
+    # append the order_id created
+    order["order_id"] = order_id
+
+    # return the order
+    return {"order": order}
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    config.register_lambda_context(context)  # see Lambda timeouts section
+    try:
+        logger.info(f"Processing order id {event.get('order_id')}")
+        return process_order(order=event.get("order"))
+    except Exception as err:
+        return {"status_code": 400, "error": f"Erro processing {str(err)}"}
diff --git a/examples/idempotency/src/working_with_response_hook_payload.json b/examples/idempotency/src/working_with_response_hook_payload.json
new file mode 100644
index 00000000000..85fdd958d59
--- /dev/null
+++ b/examples/idempotency/src/working_with_response_hook_payload.json
@@ -0,0 +1,8 @@
+{
+  "order" : {
+    "user_id": "xyz",
+    "product_id": "123456789",
+    "quantity": 2,
+    "value": 30
+  }
+}

From a04d0476e517831045d43381a644faf2dfabacaf Mon Sep 17 00:00:00 2001
From: Leandro Damascena <lcdama@amazon.pt>
Date: Thu, 4 Apr 2024 10:54:32 +0100
Subject: [PATCH 10/10] Addressing Ruben's feedback

---
 docs/utilities/idempotency.md                    | 16 +++++++---------
 .../src/working_with_response_hook.py            |  2 +-
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 288dd738695..e448d82e28e 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -482,10 +482,10 @@ sequenceDiagram
         deactivate Persistence Layer
         Note over Response hook,Persistence Layer: Record status is COMPLETE and not expired
         Response hook->>Lambda: Response hook invoked
-        Lambda-->>Client: Same response sent to client
+        Lambda-->>Client: Manipulated idempotent response sent to client
     end
 ```
-<i>Idempotent successful request with response hook</i>
+<i>Successful idempotent request with a response hook</i>
 </center>
 
 #### Expired idempotency records
@@ -946,7 +946,7 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
 ### Manipulating the Idempotent Response
 
-You can set up a `response_hook` in the `IdempotentConfig` class to access the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record.
+You can set up a `response_hook` in the `IdempotentConfig` class to manipulate the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record.
 
 === "Using an Idempotent Response Hook"
 
@@ -966,15 +966,13 @@ You can set up a `response_hook` in the `IdempotentConfig` class to access the r
 
 #### Being a good citizen
 
-Using Response hooks can add subtle improvements to manipulating returned data from idempotent operations, but also add significant complexity if you're not careful.
-
-Keep the following in mind when authoring hooks for Idempotency utility:
+When using response hooks to manipulate returned data from idempotent operations, it's important to follow best practices to avoid introducing complexity or issues. Keep these guidelines in mind:
 
-1. **Response hook works exclusively when operations are idempotent.** Carefully consider the logic within the `Response hook` and prevent any attempt to access the key from relying exclusively on idempotent operations.
+1. **Response hook works exclusively when operations are idempotent.** The hook will not be called when an operation is not idempotent, or when the idempotent logic fails.
 
-2. **Catch your own exceptions.** Catch and handle known exceptions to your logic.
+2. **Catch and Handle Exceptions.** Your response hook code should catch and handle any exceptions that may arise from your logic. Unhandled exceptions will cause the Lambda function to fail unexpectedly.
 
-3. **Watch out when you are decorating the Lambda Handler and using the Response hook.** If you don't catch and handle exceptions in your `Response hook`, your function might not run properly.
+3. **Keep Hook Logic Simple** Response hooks should consist of minimal and straightforward logic for manipulating response data. Avoid complex conditional branching and aim for hooks that are easy to reason about.
 
 ## Compatibility with other utilities
 
diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py
index a5509eecd81..725a56f32ba 100644
--- a/examples/idempotency/src/working_with_response_hook.py
+++ b/examples/idempotency/src/working_with_response_hook.py
@@ -53,4 +53,4 @@ def lambda_handler(event: dict, context: LambdaContext):
         logger.info(f"Processing order id {event.get('order_id')}")
         return process_order(order=event.get("order"))
     except Exception as err:
-        return {"status_code": 400, "error": f"Erro processing {str(err)}"}
+        return {"status_code": 400, "error": f"Error processing {str(err)}"}