Add support for response hooks in idempotency configuration

- Introduce ResponseHook delegate to IdempotencyOptions and builder
- Invoke response hook on idempotent responses in IdempotencyAspectHandler
- Update documentation with usage examples and best practices for response hooks
- Add integration tests for response hook execution and exception handling
- Improve XML doc comments for builder and persistence store

Author: hjgraca

docs/utilities/idempotency.md

@@ -846,6 +846,50 @@ Data would then be stored in DynamoDB like this:
 | idempotency#MyLambdaFunction | f091d2527ad1c78f05d54cc3f363be80 | 1636549585 | IN_PROGRESS |                                      |
 
 
+### Manipulating the Idempotent Response
+
+You can set up a response hook in the Idempotency configuration 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 `DataRecord`.
+
+#### Using Response Hooks
+
+The example below shows how to append HTTP headers to an `APIGatewayProxyResponse`:
+
+```csharp
+Idempotency.Config()
+    .WithConfig(IdempotencyOptions.Builder()
+        .WithEventKeyJmesPath("powertools_json(body).address")
+        .WithResponseHook((responseData, dataRecord) => {
+            if (responseData is APIGatewayProxyResponse proxyResponse)
+            {
+                proxyResponse.Headers ??= new Dictionary<string, string>();
+                proxyResponse.Headers["x-idempotency-response"] = "true";
+                proxyResponse.Headers["x-idempotency-expiration"] = dataRecord.ExpiryTimestamp.ToString();
+                return proxyResponse;
+            }
+            return responseData;
+        })
+        .Build())
+    .WithPersistenceStore(DynamoDBPersistenceStore.Builder()
+        .WithTableName(Environment.GetEnvironmentVariable("IDEMPOTENCY_TABLE"))
+        .Build())
+    .Configure();
+```
+
+???+ info "Info: Using custom de-serialization?"
+
+    The response hook is called after de-serialization so the payload you process will be the de-serialized Java object.
+
+#### Being a good citizen
+
+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.** The hook will not be called when an operation is not idempotent, or when the idempotent logic fails.
+
+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. **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.
+
+
 ## AOT Support
 
 Native AOT trims your application code as part of the compilation to ensure that the binary is as small as possible. .NET 8 for Lambda provides improved trimming support compared to previous versions of .NET.
@@ -921,4 +965,4 @@ When testing your code, you may wish to disable the idempotency logic altogether
 ## Extra resources
 
 If you're interested in a deep dive on how Amazon uses idempotency when building our APIs, check out
-[this article](https://aws.amazon.com/builders-library/making-retries-safe-with-idempotent-APIs/).
+[this article](https://aws.amazon.com/builders-library/making-retries-safe-with-idempotent-APIs/).

libraries/src/AWS.Lambda.Powertools.Idempotency/IdempotencyOptions.cs

@@ -13,6 +13,8 @@
  * permissions and limitations under the License.
  */
 
+using System;
+
 namespace AWS.Lambda.Powertools.Idempotency;
 
 /// <summary>
@@ -57,6 +59,10 @@ public class IdempotencyOptions
     /// as supported by <see cref="System.Security.Cryptography.HashAlgorithm"/> (eg. SHA1, SHA-256, ...)
     /// </summary>
     public string HashFunction { get; }
+    /// <summary>
+    /// Delegate for manipulating idempotent responses.
+    /// </summary>
+    public Func<object, Persistence.DataRecord, object> ResponseHook { get; }
 
     /// <summary>
     /// Constructor of <see cref="IdempotencyOptions"/>.
@@ -68,14 +74,16 @@ public class IdempotencyOptions
     /// <param name="localCacheMaxItems"></param>
     /// <param name="expirationInSeconds"></param>
     /// <param name="hashFunction"></param>
+    /// <param name="responseHook"></param>
     internal IdempotencyOptions(
         string eventKeyJmesPath, 
         string payloadValidationJmesPath, 
         bool throwOnNoIdempotencyKey, 
         bool useLocalCache, 
         int localCacheMaxItems, 
         long expirationInSeconds, 
-        string hashFunction)
+        string hashFunction,
+        Func<object, Persistence.DataRecord, object> responseHook = null)
     {
         EventKeyJmesPath = eventKeyJmesPath;
         PayloadValidationJmesPath = payloadValidationJmesPath;
@@ -84,5 +92,6 @@ internal IdempotencyOptions(
         LocalCacheMaxItems = localCacheMaxItems;
         ExpirationInSeconds = expirationInSeconds;
         HashFunction = hashFunction;
+        ResponseHook = responseHook;
     }
 }

libraries/src/AWS.Lambda.Powertools.Idempotency/IdempotencyOptionsBuilder.cs

@@ -43,21 +43,27 @@ public class IdempotencyOptionsBuilder
     private string _hashFunction = "MD5";
 
     /// <summary>
-    /// Initialize and return an instance of IdempotencyConfig.
+    /// Response hook function
+    /// </summary>
+    private Func<object, AWS.Lambda.Powertools.Idempotency.Persistence.DataRecord, object> _responseHook;
+
+    /// <summary>
+    /// Initialize and return an instance of IdempotencyOptions.
     /// Example:
-    /// IdempotencyConfig.Builder().WithUseLocalCache().Build();
-    /// This instance must then be passed to the Idempotency.Config:
-    /// Idempotency.Config().WithConfig(config).Configure();
+    /// new IdempotencyOptionsBuilder().WithUseLocalCache().Build();
+    /// This instance can then be passed to Idempotency.Configure:
+    /// Idempotency.Configure(builder => builder.WithOptions(options));
     /// </summary>
-    /// <returns>an instance of IdempotencyConfig</returns>
+    /// <returns>an instance of IdempotencyOptions</returns>
     public IdempotencyOptions Build() =>
         new(_eventKeyJmesPath,
             _payloadValidationJmesPath,
             _throwOnNoIdempotencyKey,
             _useLocalCache,
             _localCacheMaxItems,
             _expirationInSeconds,
-            _hashFunction);
+            _hashFunction,
+            _responseHook);
 
     /// <summary>
     /// A JMESPath expression to extract the idempotency key from the event record.
@@ -133,4 +139,15 @@ public IdempotencyOptionsBuilder WithHashFunction(string hashFunction)
 #endif
         return this;
     }
+
+    /// <summary>
+    /// Set a response hook function, to be called with the response and the data record.
+    /// </summary>
+    /// <param name="hook">The response hook function</param>
+    /// <returns>the instance of the builder (to chain operations)</returns>
+    public IdempotencyOptionsBuilder WithResponseHook(Func<object, AWS.Lambda.Powertools.Idempotency.Persistence.DataRecord, object> hook)
+    {
+        _responseHook = hook;
+        return this;
+    }
 }

libraries/src/AWS.Lambda.Powertools.Idempotency/Internal/IdempotencyAspectHandler.cs

@@ -208,6 +208,24 @@ private Task<T> HandleForStatus(DataRecord record)
                     {
                         throw new IdempotencyPersistenceLayerException("Unable to cast function response as " + typeof(T).Name);
                     }
+                    // Response hook logic
+                    var responseHook = Idempotency.Instance.IdempotencyOptions?.ResponseHook;
+                    if (responseHook != null)
+                    {
+                        try
+                        {
+                            var hooked = responseHook(result, record);
+                            if (hooked is T typedHooked)
+                            {
+                                return Task.FromResult(typedHooked);
+                            }
+                            // If hook returns wrong type, fallback to original result
+                        }
+                        catch (Exception)
+                        {
+                            // If hook throws, fallback to original result
+                        }
+                    }
                     return Task.FromResult(result);
                 }
                 catch (Exception e)

libraries/src/AWS.Lambda.Powertools.Idempotency/Persistence/DynamoDBPersistenceStore.cs

@@ -370,9 +370,9 @@ public class DynamoDBPersistenceStoreBuilder
     private AmazonDynamoDBClient _dynamoDbClient;
 
     /// <summary>
-    /// Initialize and return a new instance of {@link DynamoDBPersistenceStore}.
+    /// Initialize and return a new instance of <see cref="DynamoDBPersistenceStore"/>.
     /// Example:
-    ///    DynamoDBPersistenceStore.builder().withTableName("idempotency_store").build();
+    ///    new DynamoDBPersistenceStoreBuilder().WithTableName("idempotency_store").Build();
     /// </summary>
     /// <returns></returns>
     /// <exception cref="ArgumentNullException"></exception>

libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/ResponseHookTest.cs

@@ -0,0 +1,168 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Threading.Tasks;
+using Amazon.DynamoDBv2;
+using Amazon.Lambda.APIGatewayEvents;
+using AWS.Lambda.Powertools.Idempotency.Internal.Serializers;
+using AWS.Lambda.Powertools.Idempotency.Persistence;
+using AWS.Lambda.Powertools.Idempotency.Tests.Persistence;
+using FluentAssertions;
+using Xunit;
+
+namespace AWS.Lambda.Powertools.Idempotency.Tests;
+
+public class ResponseHookTest : IClassFixture<DynamoDbFixture>
+{
+    private readonly AmazonDynamoDBClient _client;
+    private readonly string _tableName;
+
+    public ResponseHookTest(DynamoDbFixture fixture)
+    {
+        _client = fixture.Client;
+        _tableName = fixture.TableName;
+    }
+
+    [Fact]
+    [Trait("Category", "Integration")]
+    public async Task ResponseHook_ShouldNotExecuteOnFirstCall()
+    {
+        // Arrange
+        var hookExecuted = false;
+        
+        Idempotency.Configure(builder => builder
+            .WithOptions(options => options
+                .WithEventKeyJmesPath("powertools_json(body).address")
+                .WithResponseHook((responseData, dataRecord) => {
+                    hookExecuted = true;
+                    if (responseData is APIGatewayProxyResponse proxyResponse)
+                    {
+                        var headers = new Dictionary<string, string>(proxyResponse.Headers ?? new Dictionary<string, string>());
+                        headers["x-idempotency-response"] = "true";
+                        headers["x-idempotency-expiration"] = dataRecord.ExpiryTimestamp.ToString();
+                        proxyResponse.Headers = headers;
+                        return proxyResponse;
+                    }
+                    return responseData;
+                }))
+            .WithPersistenceStore(new DynamoDBPersistenceStoreBuilder()
+                .WithTableName(_tableName)
+                .WithDynamoDBClient(_client)
+                .Build()));
+
+        var function = new ResponseHookTestFunction();
+        var request = IdempotencySerializer.Deserialize<APIGatewayProxyRequest>(
+            await File.ReadAllTextAsync("./resources/apigw_event2.json"));
+
+        // Act - First call
+        var response = await function.Handle(request);
+
+        // Assert - Hook should not execute on first call
+        hookExecuted.Should().BeFalse();
+        response.Headers.Should().NotContainKey("x-idempotency-response");
+        function.HandlerExecuted.Should().BeTrue();
+    }
+
+    [Fact]
+    [Trait("Category", "Integration")]
+    public async Task ResponseHook_ShouldExecuteOnIdempotentCall()
+    {
+        // Arrange
+        var hookExecuted = false;
+        
+        Idempotency.Configure(builder => builder
+            .WithOptions(options => options
+                .WithEventKeyJmesPath("powertools_json(body).address")
+                .WithResponseHook((responseData, dataRecord) => {
+                    hookExecuted = true;
+                    if (responseData is APIGatewayProxyResponse proxyResponse)
+                    {
+                        var headers = new Dictionary<string, string>(proxyResponse.Headers ?? new Dictionary<string, string>());
+                        headers["x-idempotency-response"] = "true";
+                        headers["x-idempotency-expiration"] = dataRecord.ExpiryTimestamp.ToString();
+                        proxyResponse.Headers = headers;
+                        return proxyResponse;
+                    }
+                    return responseData;
+                }))
+            .WithPersistenceStore(new DynamoDBPersistenceStoreBuilder()
+                .WithTableName(_tableName)
+                .WithDynamoDBClient(_client)
+                .Build()));
+
+        var function = new ResponseHookTestFunction();
+        var request = IdempotencySerializer.Deserialize<APIGatewayProxyRequest>(
+            await File.ReadAllTextAsync("./resources/apigw_event2.json"));
+
+        // Act - First call to populate cache
+        await function.Handle(request);
+        function.HandlerExecuted = false;
+        hookExecuted = false;
+
+        // Act - Second call (idempotent)
+        var response = await function.Handle(request);
+
+        // Assert - Hook should execute on idempotent call
+        hookExecuted.Should().BeTrue();
+        response.Headers.Should().ContainKey("x-idempotency-response");
+        response.Headers["x-idempotency-response"].Should().Be("true");
+        response.Headers.Should().ContainKey("x-idempotency-expiration");
+        function.HandlerExecuted.Should().BeFalse();
+    }
+
+    [Fact]
+    [Trait("Category", "Integration")]
+    public async Task ResponseHook_ShouldHandleExceptionsGracefully()
+    {
+        // Arrange
+        Idempotency.Configure(builder => builder
+            .WithOptions(options => options
+                .WithEventKeyJmesPath("powertools_json(body).address")
+                .WithResponseHook((responseData, dataRecord) => {
+                    throw new InvalidOperationException("Hook failed");
+                }))
+            .WithPersistenceStore(new DynamoDBPersistenceStoreBuilder()
+                .WithTableName(_tableName)
+                .WithDynamoDBClient(_client)
+                .Build()));
+
+        var function = new ResponseHookTestFunction();
+        var request = IdempotencySerializer.Deserialize<APIGatewayProxyRequest>(
+            await File.ReadAllTextAsync("./resources/apigw_event2.json"));
+
+        // Act - First call to populate cache
+        var firstResponse = await function.Handle(request);
+        function.HandlerExecuted = false;
+
+        // Act - Second call (idempotent) - should not throw despite hook exception
+        var response = await function.Handle(request);
+
+        // Assert - Should return original response despite hook exception
+        response.Should().NotBeNull();
+        response.Body.Should().Be(firstResponse.Body);
+        function.HandlerExecuted.Should().BeFalse();
+    }
+}
+
+public class ResponseHookTestFunction
+{
+    public bool HandlerExecuted { get; set; }
+
+    [Idempotent]
+    public async Task<APIGatewayProxyResponse> Handle(APIGatewayProxyRequest request)
+    {
+        HandlerExecuted = true;
+        
+        await Task.Delay(100); // Simulate some work
+        
+        return new APIGatewayProxyResponse
+        {
+            StatusCode = 200,
+            Body = "Hello World",
+            Headers = new Dictionary<string, string>
+            {
+                ["Content-Type"] = "application/json"
+            }
+        };
+    }
+}