はじめに
AWS Step Functionsは依存性のあるタスクや複雑な処理をステートマシンと呼ばれるワークフローによって宣言的に定義できるサーバーレスサービスです。特にプログラミング言語と比べて、エラーハンドリングやリトライ処理が非常に宣言的に実装しやすいです。
Step Functionsのラムダなどを実行するタスクステートにはwaitForTaskTokenという機能があり、これを使うとタスクトークンが発行され、そのトークンを使ってSendTaskSuccessが呼び出されるまでワークフローを一時停止できます。(SendTaskFailureでワークフローを失敗させて、Catchブロックなどに移ります）また、SendTaskSuccessおよびSendTaskFailureでは、タスクトークンだけでなく、任意のjson文字列をフィードバックとしてoutputフィールドに持たせて送ることができるので、そのフィードバックをタスクステートで分析して、Checkステート（条件分岐を行えるステート）で処理を分けるといったこともできます。さらに、Step Functionsはステートの遷移回数によって課金されるので、待機期間では一切料金がかからないのも大きな利点です。
しかし、タスクトークンは、承認画面や承認リンクの先にいるユーザーに渡すには扱いにくい値です。公式ドキュメントでは、SendTaskSuccess のトークンの長さは最大2048文字とされています。URLやメールにそのまま載せるのは現実的ではありません。
そこで、タスクトークンをDynamoDBにユーザーIDと紐づけて保存し、ユーザーには承認ID(短い識別子)だけを見せる構成を紹介します。API GatewayとLambdaが、承認IDからトークンを引いて処理を再開します。ユーザーはトークンを知らなくても承認できます。
なお、今回の構成では承認用LambdaからStep FunctionsのAPI(SendTaskSuccess / SendTaskFailure)を呼び出す必要があります。
前提
-
ワークフロータイプ: Standard(
.waitForTaskTokenは、Standardのみ対応) - ステートマシン: Amazon States Language、クエリ言語はJSONata
- Lambda: Python(AWS SDKはboto3を使用)
- API Gateway: HTTP API + CognitoのJWTオーソライザー
- リージョン: ap-northeast-1(例)
アーキテクチャ構成図と処理フロー
処理の流れ:
- Step Functions(Standardワークフロー)の実行を開始します。
-
lambda:invoke.waitForTaskTokenのステートが、保存用Lambdaを呼び出します。このとき、コンテキストオブジェクトからタスクトークンを渡します。 - 保存用Lambdaが、承認ID、承認者のユーザーID、タスクトークンをDynamoDBに保存し、承認者に通知します。ワークフローはここで一時停止します。
- 承認者が、API Gateway経由で承認IDを指定して承認または却下を送信します。
- 承認用Lambdaが、DynamoDBから承認IDに対応するトークンを取得し、承認者本人かを確認したうえで、
SendTaskSuccessまたはSendTaskFailureを呼び出します。 - ワークフローが再開します。
コールバックパターン(
.waitForTaskToken)は、Standardワークフローでのみ使えます。Expressワークフローは対象外です。また、DynamoDBの最適化統合は.waitForTaskTokenに対応していません(AWS SDK統合は、Standardワークフローで対応しています)。本記事では、保存後の通知処理を同じLambdaにまとめるため、Lambdaを経由してトークンを保存します。
実装のコア部分(設定とコード)
1. ステートマシン定義(JSONata)
$states.context.Task.Token でタスクトークンを取得し、Lambdaのペイロードに渡します。
{
"QueryLanguage": "JSONata",
"StartAt": "RequestApproval",
"States": {
"RequestApproval": {
"Type": "Task",
"Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
"Arguments": {
"FunctionName": "arn:aws:lambda:ap-northeast-1:123456789012:function:save-task-token",
"Payload": {
"approvalId": "{% $states.input.approvalId %}",
"approverUserId": "{% $states.input.approverUserId %}",
"taskToken": "{% $states.context.Task.Token %}"
}
},
"TimeoutSeconds": 604800,
"Next": "Approved",
"Catch": [
{
"ErrorEquals": ["Rejected"],
"Next": "Rejected"
}
]
},
"Approved": { "Type": "Succeed" },
"Rejected": { "Type": "Fail", "Error": "Rejected" }
}
}
TimeoutSeconds を設定すると、承認されないまま放置された実行を、指定時間で失敗させられます(上の例は7日間)。
2. トークンを保存するLambda
承認IDをパーティションキー、トークンとユーザーIDを属性として保存します。
import os
import time
import boto3
table = boto3.resource("dynamodb").Table(os.environ["TABLE_NAME"])
def handler(event, context):
table.put_item(
Item={
"approvalId": event["approvalId"],
"approverUserId": event["approverUserId"],
"taskToken": event["taskToken"],
# 実行の最大期間より後に消えるよう、TTLを設定する
"expiresAt": int(time.time()) + 8 * 24 * 3600,
}
)
# ここで承認者への通知(メール、Slackなど)を送る
このLambdaは、保存が終わると終了します。それでも、.waitForTaskToken のステートは、SendTaskSuccess または SendTaskFailure が呼ばれるまで待機し続けます。
DynamoDBのTTLを使うには、テーブルでTTLを有効にし、TTL属性として expiresAt(エポック秒)を指定しておく必要があります。有効にしていないと、expiresAt を入れても項目は自動では削除されません。有効にしても、期限切れの項目は通常、期限から数日以内に削除されます。削除までは残るため、TTLは即時の失効ではなく、取り残しの掃除として使います。
3. 承認を受け付けるLambda(API Gateway経由)
承認者のユーザーIDは、リクエストボディではなく、認証済みのクレームから取得します。ユーザー自身にIDを送らせると、なりすませるためです。ここでは、API GatewayのHTTP APIに、CognitoのJWTオーソライザーを設定した構成を想定しています(JWTオーソライザーはHTTP APIの機能で、クレームは event["requestContext"]["authorizer"]["jwt"]["claims"] から取得できます)。
import json
import os
import boto3
table = boto3.resource("dynamodb").Table(os.environ["TABLE_NAME"])
sfn = boto3.client("stepfunctions")
def response(status, body):
return {"statusCode": status, "body": json.dumps(body)}
def handler(event, context):
approval_id = event["pathParameters"]["approvalId"]
user_id = event["requestContext"]["authorizer"]["jwt"]["claims"]["sub"]
decision = json.loads(event["body"])["decision"] # "approve" or "reject"
item = table.get_item(Key={"approvalId": approval_id}).get("Item")
if item is None or item["approverUserId"] != user_id:
return response(404, {"message": "not found"})
try:
if decision == "approve":
sfn.send_task_success(
taskToken=item["taskToken"],
output=json.dumps({"approvedBy": user_id}),
)
else:
sfn.send_task_failure(
taskToken=item["taskToken"],
error="Rejected",
cause=f"rejected by {user_id}",
)
except (sfn.exceptions.TaskTimedOut, sfn.exceptions.InvalidToken):
return response(409, {"message": "already closed or expired"})
table.delete_item(Key={"approvalId": approval_id})
return response(200, {"message": "ok"})
承認者本人でない場合は、存在を隠すために404を返しています。
4. IAMの権限
承認用Lambdaには、Step Functionsのコールバック用APIの権限が必要です。
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": ["states:SendTaskSuccess", "states:SendTaskFailure"],
"Resource": "*"
},
{
"Effect": "Allow",
"Action": ["dynamodb:GetItem", "dynamodb:DeleteItem"],
"Resource": "arn:aws:dynamodb:ap-northeast-1:123456789012:table/approval-tokens"
}
]
}
保存用Lambdaには dynamodb:PutItem のみを付与します。
パフォーマンスと制限(クォータ)への配慮
-
トークンの長さ:
SendTaskSuccessのtaskTokenは、最大2048文字です。DynamoDBに保存するのは文字列なので問題ありませんが、ユーザーに直接見せる設計は避けます。 -
待機できる期間: 一時停止できるのは、実行が1年のサービスクォータに達するまでです。実運用では、
TimeoutSecondsで、承認期限を明示的に決めておきます。 -
タイムアウト時の挙動: コールバック用のタスクがタイムアウトすると、新しいランダムなトークンが生成されます。DynamoDBに残った古いトークンは使えなくなり、
SendTaskSuccessはTaskTimedOutで失敗します。そのため、コード例では409を返しています。 -
二重承認: 一度閉じたタスクのトークンに再度送ると、
TaskTimedOutになります。承認後はDynamoDBのレコードを削除し、TTLで取り残しも掃除します。 -
ペイロードサイズ:
SendTaskSuccessのoutputは、最大262144バイトです。
おわりに
タスクトークンをDynamoDBにユーザーIDと紐づけて保存することで、ユーザーにはトークンを見せずに、承認ステップを実装できます。承認者の確認も、トークンの管理も、サーバー側に寄せられるのが利点です。(承認用LambdaからStep FunctionsのAPIを呼び出す必要があり、states:SendTaskSuccess などの権限管理が必要になります。)
また、SendTaskSuccess,SendTaskFailureでは、outputフィールドで任意のjson文字列をフィードバックとして送ることができるので、単なる承認と拒否を行うのではなく、HITLを取り入れる手段として活用されることが多いのではないかと思います。
(私が個人的に作成した短編小説生成ワークフローはその一例で、プロットや各章の文章作成に対して、ユーザーがAPIからフィードバックを返せるようにしています。AI推論のコスト最適化など、まだまだ改善点が多いですが、興味があればぜひ覗いてみてほしいです。)