[aws-stepfunctions]: Next step in catch block of Custom step not being generated in cfn template.

[nivedita104]

Issue

I have a use case where I am using a custom step in one of our state machines and in the catch block, i want to add a new step that is not previously referenced in the state machine. When I tried deploying this, the statemachine throws an error that
"Invalid State Machine Definition: 'MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithou
tItemsDataset/Catch[0]/Next, MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithItemsDataset/Catch[0]/Next' (Service: AWSSte
pFunctions; Status Code: 400; Error Code: InvalidDefinition; Request ID: 77eb25f8-02b7-49b3-8cdd-b1ff1921ec5f; Proxy: null)" (RequestToken: 92e222c1-409c-cfb0-558a-15
f75529c9c7, HandlerErrorCode: InvalidRequest)"

Here, testLambda was the step I introduced in the catch block of the custom step.

I checked the generated template it was not generating the "testLambda" step. Hence i tried to use the .addCatch() function call that we use with the other steps and i got an error that the .addCatch handler is not supported for custom steps.

What is the solution for this? Can we not add new steps for catch block in the custom steps or am i missing something here?

Steps to Reproduce:

Default Custom state example in docs:

import * as dynamodb from '@aws-cdk/aws-dynamodb';

// create a table
const table = new dynamodb.Table(this, 'montable', {
  partitionKey: {
    name: 'id',
    type: dynamodb.AttributeType.STRING,
  },
});

const finalStatus = new sfn.Pass(this, 'final step');

// States language JSON to put an item into DynamoDB
// snippet generated from https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-code-snippet.html#tutorial-code-snippet-1
const stateJson = {
  Type: 'Task',
  Resource: 'arn:aws:states:::dynamodb:putItem',
  Parameters: {
    TableName: table.tableName,
    Item: {
      id: {
        S: 'MyEntry',
      },
    },
  },
  ResultPath: null,
};

// custom state which represents a task to insert data into DynamoDB
const custom = new sfn.CustomState(this, 'my custom task', {
  stateJson,
});

const chain = sfn.Chain.start(custom)
  .next(finalStatus);

const sm = new sfn.StateMachine(this, 'StateMachine', {
  definition: chain,
  timeout: Duration.seconds(30),
});

// don't forget permissions. You need to assign them
table.grantWriteData(sm);

Now, we add a new pass state and catch block to the custom step :

const errorHandlerLambda = new Pass(
        this,
        "errorHandlerLambda"
    );

const stateJson = {
  Type: 'Task',
  Resource: 'arn:aws:states:::dynamodb:putItem',
  Parameters: {
    TableName: table.tableName,
    Item: {
      id: {
        S: 'MyEntry',
      },
    },
  },
Catch: [
          {
            ErrorEquals: [Errors.ALL],
            ResultPath: "$.JobDetails.errorInfo",
            Next: 'errorHandlerLambda'
          }
        ],
  ResultPath: null,
};

This will throw the error I mentioned above.
I cannot add the errorHandlerLambda to the step function definition as it will only be used as the next step in case the custom step fails.