(aws-stepfunctions): CustomState generates wrong code when using Type: Choice

[cwoolum]

Describe the bug

I am using the CustomState construct to generate a Choice with JSONata. The problem is that since CustomState implements INextable, this adds an End property which is invalid for the Choice state.

Regression Issue

• Select this option if this issue appears to be a regression.

Last Known Working CDK Version

No response

Expected Behavior

For this CDK construct

 const shouldContinue = new sfn.CustomState(this, 'ShouldContinue', {
      stateJson: {
        Type: 'Choice',
        QueryLanguage: 'JSONata',
        Default: 'AnalysisComplete',
        Choices: [
          {
            Next: 'ProcessDiscoveredUrls',
            Condition:
              '{% $count($states.input.nextBatch) > 0 and $count($states.input.visitedUrls) <= $states.input.maxPages %}',
          },
        ],
      },
    });

I would expect, the proper JSON to be constructed without the End property.

Current Behavior

This is the generated JSON

{
  "End": true,
  "Type": "Choice",
  "QueryLanguage": "JSONata",
  "Default": "AnalysisComplete",
  "Choices":[{
    "Next": "ProcessDiscoveredUrls",
    "Condition": "{% $count($states.input.nextBatch) > 0 and $count($states.input.visitedUrls) <= $states.input.maxPages %}"
  }]
}

Reproduction Steps

This is easily reproduced using the construct above.

Possible Solution

One solution would be adding an if statement to conditionally call renderNextEnd() in toStateJson() but then CustomState wouldn't really implement INextable. There needs to be a construct one level higher that doesn't implement INextable.

Additional Information/Context

No response

CDK CLI Version

2.173.2 (build f8e6207)

Framework Version

No response

Node.js Version

20

OS

AL2

Language

TypeScript

Language Version

No response

Other information

No response

[khushail]

Hi @cwoolum , thanks for reaching out. Your analysis stands true as when I ran the given code, with using customState and choice, the end:true is appended in the template -

"MyStateMachine6C968CA5": {
   "Type": "AWS::StepFunctions::StateMachine",
   "Properties": {
    "DefinitionString": {
     "Fn::Join": [
      "",
      [
       "{\"StartAt\":\"MyLambdaTask\",\"States\":{\"MyLambdaTask\":{\"Next\":\"ShouldContinue\",\"Retry\":[{\"ErrorEquals\":[\"Lambda.ClientExecutionTimeoutException\",\"Lambda.ServiceException\",\"Lambda.AWSLambdaException\",\"Lambda.SdkClientException\"],\"IntervalSeconds\":2,\"MaxAttempts\":6,\"BackoffRate\":2}],\"Type\":\"Task\",\"Resource\":\"arn:",
       {
        "Ref": "AWS::Partition"
       },
       ":states:::lambda:invoke\",\"Parameters\":{\"FunctionName\":\"",
       {
        "Fn::GetAtt": [
         "MyLambdaFunction67CCA873",
         "Arn"
        ]
       },
       "\",\"Payload.$\":\"$\"}},\"ShouldContinue\":{\"End\":true,\"Type\":\"Choice\",\"QueryLanguage\":\"JSONata\",\"Default\":\"AnalysisComplete\",\"Choices\":[{\"Next\":\"ProcessDiscoveredUrls\",\"Condition\":\"{% $count($states.input.nextBatch) > 0 and $count($states.input.visitedUrls) <= $states.input.maxPages %}\"}]}}}"
      ]
     ]
    },

However as per ASL, it should not have End:true. Also mentioned in AWS DOCS here-

Important

Choice states don't support the End field. In addition, they use Next only inside their Choices field.

I did not find any workaround for this but yes, this seems to be a bug. So marking it as P2 which means it might not be immediately addressed by the team but would be on their radar. Also keeping it open for community contribution.