Getting error while getting output from Running ECS Task using step function (using callback)

Hi there, I am working on running aws ecs task using aws step functions.

The task is being executed but i am not able to get the output of the task. I am using the callback mechanism and sending the task_token as environment variable to the ecs_task and then sending the result using the send_success api function. But i am getting error

Here is an example of my Step function

{
  "Comment": "A description of my state machine",
  "StartAt": "ECS RunTask",
  "States": {
    "ECS RunTask": {
      "Type": "Task",
      "Resource": "arn:aws:states:::ecs:runTask.waitForTaskToken",
      "Parameters": {
        "LaunchType": "FARGATE",
        "Cluster": "arn:aws:ecs:ap-south-1:id:cluster/cluster_name",
        "TaskDefinition": "arn:aws:ecs:ap-south-1:id:task-definition/task_name:revision",
        "NetworkConfiguration": {
          "AwsvpcConfiguration": {
            "Subnets": [
              "subnet-XXXXXXXX",
              "subnet-XXXXXXXX",
              "subnet-XXXXXXXX"
            ],
            "SecurityGroups": [
              "sg-XXXXXXXX"
            ],
            "AssignPublicIp": "ENABLED"
          }
        },
        "Overrides": {
          "ContainerOverrides": [
            {
              "Name": "container",
              "Environment": [
                {
                  "Name": "TASK_TOKEN",
                  "Value.$": "$$.Task.Token"
                }
              ]
            }
          ]
        }
      },
      "End": true
    }
  }
}

This is my script running in the ecs task

import os
import json
import boto3
from workflow.utils.db_utils import get_connections

def main():
    # Get the job_schedule_filter from the environment variable
    job_schedule_filter = os.getenv('JOB_SCHEDULE_FILTER')
    
    # Get the task_token from the environment variable
    task_token = os.getenv('TASK_TOKEN')

    # Ensure the environment variables are set
    if job_schedule_filter is None:
        raise ValueError("The environment variable 'JOB_SCHEDULE_FILTER' is not set")
    if task_token is None:
        raise ValueError("The environment variable 'TASK_TOKEN' is not set")

    # Log the task token for debugging
    print(f"TASK_TOKEN: {task_token}")

    # Fetch the connections with the provided job_schedule_filter
    connections = get_connections(job_schedule_filter)
    
    # Convert Row objects to dictionaries for JSON serialization
    connections_dict = [row._asdict() for row in connections]

    # Prepare the output
    output = json.dumps({"connections": connections_dict})

    # Send the task success back to Step Functions
    client = boto3.client('stepfunctions')
    
    
    response = client.send_task_success(
            taskToken=task_token,
            output=output  # Sending the JSON output
        )
    print("Task success sent to Step Functions:", response)
    

if __name__ == "__main__":
    main()

This is the output i am getting

response = client.send_task_success(
              ^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/home/asrhraf/.cache/pypoetry/virtualenvs/stackgenius-QEUWgpwQ-py3.12/lib/python3.12/site-packages/botocore/client.py", line 565, in _api_call
   return self._make_api_call(operation_name, kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/home/asrhraf/.cache/pypoetry/virtualenvs/stackgenius-QEUWgpwQ-py3.12/lib/python3.12/site-packages/botocore/client.py", line 1021, in _make_api_call
   raise error_class(parsed_response, operation_name)
botocore.errorfactory.InvalidToken: An error occurred (InvalidToken) when calling the SendTaskSuccess operation: Invalid Token: 'Invalid token'```