(aws_stepfunctions_tasks): Unclear how LambdaInvoke works

[anentropic]

Describe the issue

My Lambda function returns JSON data

It's unclear what the return value of the step is though - is it just the JSON that the Lambda itself returns, or is it an "invoke lambda response" JSON? Or something else?

I'm using/attempting to use the Step Functions local testing tool
https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local-mock-cfg-file.html#mock-cfg-mckd-resp-sect:~:text=to%20the%20next.-,Return,-Return%20is%20represented

in the docs there they show example mocked response for a LambdaInvoke step as:

{
  "StatusCode": 200,
  "Payload": {
    "StatusCode": 200,
    "body": "Hello from Lambda!"
  }
}

This implies I may need to post-process the output of the lambda step with something like States.StringToJson($.Payload.body)?

There is a payloadResponseOnly option for the LambdaInvoke construct but again it's not clear what it does.

The docs say:

Invoke the Lambda in a way that only returns the payload response without additional metadata.

...which is super vague. Does it return $.Payload.body? Does it parse it as JSON too? Something else?

Whatever it does I am unable to test this option in the local mock tool because it causes it to need a Lambda arn rather than just the Lambda name(?)... I get errors like "CreateStateMachine <= Invalid State Machine Definition: ''SCHEMA_VALIDATION_FAILED: Value is not a valid resource ARN at /States/Process Batch Unit/Branches[0]/States/Invoke-ProcessBatchLambda/Resource" (...since there is no easy way to derive the ASL json from the CDK code I am having to use this other tool cdk-asl-extractor to reconstruct it from CF template, but there is no real ARN for the Lambda when testing locally)

Links

https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_stepfunctions_tasks.LambdaInvoke.html#payload

[pahud]

Yes this is a bit ambiguous.

Let's say if we submit this JSON input to the state machine:

{"foo":"123", "bar":"456"}

And our CDK code like this:

    const submitJobLambda = new lambda.Function(this, 'submitJobLambda', {
      code: lambda.Code.fromInline(`exports.handler = async (event, context) => {
            return {
              statusCode: '200',
              body: JSON.stringify(event),
            };
          };`),
      runtime: lambda.Runtime.NODEJS_LATEST,
      handler: 'index.handler',
    });

    const submitJob = new sfntasks.LambdaInvoke(this, 'Invoke Handler', {
      lambdaFunction: submitJobLambda,
      payload: sfn.TaskInput.fromJsonPathAt('$.foo'),
    });

The payload prop essentially only select and extract '$.foo' as the input to the lambda function as event. So you'll see the output like this because only 123 is passed to the lambda function.

If I just need the Payload and discard all other metadata:

    const submitJob = new sfntasks.LambdaInvoke(this, 'Invoke Handler', {
      lambdaFunction: submitJobLambda,
      payload: sfn.TaskInput.fromJsonPathAt('$.foo'),
      outputPath: '$.Payload',
    });

Now the output looks like this

I hope this helps and we welcome your pull requests to help us improve our document.

[github-actions]

This issue has not received a response in a while. If you want to keep this issue open, please leave a comment below and auto-close will be canceled.

[anentropic]

experimenting more with the mock server has helped, and it seems that thankfully I don't have to do any JSON decoding on the Lambda response

One thing still puzzles me

I want to repeat a value from the input at the output of my Lambda task (preferably without doing so from the Lambda side, because this value isn't relevant to that Lambda, only to subsequent steps in my chain)

        prepare_partitions_task = tasks.LambdaInvoke(
            self,
            "Invoke-PreparePartitionsLambda",
            lambda_function=cast(lambda_.IFunction, self._prepare_partitions_lambda),
            invocation_type=tasks.LambdaInvocationType.REQUEST_RESPONSE,
            result_selector={
                "partitioned.$": "$.Payload",
                "batch_size.$": "$.batch_size",
            },
        )

So here $.Payload is the output from the Lambda and "$.batch_size" is the input value I want to pass through

But I get:

The JSONPath '$.batch_size' specified for the field 'batch_size.$' could not be found in the input '{\"StatusCode\":200,\"Payload\":{\"job_group\":\"test_job_group\",\"job_names\":[\"test_job_group_1\",\"test_job_group_2\"]}}'"}}

So the input is not available at the output? (the "input" referenced in the error is just the mocked Lambda response)

I'm unclear if it's supposed to be available but it's up to me to manually add that to my mock server config (I hope not, it really feels like that should just be mocking the Lambda output) or if it's just not possible with a LambdaInvoke step

I had thought to maybe use $$.Execution.Input.batch_size but this fails because $$.Execution.Input is still a string of serialized JSON, and it doesn't seem to be possible to use States.JsonFromString($$.Execution.Input).batch_size i.e. to treat the result of deserializing as a JSONPath root ... I could achieve this with an extra Pass step, but feels like there ought to be a simpler way?

[anentropic]

if instead of using result_selector I do:

        prepare_partitions_task = tasks.LambdaInvoke(
            self,
            "Invoke-PreparePartitionsLambda",
            lambda_function=cast(lambda_.IFunction, self._prepare_partitions_lambda),
            invocation_type=tasks.LambdaInvocationType.REQUEST_RESPONSE,
            output_path="$.Payload",
            result_path="$.partitioned",
        )

...then I get:

Unable to apply ReferencePath $.partitioned to input \"{\\\"min_size\\\":10000,\\\"max_partitions\\\":100,\\\"batch_size\\\":100}\""}

where "input" referred in the error seems to be the contents of $.Payload

[anentropic]

This works on the mock server:

        prepare_partitions_task = tasks.LambdaInvoke(
            self,
            "Invoke-PreparePartitionsLambda",
            lambda_function=cast(lambda_.IFunction, self._prepare_partitions_lambda),
            invocation_type=tasks.LambdaInvocationType.REQUEST_RESPONSE,
            result_selector={
                "partitioned.$": "$.Payload",
                "execution_input.$": "States.StringToJson($$.Execution.Input)",
            },
        ).next(
            sfn.Pass(
                self,
                "Batch-size pass-thru",
                parameters={
                    "partitioned.$": "$.partitioned",
                    "batch_size.$": "$.execution_input.batch_size",
                }
            )
        )

however this doesn't work at all if I want to pass through some intermediate output from another step rather than the execution input as above

The only solution for that I can think of would be to use a Parallel state to wrap the LambdaInvoke as one branch and a Pass state as the other, then I'll be able to merge the results of both and carry on

I can do that - I just want to check if this is all expected behaviour and not an inconsistency in the mock server vs real executions?

[anentropic]

ugh, this part was my fault: "execution_input.$": "States.StringToJson($$.Execution.Input)"

I had ended up double serializing my json input - the mocked lambda responses were obscuring that from my other statemachine test (passing) so I missed that for a while when this one had all these issues

that also explains why result_path wasn't working... it couldn't merge new key into the input when the input was a string instead of an object, so it just silently replaced it instead

I still have a weird problem where output_path and result_path on the LambdaInvoke work fine individually, but not together

1. if I set just output_path="$.Payload" then I get an (expected) error from subsequent step that shows that the content of Lambda Payload was added at root of output data successfully
2. if I set just result_path="$.partitioned" then I get an (expected) error from subsequent step that shows the whole Lambda response i.e. {"StatusCode": 200, "Payload": {...}} was successfully added to output under partitioned key
3. if I set both then I get an unexpected error: An error occurred while executing the state 'Invoke-PreparePartitionsLambda' (entered at the event id #2). Invalid path '$.Payload' : No results for path: $['Payload']

result_selector works now but I have to explicitly name all the input values I want to propagate to the output, I was hoping I could avoid that by using output_path and result_path together

[anentropic]

I attached a debugger to the mock server and it looked like in the "both" case it had inserted the whole Lambda response under partitioned key and then was trying to apply $.Payload path to that

and if I do:

            output_path="$.partitioned.Payload",
            result_path="$.partitioned",

then in subsequent step I get:
Unable to apply step "partitioned" to input {"job_group":"test_job_group","job_names":["test_job_group_1","test_job_group_2"]}

which again looks as if the result_path is applied first and then the output_path

re-reading the OutputPath docs it seems this is intended behaviour, so my mistake again (seems like the other way around would be more useful though...)

Now I went back and tried again with the payload_response_only=True arg to LambdaInvoke

It looks like this actually does what I need!

i.e. effectively applies a $.Payload selector to the Lambda response before output_path sees it (so I don't need an output_path in my case)

The difficulty I had previously in using payload_response_only in my tests is related to how I extract the statemachine definition from the CDK stack, for which I'd been using https://github.com/nathanagez/aws-cdk-state-machine-asl which does a hacky parsing of CF template code

in particular this part:

const fnGetAtt = (expression) => {
    return expression[1] && expression[1] === 'Arn' ? expression[0] : expression.join('', expression)
}

just substitutes the resource name where an ARN is expected... this is fine for invocation_type=tasks.LambdaInvocationType.REQUEST_RESPONSE which uses a FunctionName key to specify the Lambda which accepts either ARN or name

but payload_response_only changes that and uses a Resource key which accepts only ARN

(it seems to result in a different task type, I get different events in the execution history: LambdaFunctionSucceeded when payload_response_only=True instead of TaskSucceeded)

I've now rewritten the ASL extractor in Python so that it generates fake ARN where expected:

def _fake_arn(value: str) -> str:
    return f"arn:aws:lambda:us-east-1:123456789012:function:{value}"


def fn_get_att_resolver(get_att_expr_values: list[str]) -> str:
    """
    Best effort to resolve `Fn::GetAtt` expression to string value.

    NOTE: we can't easily customise the fake ARN type, for now a
    fake Lambda ARN is returned because it's most useful for
    Step Functions local mock.
    """
    return (
        _fake_arn(get_att_expr_values[0])
        if (len(get_att_expr_values) > 0 and get_att_expr_values[1] == "Arn")
        else "".join(get_att_expr_values)
    )

i.e. I am using CDK Template.from_stack testing tool to get the StateMachine resource, then taking the resource["Properties"]["DefinitionString"] from that and passing it through the hacky CF template parser to get usable statemachine JSON, sending that to mock server in CreateStateMachine call etc

so I've eventually stumbled through all of the above into something that works enough