Can I use a lambda's output in a state machine's Assign when using waitForTaskToken?

Question

I have a JSONata state machine with a lambda task step:

```json
"My lambda step": {
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
  "Arguments": {
    "FunctionName": "arn:aws:lambda:::function:my-lambda",
    "Payload": "{% $states.input %}"
  },
  "Assign": {
    "myLambdaResult": "{% $states.result.Payload.myLambdaResult %}"
  },
  "Next": "Send ingested event"
}
```

The lambda code returns an object with `myLambdaResult` which I'd like to use in the `Assign` block. However I'm using `waitForTaskToken`, which is waiting for an external service to notify this state machine when some async work has completed. The required task output property from this notifications is overriding the output from the lambda code, so `Payload` is undefined and I can't seem to access `myLambdaResult` from the state machine context.

Is there a way I can access the lambda's output here using JSONata (or JSONPath)?

re:Post Agent · Answer

When using the `waitForTaskToken` pattern with Lambda in Step Functions, the task result structure differs from a standard Lambda invocation.

With `waitForTaskToken`, the Lambda function initially returns, but the task doesn't complete until an external process calls the Step Functions API with the task token. When this happens, the response from the `SendTaskSuccess` API call (not the original Lambda return value) becomes the task result.

In your case, when the external service calls `SendTaskSuccess`, it should include the `myLambdaResult` value in the payload it sends. The task output will contain only what is passed in the `output` parameter of the `SendTaskSuccess` API call.

To make this work, your external process needs to include the Lambda's original output when it calls `SendTaskSuccess`. For example, if your Lambda returns `{ "myLambdaResult": "some value" }`, your external process should call `SendTaskSuccess` with this same structure in the output parameter.

Then in your state machine, you can access it directly:

```json
"Assign": {
"myLambdaResult": "{% $states.result.myLambdaResult %}"
}
```

The key point is that with `waitForTaskToken`, the final task output is determined by what's passed to `SendTaskSuccess`, not by the initial Lambda return value.  
**Sources**  
[Invoke an AWS Lambda function with Step Functions - AWS Step Functions](https://docs.aws.amazon.com/step-functions/latest/dg/connect-lambda.html)  
[Discovering workflow states to use in Step Functions - AWS Step Functions](https://docs.aws.amazon.com/step-functions/latest/dg/workflow-states.html)

Can I use a lambda's output in a state machine's Assign when using waitForTaskToken?

Add your answer

Relevant content