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
- Step Functions: the default for workflows across several AWS services. You get a visual graph, execution history and optimized integrations, including Bedrock.
- Lambda durable functions: the same reliability ideas, written as ordinary code. They checkpoint progress, can run for up to a year and don't charge for compute while waiting. They suit logic that's tightly coupled to your application code. AWS has a comparison with Step Functions.
- Amazon Bedrock AgentCore: for agents where the model runs the loop. Its managed harness lets you declare the model, tools and instructions, and you can deploy your own agent code if you need more control.
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
Assignkeeps only blocks withtypeoftext. Readingcontent[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.
WriteDraftruns as a normal request and finishes when the model responds. The pause happens inRequestApproval. - Retries are deliberately broad.
States.TaskFailedalso 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,
Bodyis limited to 256 KiB, so use theInputandOutputS3 fields for anything larger. - The IAM role needs your inference profile. The execution role must allow
bedrock:InvokeModelon 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
Publishshouldn'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.