Skip to content
me.lasindu.com
All writing

Orchestrating AI workflows on AWS with Step Functions

, 5 min read

A single model call fits easily in a Lambda function. Real AI workflows tend to have several steps: draft something, check it, wait for a person to approve it, then act on it. Those steps can fail or wait for days, and a function that has to keep running in the meantime isn't a good fit.

This guide covers which AWS service to choose, then builds a draft-and-approve workflow in Step Functions.

Workflow or agent?

The first question is who decides the next step.

  • In a workflow, you decide. The steps are known in advance and the model does the work inside them. Workflows are predictable, and they're easier to test and audit.
  • In an agent, the model decides. It chooses which tools to call and when to stop. This suits open-ended tasks, but behaviour and cost are harder to predict.

If you can draw the process as a flowchart, start with a workflow. You can always make one step of it an agent later.

Choosing a service

A note on Bedrock Agents. Many older tutorials use it, but it's now called Bedrock Agents Classic and has been closed to new customers since July 30, 2026. Accounts without recent usage get an AccessDeniedException when creating agents. Existing agents keep working, but new projects should use AgentCore.

Example: draft, approve, publish

This workflow has the model write a draft, waits for a person to approve it, then publishes it. It needs a Standard workflow, because Express workflows don't support waiting for a callback.

{
  "QueryLanguage": "JSONata",
  "StartAt": "WriteDraft",
  "States": {
    "WriteDraft": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Arguments": {
        "ModelId": "us.anthropic.claude-sonnet-5",
        "ContentType": "application/json",
        "Accept": "application/json",
        "Body": {
          "anthropic_version": "bedrock-2023-05-31",
          "max_tokens": 2048,
          "messages": [
            {
              "role": "user",
              "content": "{% 'Write a short release note for: ' & $states.input.topic %}"
            }
          ]
        }
      },
      "Retry": [
        {
          "ErrorEquals": ["States.TaskFailed"],
          "IntervalSeconds": 3,
          "MaxAttempts": 3,
          "BackoffRate": 2,
          "JitterStrategy": "FULL"
        }
      ],
      "Assign": {
        "draft": "{% $join($states.result.Body.content[type = 'text'].text) %}"
      },
      "Next": "RequestApproval"
    },
    "RequestApproval": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
      "Arguments": {
        "FunctionName": "request-approval",
        "Payload": {
          "draft": "{% $draft %}",
          "taskToken": "{% $states.context.Task.Token %}"
        }
      },
      "TimeoutSeconds": 259200,
      "Catch": [
        {
          "ErrorEquals": ["States.Timeout"],
          "Next": "NotApproved"
        }
      ],
      "Next": "CheckDecision"
    },
    "CheckDecision": {
      "Type": "Choice",
      "Choices": [
        {
          "Condition": "{% $states.input.approved = true %}",
          "Next": "Publish"
        }
      ],
      "Default": "NotApproved"
    },
    "Publish": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "Arguments": {
        "FunctionName": "publish-draft",
        "Payload": {
          "draft": "{% $draft %}"
        }
      },
      "End": true
    },
    "NotApproved": {
      "Type": "Fail",
      "Error": "NotApproved",
      "Cause": "The draft was rejected or the approval timed out."
    }
  }
}

Some details worth noting:

  • Filter for text blocks. The expression in Assign keeps only blocks with type of text. Reading content[0] directly is common in examples, but it can break on models that return a thinking block first.
  • The model call and the wait are separate states. WriteDraft runs as a normal request and finishes when the model responds. The pause happens in RequestApproval.
  • Retries are deliberately broad. States.TaskFailed also retries errors that will never succeed, such as a malformed request. Once you've seen the real error names in your execution history, narrow the list.
  • Step Functions doesn't check the request body. A body that's wrong for the model only fails at runtime. Also, Body is limited to 256 KiB, so use the Input and Output S3 fields for anything larger.
  • The IAM role needs your inference profile. The execution role must allow bedrock:InvokeModel on the profile and on the foundation model, as described in the Lambda post.

The approval step

When the workflow reaches RequestApproval, it passes a task token to a Lambda function and then pauses. That function shows the draft to a person by whatever channel you use, such as email or an internal tool, and stores the token. Lambda doesn't stay running during the wait.

When the person decides, a second function resumes the workflow:

import json

import boto3

client = boto3.client("stepfunctions")


def handler(event, context):
    # Your API must authenticate the approver before this function runs.
    body = json.loads(event["body"])

    client.send_task_success(
        taskToken=body["taskToken"],
        output=json.dumps({"approved": body["approved"] is True}),
    )
    return {"statusCode": 200, "body": json.dumps({"status": "recorded"})}

The value in output becomes the input to CheckDecision.

Treat the task token like a password, because anyone who has it can complete the step. Don't put it in an unauthenticated link, and keep it out of your logs.

Always set a timeout on callback states. Without one, a task can wait until the execution reaches its one-year limit. This example gives up after three days. For a callback from an automated process, HeartbeatSeconds also helps you detect a stuck worker sooner.

Adding an agent as one step

You don't have to choose between a workflow and an agent. Step Functions has an optimized AgentCore integration, so one state can hand an open-ended task to an agent while the workflow controls everything around it.

That integration is request-response only, so a human approval still needs its own callback state, as above. For agent tasks that run longer than a single state allows, AWS describes some asynchronous patterns.

Rules that hold up in production

  • Don't let model output make irreversible decisions alone. Put a validation step or a person before anything that spends money, sends messages or deletes data.
  • Make steps safe to retry. A retried Publish shouldn't publish twice, so use an idempotency key based on the execution ID.
  • Keep prompts and model IDs out of the definition when you can. They change more often than the workflow does.
  • Log token usage per execution. Multi-step workflows multiply model costs, and retries add to them.

AWS service details and limits were checked against the AWS documentation in November 2026. Bedrock changes quickly, so confirm model IDs and availability for your Region.