fix(idempotency): complete and delete the record acquired at operation start

docs/features/idempotency.md

@@ -913,6 +913,15 @@ Below an example implementation of a custom persistence layer backed by a generi
 
     For example, the `_putRecord()` method needs to throw an error if a non-expired record already exists in the data store with a matching key.
 
+???+ tip "Overriding the public methods"
+    The protected methods above are the recommended extension points for a custom storage implementation. Subclasses that implement only these hooks need no further changes to keep the record identity guarantees of the base class.
+
+    If you override `saveInProgress()`, `saveSuccess()`, or `deleteRecord()` instead:
+
+    * Accept the optional trailing `options` argument (typed as `PersistenceOperationOptions`) and forward it unchanged to `super`, even if your application does not mutate its inputs today. This also preserves any option added in the future.
+    * The utility resolves the idempotency key and payload hash before your function runs and passes them as `options.identity` to these methods, so completion and cleanup target the record that was created even if the function mutated its input.
+    * An override that forwards only the original arguments discards the identity and falls back to hashing the payload as it is at that point.
+
 ### Manipulating the Idempotent Response
 
 You can set up a `responseHook` 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.

packages/idempotency/src/IdempotencyHandler.ts

@@ -15,6 +15,7 @@ import {
 import type { IdempotencyConfig } from './IdempotencyConfig.js';
 import type { BasePersistenceLayer } from './persistence/BasePersistenceLayer.js';
 import type { IdempotencyRecord } from './persistence/IdempotencyRecord.js';
+import type { IdempotencyRecordIdentity } from './types/BasePersistenceLayer.js';
 import type {
   AnyFunction,
   IdempotencyHandlerOptions,
@@ -43,6 +44,14 @@ export class IdempotencyHandler<Func extends AnyFunction> {
    * This is the argument that is used for the idempotency.
    */
   #functionPayloadToBeHashed: JSONValue;
+  /**
+   * The identity of the idempotency record, resolved from the payload before
+   * the operation starts.
+   *
+   * It is passed to the persistence calls that follow, so they all target the
+   * same record even if the wrapped function mutated the payload in the meantime.
+   */
+  #recordIdentity?: IdempotencyRecordIdentity;
   /**
    * Reference to the function to be made idempotent.
    */
@@ -344,7 +353,8 @@ export class IdempotencyHandler<Func extends AnyFunction> {
   readonly #deleteInProgressRecord = async (): Promise<void> => {
     try {
       await this.#persistenceStore.deleteRecord(
-        this.#functionPayloadToBeHashed
+        this.#functionPayloadToBeHashed,
+        { identity: this.#recordIdentity }
       );
     } catch (error) {
       throw new IdempotencyPersistenceLayerError(
@@ -376,9 +386,15 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       result: undefined,
     };
     try {
+      // Resolve the identity before attempting acquisition so that it is also
+      // available when the function runs without acquiring a record (replay).
+      this.#recordIdentity = this.#persistenceStore.getRecordIdentity(
+        this.#functionPayloadToBeHashed
+      );
       await this.#persistenceStore.saveInProgress(
         this.#functionPayloadToBeHashed,
-        this.#idempotencyConfig.lambdaContext?.getRemainingTimeInMillis()
+        this.#idempotencyConfig.lambdaContext?.getRemainingTimeInMillis(),
+        { identity: this.#recordIdentity }
       );
 
       return returnValue;
@@ -437,7 +453,8 @@ export class IdempotencyHandler<Func extends AnyFunction> {
     try {
       await this.#persistenceStore.saveSuccess(
         this.#functionPayloadToBeHashed,
-        result
+        result,
+        { identity: this.#recordIdentity }
       );
     } catch (error) {
       throw new IdempotencyPersistenceLayerError(

packages/idempotency/src/persistence/BasePersistenceLayer.ts

@@ -15,6 +15,8 @@ import type { IdempotencyConfig } from '../IdempotencyConfig.js';
 import type {
   BasePersistenceLayerInterface,
   BasePersistenceLayerOptions,
+  IdempotencyRecordIdentity,
+  PersistenceOperationOptions,
 } from '../types/BasePersistenceLayer.js';
 import { IdempotencyRecord } from './IdempotencyRecord.js';
 
@@ -158,10 +160,15 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
    * Deletes a record from the persistence store for the persistence key generated from the data passed in.
    *
    * @param data - the data payload that will be hashed to create the hash portion of the idempotency key
+   * @param options - operation options; when `options.identity` is provided its key is used instead of hashing `data` again, so the record is deleted even if `data` was mutated after the identity was resolved by {@link BasePersistenceLayer.getRecordIdentity | `getRecordIdentity()`}
    */
-  public async deleteRecord(data: JSONValue): Promise<void> {
+  public async deleteRecord(
+    data: JSONValue,
+    options?: PersistenceOperationOptions
+  ): Promise<void> {
     const idempotencyRecord = new IdempotencyRecord({
-      idempotencyKey: this.getHashedIdempotencyKey(data),
+      idempotencyKey:
+        options?.identity?.idempotencyKey ?? this.getHashedIdempotencyKey(data),
       status: IdempotencyRecordStatus.EXPIRED,
     });
 
@@ -177,6 +184,22 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
     return this.expiresAfterSeconds;
   }
 
+  /**
+   * Resolves the idempotency key and payload hash that identify the record for the provided data.
+   *
+   * Call it before the operation starts and pass the result to {@link BasePersistenceLayer.saveInProgress | `saveInProgress()`},
+   * {@link BasePersistenceLayer.saveSuccess | `saveSuccess()`}, and {@link BasePersistenceLayer.deleteRecord | `deleteRecord()`},
+   * so they all target the same record even if `data` is mutated while the operation runs.
+   *
+   * @param data - the data payload that will be hashed to create the hash portion of the idempotency key
+   */
+  public getRecordIdentity(data: JSONValue): IdempotencyRecordIdentity {
+    return {
+      idempotencyKey: this.getHashedIdempotencyKey(data),
+      payloadHash: this.getHashedPayload(data),
+    };
+  }
+
   /**
    * Retrieves idempotency key for the provided data and fetches data for that key from the persistence store
    *
@@ -233,16 +256,20 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
    *
    * @param data - the data payload that will be hashed to create the hash portion of the idempotency key
    * @param remainingTimeInMillis - the remaining time left in the lambda execution context
+   * @param options - operation options; when `options.identity` is provided it is used instead of hashing `data` again, see {@link BasePersistenceLayer.getRecordIdentity | `getRecordIdentity()`}
    */
   public async saveInProgress(
     data: JSONValue,
-    remainingTimeInMillis?: number
+    remainingTimeInMillis?: number,
+    options?: PersistenceOperationOptions
   ): Promise<void> {
+    const { idempotencyKey, payloadHash } =
+      options?.identity ?? this.getRecordIdentity(data);
     const idempotencyRecord = new IdempotencyRecord({
-      idempotencyKey: this.getHashedIdempotencyKey(data),
+      idempotencyKey,
       status: IdempotencyRecordStatus.INPROGRESS,
       expiryTimestamp: this.getExpiryTimestamp(),
-      payloadHash: this.getHashedPayload(data),
+      payloadHash,
     });
 
     if (remainingTimeInMillis) {
@@ -271,14 +298,21 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
    *
    * @param data - the data payload that will be hashed to create the hash portion of the idempotency key
    * @param result - the result of the successfully completed function
+   * @param options - operation options; when `options.identity` is provided its key and payload hash are used instead of hashing `data` again, so the in-progress record is completed even if `data` was mutated after the identity was resolved by {@link BasePersistenceLayer.getRecordIdentity | `getRecordIdentity()`}
    */
-  public async saveSuccess(data: JSONValue, result: JSONValue): Promise<void> {
+  public async saveSuccess(
+    data: JSONValue,
+    result: JSONValue,
+    options?: PersistenceOperationOptions
+  ): Promise<void> {
+    const { idempotencyKey, payloadHash } =
+      options?.identity ?? this.getRecordIdentity(data);
     const idempotencyRecord = new IdempotencyRecord({
-      idempotencyKey: this.getHashedIdempotencyKey(data),
+      idempotencyKey,
       status: IdempotencyRecordStatus.COMPLETED,
       expiryTimestamp: this.getExpiryTimestamp(),
       responseData: result,
-      payloadHash: this.getHashedPayload(data),
+      payloadHash,
     });
 
     await this._updateRecord(idempotencyRecord);

packages/idempotency/src/types/BasePersistenceLayer.ts

@@ -7,12 +7,43 @@ type BasePersistenceLayerOptions = {
   keyPrefix?: string;
 };
 
+/**
+ * The idempotency key and payload hash that identify the record of an operation.
+ *
+ * It is resolved from the payload before the operation starts, so completion and cleanup
+ * target the same record even if the payload is mutated while the operation runs.
+ */
+type IdempotencyRecordIdentity = Pick<
+  IdempotencyRecord,
+  'idempotencyKey' | 'payloadHash'
+>;
+
+/**
+ * Options accepted by the persistence operations that write or delete a record.
+ *
+ * @property identity - the record identity resolved before the operation started; when provided it is used instead of hashing the payload again
+ */
+type PersistenceOperationOptions = {
+  identity?: IdempotencyRecordIdentity;
+};
+
 interface BasePersistenceLayerInterface {
   configure(options?: BasePersistenceLayerOptions): void;
   isPayloadValidationEnabled(): boolean;
-  saveInProgress(data: unknown, remainingTimeInMillis?: number): Promise<void>;
-  saveSuccess(data: unknown, result: unknown): Promise<void>;
-  deleteRecord(data: unknown): Promise<void>;
+  saveInProgress(
+    data: unknown,
+    remainingTimeInMillis?: number,
+    options?: PersistenceOperationOptions
+  ): Promise<void>;
+  saveSuccess(
+    data: unknown,
+    result: unknown,
+    options?: PersistenceOperationOptions
+  ): Promise<void>;
+  deleteRecord(
+    data: unknown,
+    options?: PersistenceOperationOptions
+  ): Promise<void>;
   getRecord(data: unknown): Promise<IdempotencyRecord>;
 }
 
@@ -38,4 +69,6 @@ export type {
   BasePersistenceAttributes,
   BasePersistenceLayerInterface,
   BasePersistenceLayerOptions,
+  IdempotencyRecordIdentity,
+  PersistenceOperationOptions,
 };

packages/idempotency/src/types/index.ts

@@ -2,6 +2,8 @@ export type {
   BasePersistenceAttributes,
   BasePersistenceLayerInterface,
   BasePersistenceLayerOptions,
+  IdempotencyRecordIdentity,
+  PersistenceOperationOptions,
 } from './BasePersistenceLayer.js';
 export type {
   CacheClient,