如何在等待任务令牌的Step Functions Lambda调用状态中处理异常？

我有一个Step Functions状态机定义，其中一个状态使用arn:aws:states:::lambda:invoke.waitForTaskToken调用Lambda并等待任务令牌。我已添加Catch规则以捕获所有异常，状态定义如下：

"Invok Handler Lambda": {
        "Type": "Task",
        "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
        "Parameters": {
            "FunctionName": "${my_lambda}",
            "Payload": {
                "TaskToken.$": "$$.Task.Token",
                "Input.$": "$"
            }
        },
        "ResultPath": null,
        "Next": "Invoke Completer Lambda",
        "TimeoutSeconds": 600,
        "Retry": [
            {
                "ErrorEquals": [
                    "ErrorA",
                    "ErrorB"
                ],
                "IntervalSeconds": 1,
                "BackoffRate": 2,
                "MaxAttempts": 2
            },
            {
                "ErrorEquals": [
                    "ErrorC"
                ],
                "IntervalSeconds": 5
            }
        ],
        "Catch": [
            {
                "ErrorEquals": [
                    "States.ALL"
                ],
                "Next": "Handle Error"
            }
        ]
    }

状态机运行正常，但当my_lambda自身因某些原因失败时，不会进入Handle Error状态。请问是否有办法让Step Functions在此情况下进入错误处理状态？

出现这个问题的核心原因是：使用waitForTaskToken模式时，Step Functions默认只会等待任务令牌的回调（成功或失败），如果Lambda函数本身执行失败（比如未返回任务令牌就抛出异常），Step Functions不会自动触发Catch规则，因为它认为任务还在等待中。

你可以通过以下两种方法解决：

方法1：在Lambda函数内部捕获异常并主动回调任务令牌 #

修改你的my_lambda代码，在函数执行失败时，调用SendTaskFailure API，将错误信息和任务令牌一起发送给Step Functions。示例Python代码：

import boto3
import json

client = boto3.client('stepfunctions')

def lambda_handler(event, context):
    try:
        # 你的业务逻辑代码
        raise Exception("Lambda执行失败")
    except Exception as e:
        # 捕获异常后主动发送失败回调
        client.send_task_failure(
            taskToken=event['TaskToken'],
            error="LambdaExecutionFailed",
            cause=str(e)
        )
        raise e  # 可选，保持Lambda的失败状态记录

这样Step Functions会收到失败回调，进而触发Catch规则进入Handle Error状态。

方法2：为Lambda任务添加Lambda调用错误的捕获规则 #

Step Functions的waitForTaskToken资源会将Lambda调用本身的错误封装为Lambda.ServiceException、Lambda.AWSLambdaException或Lambda.SdkClientException这类错误类型，你可以在Catch规则中显式添加这些错误类型，确保Lambda执行失败时被捕获：

修改状态定义的Catch规则：

"Catch": [
    {
        "ErrorEquals": [
            "Lambda.ServiceException",
            "Lambda.AWSLambdaException",
            "Lambda.SdkClientException"
        ],
        "Next": "Handle Error"
    },
    {
        "ErrorEquals": [
            "States.ALL"
        ],
        "Next": "Handle Error"
    }
]

注意：如果你的Retry规则中包含了这些Lambda错误类型，需要确保MaxAttempts设置合理，避免重试后才进入Catch。当前你的Retry规则只处理ErrorA、ErrorB、ErrorC，所以Lambda相关错误不会被重试，会直接触发Catch。

两种方法可以结合使用，确保无论是Lambda执行失败未回调，还是回调失败的情况，都能被Step Functions正确捕获并进入错误处理状态。

内容的提问来源于stack exchange，提问作者Md. Shohag Mia