Step Function - Assign command is not assigning variables

Question

Hi,

I am currently working in a StepMachine to run our custom OpenAI Assistant. There are some common data (like thread_id and run_id) to nearly all states, and to avoid to keep passing the same data with ResultPath, i have added a "Assign" after the relevant tasks (https://docs.aws.amazon.com/step-functions/latest/dg/workflow-variables.html). The task is called (ExecyteShortInput), it runs completely on the lambda, but the variables are not assigned to the next state (StartExecution)

I am doing something wrong here on the Assign command, or in the lambda output params? The same output on the lambda function is called correctly by the previous state, so I am not sure what is wrong here. The variables should be assigned on the end of the "ExecuteShortInputWorkflow", and be available on the beggining of the "StartExecution" state as far as I could understand from the documentation. However, on the StartExecution task I am receiving the following input: `{'params': '$params', 'token_counters': '$token_counters', 'thread_id': '$threadId', 'x_user_email': '$xUserEmail', 'x_user_token': '$xUserToken', 'result_call': '$resultCall', 'mode': 'start_run'}`, that is, the values are not being assigned at all. Tried to put a Pass state in between, but got the same result. Can anyone help me on this issue, please? I am doing something wrong, or the Assign is still in beta for some specific clients?

And a doubt 2: How can I set up the return for exceptions to actualy break the execution? I am raising the exception, but the execution keep running forever even when there is a mistake.

Lambda task "initialize_short_input" function output:
```
return { 'thread_id': thread.id, token_counters': token_counters,  'user_id': user_id,  'usage': self.usage, 'assistant_id': self.assistant_id, 'x_user_email': self.event.get('xUserEmail'), 'x_user_token': self.event.get('xUserToken'), 'params': self.params }
```
Step Machine relevant steps:

```
"ExecuteShortInputWorkflow": {
      "Comment": "Retrieve or create a thread_id, initialize the TokenCounters, adds the instructions and send the initial message.",
      "Type": "Task",
      "Resource": "lambda_arn",
      "Parameters": {
        "params.$": "$.params",
        "xUserEmail.$": "$.x_user_email",
        "xUserToken.$": "$.x_user_token",
        "mode": "initialize_short_input"
      },
      "Assign": {
        "threadId": "{% $states.result.Payload.thread_id %}",
        "tokenCounters": "{% $states.result.Payload.token_counters %}",
        "userId": "{% $states.result.Payload.user_id %}",
        "usage": "{% $states.result.Payload.usage %}",
        "assistantId": "{% $states.result.Payload.assistand_id %}",
        "xUserEmail": "{% $states.result.Payload.x_user_email %}",
        "xUserToken": "{% $states.result.Payload.x_user_id %}",
        "params": "{% $states.result.Payload.params %}",
        "resultCall": "{ % $states.result %}"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "OpenAIError"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 3
        },
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 3
        }
      ],
      "Catch": [
        {
          "ErrorEquals": [
            "OpenAIError"
          ],
          "Next": "HandleOpenAIError"
        },
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "HandleUnexpectedError"
        }
      ],
      "Next": "PassTest"
    },
    "PassTest": {
      "Type": "Pass",
      "Next": "StartAssistantRun"
    },
    "StartAssistantRun": {
      "Comment": "Create a new run with the given thread_id, and return the run id",
      "Type": "Task",
      "Resource": "lambda_arn",
      "Parameters": {
        "params": "$params",
        "token_counters": "$token_counters",
        "thread_id": "$threadId",
        "x_user_email": "$xUserEmail",
        "x_user_token": "$xUserToken",
        "result_call": "$resultCall",
        "mode": "start_run"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "OpenAIError"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 3
        },
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 3
        }
      ],
      "Catch": [
        {
          "ErrorEquals": [
            "OpenAIError"
          ],
          "Next": "HandleOpenAIError"
        },
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "HandleUnexpectedError"
        }
      ],
      "Assign": {
        "runId": "{% $states.result.Payload.run_id %}"
      },
      "Next": "WaitRunCompletion"
    }
```
Thanks in advance,
Daniel.

BrianDAnderson_AWS · Accepted Answer

Fixing Variable Assignment in Step Functions for Your OpenAI Assistant Workflow

I see the issue with your Step Functions workflow where variables aren't being properly assigned between the ExecuteShortInput
and StartExecution
states. Let me help you fix this and also address your question about exception handling.
The Variable Assignment Problem

The issue you're experiencing is that your variables are being passed as literal strings (like '$params'
) instead of their actual values. This happens because of how you're referencing the variables in your workflow definition.
Current Issue

Your current input to StartExecution
looks like:

{
  'params': '$params', 
  'token_counters': '$token_counters', 
  'thread_id': '$threadId', 
  'x_user_email': '$xUserEmail', 
  'x_user_token': '$xUserToken', 
  'result_call': '$resultCall', 
  'mode': 'start_run'
}

Solution: Fix the Variable Assignment

Here's how to properly set up the variable assignment between your states:

In your ExecuteShortInput state:

"ExecuteShortInput": {
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "YourLambdaFunction",
    "Payload.$": "$"
  },
  "ResultPath": "$.lambdaResult",
  "Next": "AssignVariables"
}

Add an explicit AssignVariables state:

"AssignVariables": {
  "Type": "Pass",
  "Parameters": {
    "params.$": "$.lambdaResult.Payload.params",
    "token_counters.$": "$.lambdaResult.Payload.token_counters",
    "thread_id.$": "$.lambdaResult.Payload.thread_id",
    "x_user_email.$": "$.lambdaResult.Payload.x_user_email",
    "x_user_token.$": "$.lambdaResult.Payload.x_user_token",
    "result_call.$": "$.lambdaResult.Payload.result_call",
    "mode": "start_run"
  },
  "Next": "StartExecution"
}

Then in your StartExecution state:

"StartExecution": {
  "Type": "Task",
  "Resource": "arn:aws:states:::states:startExecution",
  "Parameters": {
    "StateMachineArn": "arn:aws:states:region:account:stateMachine:YourStateMachine",
    "Input.$": "$"
  },
  "Next": "NextState"
}

Alternative: Using the Assign Field (AWS Step Functions SDK 2.0)

If you're using the newer AWS Step Functions SDK 2.0, you can use the Assign
field directly:

"ExecuteShortInput": {
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "YourLambdaFunction",
    "Payload.$": "$"
  },
  "ResultPath": "$.lambdaResult",
  "Assign": [
    {
      "Variable": "$.params",
      "Value.$": "$.lambdaResult.Payload.params"
    },
    {
      "Variable": "$.token_counters",
      "Value.$": "$.lambdaResult.Payload.token_counters"
    },
    {
      "Variable": "$.thread_id",
      "Value.$": "$.lambdaResult.Payload.thread_id"
    }
    // Add other variables similarly
  ],
  "Next": "StartExecution"
}

Setting Up Exception Handling to Break Execution

To handle exceptions and break the execution flow when needed:

Add a Catch block to your ExecuteShortInput state:

"ExecuteShortInput": {
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "YourLambdaFunction",
    "Payload.$": "$"
  },
  "ResultPath": "$.lambdaResult",
  "Catch": [
    {
      "ErrorEquals": ["States.ALL"],
      "Next": "HandleError"
    }
  ],
  "Next": "AssignVariables"
}

Create an error handling state:

"HandleError": {
  "Type": "Pass",
  "Parameters": {
    "error.$": "$.error",
    "cause.$": "$.cause",
    "status": "failed"
  },
  "End": true
}

For Lambda-specific errors, you can be more specific:

"Catch": [
  {
    "ErrorEquals": ["Lambda.ServiceException", "Lambda.AWSLambdaException", "Lambda.SdkClientException"],
    "Next": "HandleLambdaError"
  },
  {
    "ErrorEquals": ["States.TaskFailed"],
    "Next": "HandleTaskError"
  },
  {
    "ErrorEquals": ["States.ALL"],
    "Next": "HandleGenericError"
  }
]

Debugging Tips

If you're still having issues after implementing these changes:

Add a Pass state with ResultPath: "$" before your problematic state to see exactly what data is available:

"DebugState": {
  "Type": "Pass",
  "ResultPath": "$",
  "Next": "StartExecution"
}

Check your Lambda function's return structure to ensure it's returning the expected data format.

Use CloudWatch Logs to see the exact input and output of each state.

Complete Example

Here's a complete example that puts everything together:

{
  "Comment": "OpenAI Assistant Workflow",
  "StartAt": "ExecuteShortInput",
  "States": {
    "ExecuteShortInput": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "Parameters": {
        "FunctionName": "YourLambdaFunction",
        "Payload.$": "$"
      },
      "ResultPath": "$.lambdaResult",
      "Catch": [
        {
          "ErrorEquals": ["States.ALL"],
          "Next": "HandleError"
        }
      ],
      "Next": "AssignVariables"
    },
    "AssignVariables": {
      "Type": "Pass",
      "Parameters": {
        "params.$": "$.lambdaResult.Payload.params",
        "token_counters.$": "$.lambdaResult.Payload.token_counters",
        "thread_id.$": "$.lambdaResult.Payload.thread_id",
        "x_user_email.$": "$.lambdaResult.Payload.x_user_email",
        "x_user_token.$": "$.lambdaResult.Payload.x_user_token",
        "result_call.$": "$.lambdaResult.Payload.result_call",
        "mode": "start_run"
      },
      "Next": "StartExecution"
    },
    "StartExecution": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:region:account:stateMachine:YourStateMachine",
        "Input.$": "$"
      },
      "Next": "FinalState"
    },
    "HandleError": {
      "Type": "Pass",
      "Parameters": {
        "error.$": "$.error",
        "cause.$": "$.cause",
        "status": "failed"
      },
      "End": true
    },
    "FinalState": {
      "Type": "Pass",
      "End": true
    }
  }
}

This should resolve your variable assignment issues and provide proper exception handling to break execution when needed.

Sources:
https://docs.aws.amazon.com/step-functions/latest/dg/workflow-variables.html

Daniel Domingueti · Answer

Hi Brian,

Thanks for your feedback! It worked with ResultPath instead of using the Assign command directly on the Tasks. I thought that editing on the console would enable me the AWS Step Functions SDK 2.0, but it looks like I was wrong.

Thanks for your tips on debugging and exception handling as well! I could finally run the entire step machine and retrieve an answer from OpenAI.

See ya.

Step Function - Assign command is not assigning variables

Add your answer

Relevant content