AWS StepFunction Lambda wait for token - when does wait occur?

Question

AWS StepFunctions support a callback model. The documentation reads:

Callback tasks provide a way to pause a workflow until a task token is returned. A task might need to wait for a human approval, integrate with a third party, or call legacy systems. For tasks like these, you can pause Step Functions indefinitely, and wait for an external process or workflow to complete. For these situations Step Functions allows you to pass a task token to some integrated services. The task will pause until it receives that task token back with a SendTaskSuccess or SendTaskFailure call.

If a state machine is defined using CDK to call a lambda as part of a Step Function map (i.e. it's in a loop):

LambdaInvoke myLambdaStepFunction = LambdaInvoke.Builder.create(this, "MyLambdaStepFunction ")
                                                      .lambdaFunction(myLambdaFunction)
                                                      .integrationPattern(IntegrationPattern.WAIT_FOR_TASK_TOKEN)
                                                      .build();

stepfunctions.Map loopSteps = stepfunctions.Map.Builder.create(this, "LambdaLoop")
                                                    .maxConcurrency(1)
                                                    .build();
loopSteps.iterator(myLambdaStepFunction);

StateMachine stateMachine = StateMachine.Builder.create(this, "StateMachine")
                                             .stateMachineType(StateMachineType.STANDARD)
                                             .timeout(Duration.minutes(5))
                                             .definition(deployAllIncrementallySteps)
                                             .build();

when does the wait for token occur?

Is it after the lambda is invoked for the first time?
Is it before the lambda is invoked for the first time?

Hi Oscar, yes, details below: TL;DR the wait for the token occurs after the Lambda has executed and before the next step is able to occur. — John

John John · Accepted Answer · 2020-08-05T08:44:24

The wait for the token occurs after the Lambda with the IntegrationPattern.WAIT_FOR_TASK_TOKEN integration has executed. Here's the pattern:

Your lambda when called will have a token passed to it by AWS in the Step Function event if you specify it in the payload, e.g. here I specify the token to be in the JSON payload supplied to the Lambda at the path of $.token and the remainder of the input at $.input:

LambdaInvoke myLambdaStep= LambdaInvoke.Builder.create(this, "MyLambdaStep")
                                .lambdaFunction(myLambdaFunction)
                                .integrationPattern(IntegrationPattern.WAIT_FOR_TASK_TOKEN)
                                .payload(TaskInput.fromObject(Map.of("token", JsonPath.getTaskToken(), "input", JsonPath.stringAt("$.input"))))
                                .build();

Within your Lambda function, you are responsible for parsing the token from the event and passing this token to some outside agent: at some point, that agent needs to call sendTaskSuccess or sendTaskFailure with that token to indicate that the work being carried out by that agent is complete. The external agent can be anything at all (another Lambda, an EC2 server doing some work for you, an external webhook, anything - it just needs to make the call back to AWS when its work is complete).
When AWS receives the call from your agent, it will call the appropriate next step in your step-machine, based on the success or failure call made by the agent.

AWS StepFunction Lambda wait for token - when does wait occur?

1 Answers