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"
  }
}

Answer 1

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' } 

Answer 2

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