sfn-saga-compensation-orchestrator

import os, time, boto3

ENDPOINT = os.environ.get("AWS_ENDPOINT_URL")
ddb = boto3.resource("dynamodb", endpoint_url=ENDPOINT)
inv = ddb.Table("inventory")
idem = ddb.Table("saga-idempotency")


def lambda_handler(event, context):
    key = event["idempotency_key"]
    step_key = f"{key}:reserve"

    if idem.get_item(Key={"idempotency_key": step_key}).get("Item"):
        return {"status": "already_reserved", "sku": event["sku"], "qty": event["qty"]}

    sku = event["sku"]
    qty = int(event["qty"])
    inv.update_item(
        Key={"sku": sku},
        UpdateExpression="ADD reserved_qty :q",
        ExpressionAttributeValues={":q": qty},
    )

    idem.put_item(
        Item={
            "idempotency_key": step_key,
            "ttl": int(time.time()) + 86400,
            "step": "reserve",
        }
    )
    return {"status": "reserved", "sku": sku, "qty": qty}

import os, time, boto3

ENDPOINT = os.environ.get("AWS_ENDPOINT_URL")
ddb = boto3.resource("dynamodb", endpoint_url=ENDPOINT)
sm = boto3.client("secretsmanager", endpoint_url=ENDPOINT)
payments = ddb.Table("payments")
idem = ddb.Table("saga-idempotency")


def lambda_handler(event, context):
    if event.get("force_failure"):
        raise RuntimeError("forced payment failure")

    key = event["idempotency_key"]
    step_key = f"{key}:charge"

    if idem.get_item(Key={"idempotency_key": step_key}).get("Item"):
        return {"status": "already_charged", "payment_id": event["order_id"]}

    secret = sm.get_secret_value(SecretId="payment/gateway-apikey")
    print(f"charge-payment: loaded gateway key (len={len(secret['SecretString'])})")

    payment_id = event["order_id"]
    payments.put_item(
        Item={
            "payment_id": payment_id,
            "order_id": event["order_id"],
            "amount": int(event["amount"]),
            "status": "CAPTURED",
        }
    )

    idem.put_item(
        Item={
            "idempotency_key": step_key,
            "ttl": int(time.time()) + 86400,
            "step": "charge",
        }
    )
    return {"status": "CAPTURED", "payment_id": payment_id}

import os, time, boto3

ENDPOINT = os.environ.get("AWS_ENDPOINT_URL")
ddb = boto3.resource("dynamodb", endpoint_url=ENDPOINT)
shipments = ddb.Table("shipments")
idem = ddb.Table("saga-idempotency")


def lambda_handler(event, context):
    key = event["idempotency_key"]
    step_key = f"{key}:ship"

    if idem.get_item(Key={"idempotency_key": step_key}).get("Item"):
        return {"status": "already_shipped", "shipment_id": event["order_id"]}

    shipment_id = event["order_id"]
    shipments.put_item(
        Item={
            "shipment_id": shipment_id,
            "order_id": event["order_id"],
            "sku": event["sku"],
            "qty": int(event["qty"]),
            "status": "CREATED",
        }
    )

    idem.put_item(
        Item={
            "idempotency_key": step_key,
            "ttl": int(time.time()) + 86400,
            "step": "ship",
        }
    )
    return {"status": "CREATED", "shipment_id": shipment_id}

import os, time, boto3

ENDPOINT = os.environ.get("AWS_ENDPOINT_URL")
ddb = boto3.resource("dynamodb", endpoint_url=ENDPOINT)
inv = ddb.Table("inventory")
idem = ddb.Table("saga-idempotency")


def lambda_handler(event, context):
    key = event["idempotency_key"]
    comp_key = f"{key}:release"

    if idem.get_item(Key={"idempotency_key": comp_key}).get("Item"):
        return {"status": "already_released"}

    reserve_marker = idem.get_item(Key={"idempotency_key": f"{key}:reserve"}).get("Item")
    if reserve_marker:
        sku = event["sku"]
        qty = int(event["qty"])
        inv.update_item(
            Key={"sku": sku},
            UpdateExpression="ADD reserved_qty :q",
            ExpressionAttributeValues={":q": -qty},
        )

    idem.put_item(
        Item={
            "idempotency_key": comp_key,
            "ttl": int(time.time()) + 86400,
            "step": "release",
        }
    )
    return {"status": "released"}

import os, time, boto3

ENDPOINT = os.environ.get("AWS_ENDPOINT_URL")
ddb = boto3.resource("dynamodb", endpoint_url=ENDPOINT)
payments = ddb.Table("payments")
idem = ddb.Table("saga-idempotency")


def lambda_handler(event, context):
    key = event["idempotency_key"]
    comp_key = f"{key}:refund"

    if idem.get_item(Key={"idempotency_key": comp_key}).get("Item"):
        return {"status": "already_refunded"}

    charge_marker = idem.get_item(Key={"idempotency_key": f"{key}:charge"}).get("Item")
    if charge_marker:
        payment_id = event["order_id"]
        payments.update_item(
            Key={"payment_id": payment_id},
            UpdateExpression="SET #s = :s",
            ExpressionAttributeNames={"#s": "status"},
            ExpressionAttributeValues={":s": "REFUNDED"},
        )

    idem.put_item(
        Item={
            "idempotency_key": comp_key,
            "ttl": int(time.time()) + 86400,
            "step": "refund",
        }
    )
    return {"status": "refunded"}

def lambda_handler(event, context):
    detail_type = event.get("detail-type", "Unknown")
    print(f"saga-terminal:{detail_type}")
    return {"ok": True}

{
  "query": "select:TodoWrite",
  "max_results": 1
}

, output, 
{"type": "tool_reference", "tool_name": "TodoWrite"}

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {"Service": "lambda.amazonaws.com"},
      "Action": "sts:AssumeRole"
    }
  ]
}

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {"Service": "states.amazonaws.com"},
      "Action": "sts:AssumeRole"
    }
  ]
}

{
  "Comment": "order-saga: 3 forward steps with reverse-order compensation",
  "StartAt": "ReserveInventory",
  "States": {
    "ReserveInventory": {
      "Type": "Task",
      "Resource": "__RESERVE_ARN__",
      "ResultPath": "$.reserveResult",
      "Catch": [
        {
          "ErrorEquals": ["States.ALL"],
          "ResultPath": "$.error",
          "Next": "PublishFailed"
        }
      ],
      "Next": "ChargePayment"
    },
    "ChargePayment": {
      "Type": "Task",
      "Resource": "__CHARGE_ARN__",
      "ResultPath": "$.chargeResult",
      "Catch": [
        {
          "ErrorEquals": ["States.ALL"],
          "ResultPath": "$.error",
          "Next": "CompReleaseInventoryOnly"
        }
      ],
      "Next": "CreateShipment"
    },
    "CreateShipment": {
      "Type": "Task",
      "Resource": "__SHIP_ARN__",
      "ResultPath": "$.shipResult",
      "Catch": [
        {
          "ErrorEquals": ["States.ALL"],
          "ResultPath": "$.error",
          "Next": "CompRefundPayment"
        }
      ],
      "Next": "PublishCompleted"
    },
    "CompRefundPayment": {
      "Type": "Task",
      "Resource": "__REFUND_ARN__",
      "ResultPath": "$.refundResult",
      "Next": "CompReleaseInventoryAfterRefund"
    },
    "CompReleaseInventoryAfterRefund": {
      "Type": "Task",
      "Resource": "__RELEASE_ARN__",
      "ResultPath": "$.releaseResult",
      "Next": "PublishFailed"
    },
    "CompReleaseInventoryOnly": {
      "Type": "Task",
      "Resource": "__RELEASE_ARN__",
      "ResultPath": "$.releaseResult",
      "Next": "PublishFailed"
    },
    "PublishCompleted": {
      "Type": "Task",
      "Resource": "arn:aws:states:::events:putEvents",
      "Parameters": {
        "Entries": [
          {
            "EventBusName": "saga-events",
            "Source": "order.saga",
            "DetailType": "OrderCompleted",
            "Detail": "{}"
          }
        ]
      },
      "ResultPath": null,
      "End": true
    },
    "PublishFailed": {
      "Type": "Task",
      "Resource": "arn:aws:states:::events:putEvents",
      "Parameters": {
        "Entries": [
          {
            "EventBusName": "saga-events",
            "Source": "order.saga",
            "DetailType": "OrderFailed",
            "Detail": "{}"
          }
        ]
      },
      "ResultPath": null,
      "End": true
    }
  }
}

#!/usr/bin/env bash
set -euo pipefail

export AWS_ACCESS_KEY_ID=test
export AWS_SECRET_ACCESS_KEY=test
export AWS_DEFAULT_REGION=us-east-1
export AWS_PAGER=""
ENDPOINT="http://localstack:4566"
AWSL="aws --endpoint-url $ENDPOINT"
REGION=us-east-1
ACCOUNT=000000000000

cd /app

echo "==> KMS CMK + alias"
KEY_ID=$($AWSL kms create-key --description "saga CMK" --query 'KeyMetadata.KeyId' --output text)
KEY_ARN="arn:aws:kms:${REGION}:${ACCOUNT}:key/${KEY_ID}"
# alias may exist from a prior run; delete first
$AWSL kms delete-alias --alias-name alias/saga-cmk 2>/dev/null || true
$AWSL kms create-alias --alias-name alias/saga-cmk --target-key-id "$KEY_ID"
echo "    KEY_ARN=$KEY_ARN"

echo "==> Secrets Manager (CMK-encrypted)"
$AWSL secretsmanager delete-secret --secret-id payment/gateway-apikey --force-delete-without-recovery 2>/dev/null || true
SECRET_ARN=$($AWSL secretsmanager create-secret \
  --name payment/gateway-apikey \
  --kms-key-id "$KEY_ARN" \
  --secret-string "sk_test_abcdef0123456789" \
  --query 'ARN' --output text)
echo "    SECRET_ARN=$SECRET_ARN"

echo "==> DynamoDB tables"
for t in inventory payments shipments saga-idempotency; do
  $AWSL dynamodb delete-table --table-name "$t" 2>/dev/null || true
done
# wait briefly for deletes to settle
sleep 1

$AWSL dynamodb create-table --table-name inventory \
  --attribute-definitions AttributeName=sku,AttributeType=S \
  --key-schema AttributeName=sku,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST >/dev/null

$AWSL dynamodb create-table --table-name payments \
  --attribute-definitions AttributeName=payment_id,AttributeType=S \
  --key-schema AttributeName=payment_id,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST >/dev/null

$AWSL dynamodb create-table --table-name shipments \
  --attribute-definitions AttributeName=shipment_id,AttributeType=S \
  --key-schema AttributeName=shipment_id,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST >/dev/null

$AWSL dynamodb create-table --table-name saga-idempotency \
  --attribute-definitions AttributeName=idempotency_key,AttributeType=S \
  --key-schema AttributeName=idempotency_key,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST >/dev/null

$AWSL dynamodb update-time-to-live --table-name saga-idempotency \
  --time-to-live-specification "Enabled=true,AttributeName=ttl" >/dev/null

INV_ARN="arn:aws:dynamodb:${REGION}:${ACCOUNT}:table/inventory"
PAY_ARN="arn:aws:dynamodb:${REGION}:${ACCOUNT}:table/payments"
SHIP_ARN="arn:aws:dynamodb:${REGION}:${ACCOUNT}:table/shipments"
IDEM_ARN="arn:aws:dynamodb:${REGION}:${ACCOUNT}:table/saga-idempotency"

echo "==> Lambda execution roles (least-privilege, each to its own table)"

create_role_with_policy() {
  local role_name="$1"
  local policy_json="$2"
  $AWSL iam delete-role-policy --role-name "$role_name" --policy-name inline 2>/dev/null || true
  $AWSL iam delete-role --role-name "$role_name" 2>/dev/null || true
  $AWSL iam create-role --role-name "$role_name" \
    --assume-role-policy-document file:///app/policies/lambda-trust.json >/dev/null
  echo "$policy_json" > "/tmp/${role_name}-policy.json"
  $AWSL iam put-role-policy --role-name "$role_name" --policy-name inline \
    --policy-document "file:///tmp/${role_name}-policy.json" >/dev/null
  echo "arn:aws:iam::${ACCOUNT}:role/${role_name}"
}

LOGS='{"Effect":"Allow","Action":["logs:CreateLogGroup","logs:CreateLogStream","logs:PutLogEvents"],"Resource":"*"}'

RESERVE_ROLE_ARN=$(create_role_with_policy "reserve-inventory-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["dynamodb:UpdateItem","dynamodb:GetItem","dynamodb:PutItem"],"Resource":["${INV_ARN}","${IDEM_ARN}"]}
]}
EOF
)")

CHARGE_ROLE_ARN=$(create_role_with_policy "charge-payment-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["dynamodb:PutItem","dynamodb:GetItem","dynamodb:UpdateItem"],"Resource":["${PAY_ARN}","${IDEM_ARN}"]},
  {"Effect":"Allow","Action":["secretsmanager:GetSecretValue"],"Resource":"${SECRET_ARN}"},
  {"Effect":"Allow","Action":["kms:Decrypt"],"Resource":"${KEY_ARN}"}
]}
EOF
)")

SHIP_ROLE_ARN=$(create_role_with_policy "create-shipment-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["dynamodb:PutItem","dynamodb:GetItem"],"Resource":["${SHIP_ARN}","${IDEM_ARN}"]}
]}
EOF
)")

RELEASE_ROLE_ARN=$(create_role_with_policy "release-inventory-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["dynamodb:UpdateItem","dynamodb:GetItem","dynamodb:PutItem"],"Resource":["${INV_ARN}","${IDEM_ARN}"]}
]}
EOF
)")

REFUND_ROLE_ARN=$(create_role_with_policy "refund-payment-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["dynamodb:UpdateItem","dynamodb:GetItem","dynamodb:PutItem"],"Resource":["${PAY_ARN}","${IDEM_ARN}"]}
]}
EOF
)")

EVENT_BUS_ARN="arn:aws:events:${REGION}:${ACCOUNT}:event-bus/saga-events"
NOTIFIER_ROLE_ARN=$(create_role_with_policy "saga-notifier-role" "$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  ${LOGS},
  {"Effect":"Allow","Action":["events:PutEvents"],"Resource":"${EVENT_BUS_ARN}"}
]}
EOF
)")

echo "    roles created"

echo "==> Package lambdas"
mkdir -p /app/build
for fn in reserve_inventory charge_payment create_shipment release_inventory refund_payment saga_notifier; do
  (cd /app/lambdas && zip -q -j "/app/build/${fn}.zip" "${fn}.py")
done

create_lambda() {
  local name="$1"
  local file="$2"
  local role="$3"
  $AWSL lambda delete-function --function-name "$name" 2>/dev/null || true
  $AWSL lambda create-function \
    --function-name "$name" \
    --runtime python3.11 \
    --role "$role" \
    --handler "${file}.lambda_handler" \
    --zip-file "fileb:///app/build/${file}.zip" \
    --environment "Variables={AWS_ENDPOINT_URL=${ENDPOINT}}" \
    --timeout 30 >/dev/null
  echo "arn:aws:lambda:${REGION}:${ACCOUNT}:function:${name}"
}

echo "==> Deploy lambdas"
RESERVE_ARN=$(create_lambda reserve-inventory reserve_inventory "$RESERVE_ROLE_ARN")
CHARGE_ARN=$(create_lambda charge-payment charge_payment "$CHARGE_ROLE_ARN")
SHIP_LAMBDA_ARN=$(create_lambda create-shipment create_shipment "$SHIP_ROLE_ARN")
RELEASE_ARN=$(create_lambda release-inventory release_inventory "$RELEASE_ROLE_ARN")
REFUND_ARN=$(create_lambda refund-payment refund_payment "$REFUND_ROLE_ARN")
NOTIFIER_ARN=$(create_lambda saga-notifier saga_notifier "$NOTIFIER_ROLE_ARN")
echo "    lambdas deployed"

echo "==> EventBridge bus + rule + target"
$AWSL events delete-event-bus --name saga-events 2>/dev/null || true
$AWSL events create-event-bus --name saga-events >/dev/null

$AWSL events remove-targets --event-bus-name saga-events --rule on-saga-terminal --ids 1 2>/dev/null || true
$AWSL events delete-rule --event-bus-name saga-events --name on-saga-terminal 2>/dev/null || true

$AWSL events put-rule \
  --name on-saga-terminal \
  --event-bus-name saga-events \
  --event-pattern '{"source":["order.saga"],"detail-type":["OrderCompleted","OrderFailed"]}' >/dev/null

RULE_ARN="arn:aws:events:${REGION}:${ACCOUNT}:rule/saga-events/on-saga-terminal"

# allow eventbridge to invoke notifier
$AWSL lambda remove-permission --function-name saga-notifier --statement-id ebInvoke 2>/dev/null || true
$AWSL lambda add-permission \
  --function-name saga-notifier \
  --statement-id ebInvoke \
  --action lambda:InvokeFunction \
  --principal events.amazonaws.com \
  --source-arn "$RULE_ARN" >/dev/null

$AWSL events put-targets \
  --event-bus-name saga-events \
  --rule on-saga-terminal \
  --targets "Id=1,Arn=${NOTIFIER_ARN}" >/dev/null
echo "    event bus/rule ready"

echo "==> State machine role"
# scoped: invoke only the six saga lambdas, putEvents only on saga-events bus
SFN_POLICY=$(cat <<EOF
{"Version":"2012-10-17","Statement":[
  {"Effect":"Allow","Action":["lambda:InvokeFunction"],"Resource":[
    "${RESERVE_ARN}","${CHARGE_ARN}","${SHIP_LAMBDA_ARN}","${RELEASE_ARN}","${REFUND_ARN}"
  ]},
  {"Effect":"Allow","Action":["events:PutEvents"],"Resource":"${EVENT_BUS_ARN}"},
  {"Effect":"Allow","Action":["logs:CreateLogGroup","logs:CreateLogStream","logs:PutLogEvents"],"Resource":"*"}
]}
EOF
)
$AWSL iam delete-role-policy --role-name order-saga-role --policy-name inline 2>/dev/null || true
$AWSL iam delete-role --role-name order-saga-role 2>/dev/null || true
$AWSL iam create-role --role-name order-saga-role \
  --assume-role-policy-document file:///app/policies/sfn-trust.json >/dev/null
echo "$SFN_POLICY" > /tmp/sfn-policy.json
$AWSL iam put-role-policy --role-name order-saga-role --policy-name inline \
  --policy-document file:///tmp/sfn-policy.json >/dev/null
SFN_ROLE_ARN="arn:aws:iam::${ACCOUNT}:role/order-saga-role"

echo "==> State machine"
sed -e "s#__RESERVE_ARN__#${RESERVE_ARN}#g" \
    -e "s#__CHARGE_ARN__#${CHARGE_ARN}#g" \
    -e "s#__SHIP_ARN__#${SHIP_LAMBDA_ARN}#g" \
    -e "s#__RELEASE_ARN__#${RELEASE_ARN}#g" \
    -e "s#__REFUND_ARN__#${REFUND_ARN}#g" \
    /app/asl.template.json > /app/build/asl.json

# clean up an existing state machine if present
EXIST=$($AWSL stepfunctions list-state-machines --query "stateMachines[?name=='order-saga'].stateMachineArn" --output text || true)
if [[ -n "$EXIST" && "$EXIST" != "None" ]]; then
  $AWSL stepfunctions delete-state-machine --state-machine-arn "$EXIST" >/dev/null
  sleep 1
fi

SM_ARN=$($AWSL stepfunctions create-state-machine \
  --name order-saga \
  --definition file:///app/build/asl.json \
  --role-arn "$SFN_ROLE_ARN" \
  --query 'stateMachineArn' --output text)
echo "    SM_ARN=$SM_ARN"

echo "ALL_PROVISIONED"

#!/usr/bin/env bash
set -euo pipefail
export AWS_PAGER=""
ENDPOINT="http://localstack:4566"
AWSL="aws --endpoint-url $ENDPOINT"
SM_ARN="arn:aws:states:us-east-1:000000000000:stateMachine:order-saga"

run_and_wait() {
  local name="$1"
  local input="$2"
  local exec_arn
  exec_arn=$($AWSL stepfunctions start-execution \
    --state-machine-arn "$SM_ARN" \
    --name "$name" \
    --input "$input" \
    --query 'executionArn' --output text)
  echo "exec=$exec_arn"
  for i in $(seq 1 60); do
    status=$($AWSL stepfunctions describe-execution --execution-arn "$exec_arn" --query 'status' --output text)
    if [[ "$status" != "RUNNING" ]]; then
      echo "status=$status after ${i}s"
      return
    fi
    sleep 1
  done
  echo "TIMEOUT"; return 1
}

echo "== HAPPY PATH =="
run_and_wait "happy-$(date +%s)" '{"order_id":"o-1","sku":"x","qty":2,"amount":100,"idempotency_key":"k-1"}'

echo "== FAILURE PATH =="
run_and_wait "fail-$(date +%s)" '{"order_id":"o-2","sku":"x","qty":2,"amount":100,"idempotency_key":"k-2","force_failure":true}'

echo "== INVENTORY =="
$AWSL dynamodb get-item --table-name inventory --key '{"sku":{"S":"x"}}' | cat
echo "== PAYMENTS o-1 =="
$AWSL dynamodb get-item --table-name payments --key '{"payment_id":{"S":"o-1"}}' | cat
echo "== PAYMENTS o-2 =="
$AWSL dynamodb get-item --table-name payments --key '{"payment_id":{"S":"o-2"}}' | cat
echo "== SHIPMENTS o-1 =="
$AWSL dynamodb get-item --table-name shipments --key '{"shipment_id":{"S":"o-1"}}' | cat
echo "== SHIPMENTS o-2 =="
$AWSL dynamodb get-item --table-name shipments --key '{"shipment_id":{"S":"o-2"}}' | cat

echo "== NOTIFIER LOGS =="
groups=$($AWSL logs describe-log-groups --log-group-name-prefix /aws/lambda/saga-notifier --query 'logGroups[].logGroupName' --output text)
for g in $groups; do
  streams=$($AWSL logs describe-log-streams --log-group-name "$g" --query 'logStreams[].logStreamName' --output text)
  for s in $streams; do
    $AWSL logs get-log-events --log-group-name "$g" --log-stream-name "$s" --query 'events[].message' --output text
  done
done