How to get Task Token in AWS Lambda Function

Question

I have an AWS Step Function with the handlers implemented in Java.

My step function definition:

definition:
  Comment: Steps for issuing a card
  StartAt: RecipientFraudChecks
  States:
    RecipientFraudChecks:
      Type: Task
      Next: SaveTaskToken
      Resource: arn:aws:lambda:eu-west-1:099720403855:RecipientFraudChecks
    SaveTaskToken:
      Type: Task
      Resource: arn:aws:lambda:eu-west-1:12345678:function:SaveTaskToken
      End: true

I have a Java project and all the Lambda Function handlers are defined there:

public class SaveTaskToken implements RequestHandler<Map<String,String>, String> {

   ....

   @Override
   public String handleRequest(Map<String, String> input, final Context context) {

      // do the fraud checks
      System.out.println("the context is: " + gson.toJson(context));
      System.out.println("input: " + gson.toJson(input));

}

I'm running the step function locally using AWS SAM, and triggering according to this: https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local-lambda.html#install-sam

Within context I would expect to see the Task Token, but I do not. Logs show:

the context is: {
  "memoryLimit": 512,
  "awsRequestId": "5065a9aa-1a4a-46fe-9b58-7dc2194f92b7",
  "logGroupName": "aws/lambda/SaveTaskToken",
  "logStreamName": "$LATEST",
  "functionName": "SaveTaskToken",
  "functionVersion": "$LATEST",
  "invokedFunctionArn": "",
  "cognitoIdentity": {
    "identityId": "",
    "poolId": ""
  },
  "logger": {}
}

In fact its nothing like the global Context I should expect in the docs.

What am I doing wrong? How can I get the Task Token?

EDIT

I added Parameters property to 'SaveTaskToken' and changed resource to arn:aws:states:::lambda:invoke.waitForTaskToken and know I can get the Task Token:

definition:
  Comment: Steps for issuing a card
  StartAt: RecipientFraudChecks
  States:
    RecipientFraudChecks:
      Type: Task
      Next: SaveTaskToken
      Resource: arn:aws:lambda:eu-west-1:099720403855:RecipientFraudChecks
    SaveTaskToken:
      Type: Task
      Resource: arn:aws:states:::lambda:invoke.waitForTaskToken
      Parameters:
        FunctionName: arn:aws:lambda:eu-west-1:12345678:function:SaveTaskToken
        Payload:
          taskToken
      End: true

In the logs I can see:

the input is: {
  "taskToken": "5286"
}

It has caused another problem - it overrides the input to the state machine. Im passing in the input:

{"giftCode": "xxx"}

In the first Lambda function, RecipientFraudChecks, I can get the input. However, in the second, since adding the Parameters property, I now can no longer get the input to the state machine, only the task token...

EDIT Have implemented the answer here: https://stackoverflow.com/a/66995869/1246159

{
  "Comment": "Steps for issuing a card",
  "StartAt": "RecipientFraudChecks",
  "States": {
    "RecipientFraudChecks": {
      "Type": "Task",
      "Next": "PauseCardIfNecessary",
      "ResultPath": "$.firstLambdaOutput",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:RecipientFraudChecks"
    },
    "PauseCardIfNecessary": {
      "Type": "Task",
      "Next": "GetOrCreateClient",
      "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
      "Parameters": {
        "FunctionName": "arn:aws:lambda:us-east-1:123456789012:function:PauseCardIfNecessary",
        "Payload": {
          "token.$": "$$.Task.Token",
          "otherInput.$": "$"
        }
      }
    },
    "GetOrCreateClient": {
      "Type": "Task",
      "Next": "GetOrAccountClient",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:GetOrCreateClient"
    },
    "GetOrAccountClient": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:GetOrAccountClient",
      "End": true
    }
  }
}

But I get another error, here are the logs:

arn: aws: states: eu-west-1: 123456789012: execution: HelloWorld5: cardIssue: {
  "Type": "TaskStateExited",
  "PreviousEventId": 5,
  "StateExitedEventDetails": {
    "Name": "RecipientFraudChecks",
    "Output": "{\"inputToStep\":\"xxxx\",\"firstLambdaOutput\":\"output of recipient lambda\"}"
  }
} arn: aws: states: eu-west-1: 123456789012: execution: HelloWorld5: cardIssue: {
  "Type": "TaskStateEntered",
  "PreviousEventId": 6,
  "StateEnteredEventDetails": {
    "Name": "PauseCardIfNecessary",
    "Input": "{\"inputToStep\":\"xxxx\",\"firstLambdaOutput\":\"output of recipient lambda\"}"
  }
} arn: aws: states: eu-west-1: 123456789012: execution: HelloWorld5: cardIssue: {
  "Type": "ExecutionFailed",
  "PreviousEventId": 7,
  "ExecutionFailedEventDetails": {
    "Error": "States.Runtime",
    "Cause": "An error occurred while executing the state 'PauseCardIfNecessary' (entered at the event id #7). The value for the field 'token.$' must be a valid JSONPath expression"
  }
}

Balu Vyamajala · Accepted Answer · 2021-04-08T00:53:25.860

Task token is not automatically passed in lambda context, it needs to be passed as input.

The context here is not context of lambda function, but the context of the task and to grab details from context, we can use $$. within step function definition, Ex:

To get Task Token $$.Task.Token
To get start time $$.Execution.StartTime

Example Step function:

First Task executes a Lambda with step function input.
Appends the output of first lambda with step function input.
Second Task executes another lambda with resource waitForTaskToken, right here, we grab the task token and pass it along with output of previous step as input.
Step functions waits until it gets a SendTaskSuccess or SendTaskFailure

{
  "StartAt": "fist-lambda-invoke-sync",
  "States": {
    "fist-lambda-invoke-sync": {
      "Next": "second-lambda-invoke-task-token",
      "Type": "Task",
      "ResultPath": "$.firstLambdaOutput",
      "Resource": "arn:aws:lambda:us-east-1:11112223333:function:myfirstlambda"
    },
    "second-lambda-invoke-task-token": {
      "End": true,
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
      "Parameters": {
        "FunctionName": "arn:aws:lambda:us-east-1:11112223333:function:mysecondlambda",
        "Payload": {
          "token.$": "$$.Task.Token",
          "otherInput.$": "$"
        }
      }
    }
  }
}

Input to Step function:

{
  "inputToStep": "myValue"
}

Output of First Lambda: assuming it is a string value "10". It will be appended to input Json because of "ResultPath": "$.firstLambdaOutput"

{
  "inputToStep": "myValue",
  "firstLambdaOutput": "10"
}

Input to Second Lambda: Will receive below json as input with task token appended, because of "token.$": "$$.Task.Token" and "otherInput.$": "$"

{ otherInput: { inputToStep: 'myValue', firstLambdaOutput: '10' },
  token: 'This is where Task Token generated by step function will be sent' }

Thanks, I implemented this but getting more errors, will update question — Mark, Apr 08 '21 at 10:56
@Mark Are you still testing this locally? because, I tested both my definition and your definition from AWS console, both are working as expected. — Balu Vyamajala, Apr 08 '21 at 16:16
@Mark are you using express step function or standard step function? — Balu Vyamajala, Apr 08 '21 at 16:28
@Mark can you confirm where you are testing, i tested start to finish. All Lambdas I am calling are simple 1 line lambdas, which just prints the input and returns success. For the wait step , i called `aws stepfunctions send-task-success --task-token AAAAKgAA... --task-output '{"out":"success"}'` while it was waiting, that completed the wait step. — Balu Vyamajala, Apr 09 '21 at 11:49
@Mark Also are you getting error when creating the step function using cloudformation/sam or you are getting the error while running the step function? — Balu Vyamajala, Apr 09 '21 at 11:51
Sorry forgot to mark as correct, in the end your solution worked, thank you;) — Mark, Apr 14 '21 at 06:26

score 1 · Answer 2 · answered Apr 06 '21 at 20:11

Regarding the input, there seems to be an easy solution in the docs [1]:

Instead of hard-coding the event payload in the state machine definition, you can use the input from the state machine execution. The following example uses the input specified when you run the state machine as the event payload:
"Payload.$": "$"

There is another example in the docs [2]:

{  
   "StartAt":"GetManualReview",
   "States":{  
      "GetManualReview":{  
         "Type":"Task",
         "Resource":"arn:aws:states:::lambda:invoke.waitForTaskToken",
         "Parameters":{  
            "FunctionName":"get-model-review-decision",
            "Payload":{  
               "model.$":"$.new_model",
               "token.$":"$$.Task.Token"
            },
            "Qualifier":"prod-v1"
         },
         "End":true
      }
   }
}

You can possibly change "model.$":"$.new_model" to something like "input.$":"$" and get the desired Lambda payload.

Which translates to the following in YAML:

Parameters:
  Payload:
    input.$: "$"
    token.$: "$$.Task.Token"

[1] https://docs.aws.amazon.com/lambda/latest/dg/services-stepfunctions.html#services-stepfunctions-setup
[2] https://docs.amazonaws.cn/en_us/step-functions/latest/dg/connect-lambda.html

Thanks, implemented something close to this but still getting errors, updated question — Mark, Apr 08 '21 at 11:49

How to get Task Token in AWS Lambda Function

2 Answers2