aws_stepfunctions_tasks: payloadResponseOnly cannot be used with TaskInput.fromJsonPathAt

[cscanlin]

Describe the bug

@jakubknejzlik Described the issue perfectly here in a now closed issue:
#22390 (comment)
(Copied for visibility)

CDK code:

new LambdaInvoke(scope, id, {
    lambdaFunction: ...,
    payload: TaskInput.fromJsonPathAt("$.input"),
    payloadResponseOnly: true,
})

produces following definition:

"state": {
     ...
      "Type": "Task",
      "Resource": "arn:aws:lambda:{region}:{account}:function:{function_name}",
      "Parameters": "$.input" // this is the problem
}

When payloadResponseOnly: false it produces following definition:

"state": {
     ...
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "Parameters": {
        "FunctionName": "arn:aws:lambda:{region}:{account}:function:{function_name}",
        "Payload.$": "$.input"
      }
}

When I use TaskInput.fromObject like so:

new LambdaInvoke(scope, id, {
    lambdaFunction: ...,
    payload: TaskInput.fromObject({
      input: JsonPath.stringAt("$.input"),
    }),
    payloadResponseOnly: true,
})

The definition is ok:

"state": {
     ...
      "Type": "Task",
      "Resource": "arn:aws:lambda:{region}:{account}:function:{function_name}",
      "Parameters": {
        "input.$": "$.input"
      }
}

And another closed issue:
#23996

Regression Issue

• Select this option if this issue appears to be a regression.

Last Known Working CDK Version

No response

Expected Behavior

See above

Current Behavior

See above

Reproduction Steps

See above

Possible Solution

See above

Additional Information/Context

No response

CDK CLI Version

2.1012.0

Framework Version

No response

Node.js Version

22.15.0

OS

MacOS

Language

Python

Language Version

3.10

Other information

No response

[pahud]

Hi @cscanlin,

Thanks for reporting this issue and providing clear reproduction details, including references to previous related issues (#22390, #23996).

You've correctly identified that when using the LambdaInvoke task with payloadResponseOnly: true and payload: TaskInput.fromJsonPathAt(...), the generated Step Functions Parameters field incorrectly receives the literal JSONPath string instead of using it to reference the input state.

Summary:

The core problem is the handling of TaskInput.fromJsonPathAt specifically when payloadResponseOnly is set to true. The construct should generate Parameters that use the JSONPath to select input (e.g., { "Payload.$": "$.input" }), but instead, it generates Parameters: "$.input".

Root Cause Analysis:

The issue stems from the _renderTask method within the LambdaInvoke class. When payloadResponseOnly is true, the code directly uses the payload.value (which is the JSONPath string "$.input" in this case) as the parameter argument, instead of wrapping it in the necessary structure like { "Payload.$": "$.input" }.

aws-cdk/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/lambda/invoke.ts

Lines 176 to 194 in 7f378b6

	protected _renderTask(topLevelQueryLanguage?: sfn.QueryLanguage): any {
	const queryLanguage = sfn._getActualQueryLanguage(topLevelQueryLanguage, this.props.queryLanguage);
	const [resource, paramOrArg] = this.props.payloadResponseOnly ?
	[this.props.lambdaFunction.functionArn, this.props.payload?.value] :
	[integrationResourceArn('lambda', 'invoke', this.integrationPattern), {
	FunctionName: this.props.lambdaFunction.functionArn,
	Payload: this.props.payload?.value ??
	(queryLanguage === sfn.QueryLanguage.JSONATA ? '{% $states.input %}'
	: sfn.TaskInput.fromJsonPathAt('$').value),
	InvocationType: this.props.invocationType,
	ClientContext: this.props.clientContext,
	Qualifier: this.props.qualifier,
	}];
	return {
	Resource: resource,
	...this._renderParametersOrArguments(paramOrArg, queryLanguage),
	};
	}
	}

Possible Solution:

The logic within _renderTask needs adjustment. When payloadResponseOnly is true and the payload is derived from TaskInput.fromJsonPathAt, it should construct a parameter object that correctly incorporates the JSONPath (e.g., { "Payload.$": payload.value }) before passing it to _renderParametersOrArguments.

Off the top of my head, here's how the _renderTask method could be modified:

 /**
   * @internal
   */
  protected _renderTask(topLevelQueryLanguage?: sfn.QueryLanguage): any {
    const queryLanguage = sfn._getActualQueryLanguage(topLevelQueryLanguage, this.props.queryLanguage);

    let resource: string;
    let parameters: any;

    if (this.props.payloadResponseOnly) {
      resource = this.props.lambdaFunction.functionArn;
      // Check if the payload is specifically from fromJsonPathAt
      // (This check might need refinement based on internal TaskInput implementation details)
      // A simple check could be if payload.type is TaskInputType.JSON_PATH
      // or if payload.value is a string starting with '$'
      if (this.props.payload && typeof this.props.payload.value === 'string' && this.props.payload.value.startsWith('$')) {
         // If it's a JSONPath, wrap it correctly for Parameters
         parameters = { 'Payload.$': this.props.payload.value };
      } else {
         // Otherwise, use the value directly (for TaskInput.fromObject etc.)
         parameters = this.props.payload?.value;
      }
    } else {
      // Original logic for payloadResponseOnly: false
      resource = integrationResourceArn('lambda', 'invoke', this.integrationPattern);
      parameters = {
        FunctionName: this.props.lambdaFunction.functionArn,
        Payload: this.props.payload?.value ??
          (queryLanguage === sfn.QueryLanguage.JSONATA ? '{% $states.input %}'
            : sfn.TaskInput.fromJsonPathAt('$').value),
        InvocationType: this.props.invocationType,
        ClientContext: this.props.clientContext,
        Qualifier: this.props.qualifier,
      };
    }

    return {
      Resource: resource,
      ...this._renderParametersOrArguments(parameters, queryLanguage),
    };
  }

Explanation of the change:

• Inside the if (this.props.payloadResponseOnly) block, we add a check.
• This check determines if this.props.payload.value looks like a JSONPath string (e.g., starts with $). Note: A more robust check might involve inspecting an internal property of TaskInput if available, to definitively know it came from fromJsonPathAt.
• If it is a JSONPath, we create the parameters object as { 'Payload.$': this.props.payload.value }.
• If it's not a JSONPath (e.g., it came from TaskInput.fromObject), we use the payload.value directly as before.
• This ensures the correct Parameters structure is generated for the payloadResponseOnly: true case when using TaskInput.fromJsonPathAt.

This illustrates the necessary logic adjustment.

We appreciate you bringing this to our attention. We encourage contributions! If you're interested in submitting a Pull Request to fix this, the relevant code is linked above. Please let us know if you have any further questions.