MCP Server Step Functions Wait-for-Callback — task tokens, heartbeats, and human approval gates

Task token delivery: SQS queue as the handoff channel

The most common delivery mechanism for task tokens is an SQS queue. The state machine's Task state is defined with "Resource": "arn:aws:states:::sqs:sendMessage.waitForTaskToken" — the .waitForTaskToken suffix activates the pattern. Step Functions delivers a message to the specified SQS queue containing the task token in the TaskToken field of the message body, then pauses the execution. An external consumer reads the message, extracts the token, processes the work, and calls back.

// ASL Task state definition using SQS waitForTaskToken
const approvalTaskState = {
  "Type": "Task",
  "Resource": "arn:aws:states:::sqs:sendMessage.waitForTaskToken",
  "Parameters": {
    "QueueUrl": "https://sqs.us-east-1.amazonaws.com/123456789/approval-queue",
    "MessageBody": {
      "TaskToken.$": "$$.Task.Token",        // $$.Task.Token = the task token
      "Input.$": "$",                        // current execution input
      "ExecutionId.$": "$$.Execution.Id",
      "ExecutionName.$": "$$.Execution.Name",
      "StateName.$": "$$.State.Name",
    }
  },
  "HeartbeatSeconds": 3600,      // fail with States.HeartbeatTimeout if no heartbeat for 1 hour
  "TimeoutSeconds": 86400,       // fail with States.Timeout if not resolved within 24 hours
  "Catch": [
    {
      "ErrorEquals": ["States.HeartbeatTimeout", "States.Timeout"],
      "Next": "HandleApprovalTimeout",
      "ResultPath": "$.approvalError"
    }
  ],
  "ResultPath": "$.approvalResult", // merge result into execution state
  "Next": "ProcessApprovalDecision"
};

// Alternative: Lambda waitForTaskToken (Lambda receives token and stores it for later callback)
const lambdaWaitState = {
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke.waitForTaskToken",
  "Parameters": {
    "FunctionName": "arn:aws:lambda:us-east-1:123:function:approval-dispatcher",
    "Payload": {
      "TaskToken.$": "$$.Task.Token",
      "requestContext.$": "$"
    }
  },
  "HeartbeatSeconds": 1800,
  "TimeoutSeconds": 43200,
  "Next": "AfterApproval"
};

The task token value — accessible in ASL via $$.Task.Token — is available only within the Task state that activates waitForTaskToken. It cannot be read from other states or from GetExecutionHistory after the fact (for security reasons, the token is not stored in execution history).

Storing and retrieving task tokens: the DynamoDB handoff store

The external system that processes the approval or long-running work needs access to the task token after it has consumed and deleted the SQS message. If the SQS message is deleted before the work is complete (e.g., to prevent re-delivery), the token must be stored durably. DynamoDB with a TTL matching TimeoutSeconds is the standard pattern.

import {
  SFNClient,
  SendTaskSuccessCommand,
  SendTaskFailureCommand,
  SendTaskHeartbeatCommand,
} from '@aws-sdk/client-sfn';
import { DynamoDBClient, PutItemCommand, GetItemCommand, DeleteItemCommand } from '@aws-sdk/client-dynamodb';
import { marshall, unmarshall } from '@aws-sdk/util-dynamodb';
import { SQSClient, ReceiveMessageCommand, DeleteMessageCommand } from '@aws-sdk/client-sqs';

const sfn = new SFNClient({ region: 'us-east-1' });
const dynamo = new DynamoDBClient({ region: 'us-east-1' });
const sqs = new SQSClient({ region: 'us-east-1' });

const TOKEN_TABLE = process.env.TOKEN_TABLE!;     // DynamoDB table for task token storage
const APPROVAL_QUEUE_URL = process.env.APPROVAL_QUEUE_URL!;

// Consumer: read SQS message, store token, delete SQS message
async function consumeApprovalMessage() {
  const { Messages } = await sqs.send(new ReceiveMessageCommand({
    QueueUrl: APPROVAL_QUEUE_URL,
    MaxNumberOfMessages: 1,
    WaitTimeSeconds: 20,
  }));

  if (!Messages?.length) return null;

  const message = Messages[0];
  const body = JSON.parse(message.Body!);

  // Store task token in DynamoDB keyed by a stable business ID
  const requestId = body.Input?.requestId ?? body.ExecutionName;
  const ttlEpoch = Math.floor(Date.now() / 1000) + 86400; // 24-hour TTL

  await dynamo.send(new PutItemCommand({
    TableName: TOKEN_TABLE,
    Item: marshall({
      requestId,
      taskToken: body.TaskToken,
      executionId: body.ExecutionId,
      input: body.Input,
      createdAt: new Date().toISOString(),
      ttl: ttlEpoch,
    }),
    ConditionExpression: 'attribute_not_exists(requestId)', // prevent overwrite
  }));

  // Delete the SQS message after storing the token
  await sqs.send(new DeleteMessageCommand({
    QueueUrl: APPROVAL_QUEUE_URL,
    ReceiptHandle: message.ReceiptHandle!,
  }));

  return { requestId, input: body.Input, executionId: body.ExecutionId };
}

// Approve or reject an approval request by looking up the task token
async function resolveApproval(
  requestId: string,
  approved: boolean,
  reason?: string,
): Promise {
  const item = await dynamo.send(new GetItemCommand({
    TableName: TOKEN_TABLE,
    Key: marshall({ requestId }),
    ConsistentRead: true,
  }));

  if (!item.Item) {
    throw new Error(`No pending approval found for requestId: ${requestId}`);
  }

  const { taskToken } = unmarshall(item.Item);

  try {
    if (approved) {
      await sfn.send(new SendTaskSuccessCommand({
        taskToken,
        output: JSON.stringify({ approved: true, resolvedBy: 'mcp-tool', reason }),
      }));
    } else {
      await sfn.send(new SendTaskFailureCommand({
        taskToken,
        error: 'ApprovalDenied',
        cause: reason ?? 'Request was rejected',
      }));
    }
  } catch (err: any) {
    // TaskTimedOut = token expired; execution already failed — nothing to do
    if (err.name !== 'TaskTimedOut' && err.name !== 'InvalidToken') throw err;
  }

  // Clean up token from DynamoDB regardless of outcome
  await dynamo.send(new DeleteItemCommand({
    TableName: TOKEN_TABLE,
    Key: marshall({ requestId }),
  }));
}

// Heartbeat: keep a waiting task token alive during long external work
async function heartbeatApproval(requestId: string): Promise {
  const item = await dynamo.send(new GetItemCommand({
    TableName: TOKEN_TABLE,
    Key: marshall({ requestId }),
  }));

  if (!item.Item) return; // already resolved

  const { taskToken } = unmarshall(item.Item);

  try {
    await sfn.send(new SendTaskHeartbeatCommand({ taskToken }));
  } catch (err: any) {
    if (err.name === 'TaskTimedOut') {
      // Step Functions timed out the task — clean up local token record
      await dynamo.send(new DeleteItemCommand({
        TableName: TOKEN_TABLE,
        Key: marshall({ requestId }),
      }));
    }
  }
}

MCP tool: approval gate as a tool pair

The natural MCP tool design for waitForTaskToken is a two-tool pattern: one tool submits a request and returns a requestId; a second tool polls for resolution. The agent calls the submit tool, continues other work, then periodically calls the poll tool to check whether the approval was granted or denied.

// MCP tool: submit approval request (kicks off Step Functions execution)
// Returns a requestId the agent can use to poll for the decision
const submitApprovalTool = {
  name: 'submit_approval_request',
  description: 'Submit a request for human approval. Returns a requestId to poll for the decision.',
  inputSchema: {
    type: 'object',
    properties: {
      action: { type: 'string' },
      context: { type: 'object' },
      timeout_hours: { type: 'number', default: 24 },
    },
    required: ['action'],
  },
  async handler(input: { action: string; context?: unknown; timeout_hours?: number }) {
    const { StartExecutionCommand } = await import('@aws-sdk/client-sfn');
    const requestId = `req-${Date.now()}-${Math.random().toString(36).slice(2)}`;

    await sfn.send(new StartExecutionCommand({
      stateMachineArn: process.env.APPROVAL_SM_ARN!,
      name: requestId,
      input: JSON.stringify({ requestId, action: input.action, context: input.context }),
    }));

    return {
      requestId,
      status: 'pending',
      message: `Approval request submitted. Poll with requestId: ${requestId}`,
    };
  },
};

// MCP tool: check approval status
const checkApprovalTool = {
  name: 'check_approval_status',
  description: 'Check the status of a previously submitted approval request.',
  inputSchema: {
    type: 'object',
    properties: { requestId: { type: 'string' } },
    required: ['requestId'],
  },
  async handler(input: { requestId: string }) {
    const item = await dynamo.send(new GetItemCommand({
      TableName: TOKEN_TABLE,
      Key: marshall({ requestId: input.requestId }),
    }));

    if (item.Item) {
      return { requestId: input.requestId, status: 'pending' };
    }

    // Token gone from DynamoDB — look up execution result
    const { DescribeExecutionCommand } = await import('@aws-sdk/client-sfn');
    const exec = await sfn.send(new DescribeExecutionCommand({
      executionArn: `arn:aws:states:us-east-1:123:execution:approval-sm:${input.requestId}`,
    }));

    return {
      requestId: input.requestId,
      status: exec.status,
      output: exec.output ? JSON.parse(exec.output) : null,
    };
  },
};

Failure modes reference