fix(stepfunctions): retry block in CustomState is always empty

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json

@@ -20,13 +20,13 @@
   "StateMachine2E01A3A5": {
    "Type": "AWS::StepFunctions::StateMachine",
    "Properties": {
+    "DefinitionString": "{\"StartAt\":\"my custom task\",\"States\":{\"my custom task\":{\"Next\":\"final step\",\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::dynamodb:putItem\",\"Parameters\":{\"TableName\":\"my-cool-table\",\"Item\":{\"id\":{\"S\":\"my-entry\"}}},\"ResultPath\":null,\"Retry\":[{\"ErrorEquals\":[\"States.ALL\"],\"IntervalSeconds\":10,\"MaxAttempts\":5}],\"Catch\":[{\"ErrorEquals\":[\"States.ALL\"],\"Next\":\"failed\"}]},\"final step\":{\"Type\":\"Pass\",\"End\":true},\"failed\":{\"Type\":\"Fail\",\"Error\":\"DidNotWork\",\"Cause\":\"We got stuck\"}},\"TimeoutSeconds\":30}",
     "RoleArn": {
      "Fn::GetAtt": [
       "StateMachineRoleB840431D",
       "Arn"
      ]
-    },
-    "DefinitionString": "{\"StartAt\":\"my custom task\",\"States\":{\"my custom task\":{\"Next\":\"final step\",\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::dynamodb:putItem\",\"Parameters\":{\"TableName\":\"my-cool-table\",\"Item\":{\"id\":{\"S\":\"my-entry\"}}},\"ResultPath\":null,\"Catch\":[{\"ErrorEquals\":[\"States.ALL\"],\"Next\":\"failed\"}]},\"final step\":{\"Type\":\"Pass\",\"End\":true},\"failed\":{\"Type\":\"Fail\",\"Error\":\"DidNotWork\",\"Cause\":\"We got stuck\"}},\"TimeoutSeconds\":30}"
+    }
    },
    "DependsOn": [
     "StateMachineRoleB840431D"

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.ts

@@ -35,6 +35,11 @@ const custom = new sfn.CustomState(stack, 'my custom task', {
 });
 
 custom.addCatch(failure);
+custom.addRetry({
+  errors: [sfn.Errors.ALL],
+  interval: cdk.Duration.seconds(10),
+  maxAttempts: 5,
+});
 
 const chain = sfn.Chain.start(custom).next(finalStatus);
 

packages/aws-cdk-lib/aws-stepfunctions/README.md

@@ -602,6 +602,13 @@ const custom = new sfn.CustomState(this, 'my custom task', {
 const errorHandler = new sfn.Pass(this, 'handle failure');
 custom.addCatch(errorHandler);
 
+// retry the task if something goes wrong
+custom.addRetry({
+  errors: [sfn.Errors.ALL],
+  interval: Duration.seconds(10),
+  maxAttempts: 5,
+});
+
 const chain = sfn.Chain.start(custom)
   .next(finalStatus);
 

packages/aws-cdk-lib/aws-stepfunctions/lib/states/custom-state.ts

@@ -1,7 +1,7 @@
 import { Construct } from 'constructs';
 import { State } from './state';
 import { Chain } from '..';
-import { CatchProps, IChainable, INextable } from '../types';
+import { CatchProps, IChainable, INextable, RetryProps } from '../types';
 
 /**
  * Properties for defining a custom state definition
@@ -34,6 +34,17 @@ export class CustomState extends State implements IChainable, INextable {
     this.stateJson = props.stateJson;
   }
 
+  /**
+   * Add retry configuration for this state
+   *
+   * This controls if and how the execution will be retried if a particular
+   * error occurs.
+   */
+  public addRetry(props: RetryProps = {}): CustomState {
+    super._addRetry(props);
+    return this;
+  }
+
   /**
    * Add a recovery handler for this state
    *

packages/aws-cdk-lib/aws-stepfunctions/test/custom-state.test.ts

@@ -120,4 +120,47 @@ describe('Custom State', () => {
       },
     );
   });
+
+  test('can add a retry state', () => {
+    // GIVEN
+    const custom = new sfn.CustomState(stack, 'Custom', {
+      stateJson,
+    });
+    const chain = sfn.Chain.start(custom);
+
+    // WHEN
+    custom.addRetry({
+      errors: [sfn.Errors.ALL],
+      interval: cdk.Duration.seconds(10),
+      maxAttempts: 5,
+    });
+
+    // THEN
+    expect(render(stack, chain)).toStrictEqual(
+      {
+        StartAt: 'Custom',
+        States: {
+          Custom: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::dynamodb:putItem',
+            Parameters: {
+              TableName: 'MyTable',
+              Item: {
+                id: {
+                  S: 'MyEntry',
+                },
+              },
+            },
+            ResultPath: null,
+            Retry: [{
+              ErrorEquals: ['States.ALL'],
+              IntervalSeconds: 10,
+              MaxAttempts: 5,
+            }],
+            End: true,
+          },
+        },
+      },
+    );
+  });
 });