diff --git a/CHANGELOG.md b/CHANGELOG.md
index dac8528..320e21a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,14 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
+## [4.5.1] - 2024-05-06
+### Added
+- Adding tags to the Datadog agent
+- Refactor Datadog agent implementation
+
+## [4.5.0] - 2024-04-25
+### Added
+- Added optional `extended_server_config`. For instance to configure Waggle Dance Rate limiting configuration.
 
 ## [4.5.0] - 2024-04-25
 ### Added
diff --git a/ecs.tf b/ecs.tf
index 1d4ab6d..cf7164d 100644
--- a/ecs.tf
+++ b/ecs.tf
@@ -18,7 +18,9 @@ resource "aws_ecs_service" "waggledance_service" {
   task_definition = aws_ecs_task_definition.waggledance[0].arn
   desired_count   = var.wd_ecs_task_count
 
-  propagate_tags  = "SERVICE"
+  platform_version = "1.4.0"
+
+  # propagate_tags  = "SERVICE"
   tags            = var.tags
 
   network_configuration {
@@ -52,12 +54,7 @@ resource "aws_ecs_task_definition" "waggledance" {
   memory                   = var.memory
   cpu                      = var.cpu
   requires_compatibilities = ["EC2", "FARGATE"]
-  container_definitions = <<DEFINITION
-  [
-      ${data.template_file.waggledance.rendered}
-      ${var.include_datadog_agent ? data.template_file.datadog-agent.rendered : ""}
-  ]
-  DEFINITION
+  container_definitions    = data.template_file.waggledance.rendered
   tags                     = var.tags
 }
 
diff --git a/templates.tf b/templates.tf
index 7cae0d3..1832a64 100644
--- a/templates.tf
+++ b/templates.tf
@@ -7,6 +7,7 @@
 locals {
   default_exposed_endpoints = "health,info,metrics"
   exposed_endpoints         = var.prometheus_enabled ? join(",", [local.default_exposed_endpoints, "prometheus"]) : local.default_exposed_endpoints
+  datadog_tags              = join(" ", formatlist("%s:%s", keys(var.tags), values(var.tags)))
 }
 
 data "template_file" "endpoints_server_yaml" {
@@ -186,25 +187,16 @@ data "template_file" "waggledance" {
     hive_site_xml       = length(var.alluxio_endpoints) == 0 ? "" : base64encode(data.template_file.hive_site_xml.rendered)
     bastion_ssh_key_arn = var.bastion_ssh_key_secret_name == "" ? "" : join("", data.aws_secretsmanager_secret.bastion_ssh_key.*.arn)
     docker_auth         = var.docker_registry_auth_secret_name == "" ? "" : format("\"repositoryCredentials\" :{\n \"credentialsParameter\":\"%s\"\n},", join("\",\"", concat(data.aws_secretsmanager_secret.docker_registry.*.arn)))
-    tcp_keepalive_time  = var.tcp_keepalive_time
-    tcp_keepalive_intvl = var.tcp_keepalive_intvl
-    tcp_keepalive_probes = var.tcp_keepalive_probes
-    wd_instance_type = var.wd_instance_type
-    metrics_port = var.metrics_port
-    datadog_agent_version = var.datadog_agent_version
-    include_datadog_agent = var.include_datadog_agent
-  }
-}
+    wd_instance_type    = var.wd_instance_type
 
-data "template_file" "datadog-agent" {
-  template = file("${path.module}/templates/datadog-agent.json")
+    tcp_keepalive_time   = var.tcp_keepalive_time
+    tcp_keepalive_intvl  = var.tcp_keepalive_intvl
+    tcp_keepalive_probes = var.tcp_keepalive_probes
 
-  vars = {
-    region                = var.aws_region
-    loggroup              = var.wd_instance_type == "ecs" ? join("", aws_cloudwatch_log_group.waggledance_ecs.*.name) : ""
-    datadog_secret_key    = length(var.datadog_key_secret_name) > 0 ? chomp(data.external.datadog_key[0].result["api_key"]) : ""
-    wd_instance_type      = var.wd_instance_type
     metrics_port          = var.metrics_port
     datadog_agent_version = var.datadog_agent_version
+    include_datadog_agent = var.include_datadog_agent
+    datadog_tags          = local.datadog_tags
+    datadog_secret_key    = length(var.datadog_key_secret_name) > 0 ? chomp(data.external.datadog_key[0].result["api_key"]) : ""
   }
-}
+}
\ No newline at end of file
diff --git a/templates/datadog-agent.json b/templates/datadog-agent.json
deleted file mode 100644
index d2f49fb..0000000
--- a/templates/datadog-agent.json
+++ /dev/null
@@ -1,31 +0,0 @@
-,
-{
-  "name": "datadog-agent",
-  "image": "public.ecr.aws/datadog/agent:${datadog_agent_version}",
-  "essential": true,
-  "logConfiguration": {
-    "logDriver": "awslogs",
-    "options": {
-      "awslogs-group": "${loggroup}",
-      "awslogs-region": "${region}",
-      "awslogs-stream-prefix": "/"
-    }
-  },
-  "environment": [
-    {
-      "name": "DD_API_KEY",
-      "value": "${datadog_secret_key}"
-    },
-    {
-      "name": "ECS_FARGATE",
-      "value": "true"
-    }
-  ],
-  "healthCheck": {
-    "command": ["CMD-SHELL", "curl -f http://localhost:18000/actuator/health || exit 1"],
-    "interval": 5,
-    "retries": 3,
-    "startPeriod": 60,
-    "timeout": 5
-  }
-}
diff --git a/templates/waggledance.json b/templates/waggledance.json
index 84c45d7..5c2ab38 100644
--- a/templates/waggledance.json
+++ b/templates/waggledance.json
@@ -1,92 +1,123 @@
-{
-	"name": "waggledance",
-	"image": "${docker_image}:${docker_version}",
-	${docker_auth}
-	"essential": true,
-	"logConfiguration": {
-	  "logDriver": "awslogs",
-	  "options": {
-		"awslogs-group": "${loggroup}",
-		"awslogs-region": "${region}",
-		"awslogs-stream-prefix": "/"
-	  }
-	},
-	"dockerLabels": {
-		"com.datadoghq.ad.instances": "[{\"openmetrics_endpoint\": \"http://%%host%%:18000/actuator/prometheus\", \"namespace\": \"waggledance\", \"metrics\": [\"metastore_status\",\"counter_*\",\"jvm_*\",\"system_*\",\"timer_*\",\"http_*\",\"process_*\"], \"tags\": { \"region\":\"${region}\", \"instance_type\":\"${wd_instance_type}\"} }]",
-		"com.datadoghq.ad.check_names": "[\"openmetrics\"]",
-		"com.datadoghq.ad.init_configs": "[{}]"
-	},
-	"portMappings": [
-	  {
-		"containerPort": 48869,
-		"hostPort": 48869
-	  }
-	],
-	"environment":[
-	  {
-		"name": "HEAPSIZE",
-		"value": "${heapsize}"
-	  },
-	  {
-		"name": "LOGLEVEL",
-		"value": "${loglevel}"
-	  },
-	  {
-		"name": "INVOCATIONLOGLEVEL",
-		"value": "${invocationloglevel}"
-	  },
-	  {
-		"name": "SERVER_YAML",
-		"value": "${server_yaml}"
-	  },
-	  {
-		"name": "FEDERATION_YAML",
-		"value": "${federation_yaml}"
-	  },
-	  {
-		"name": "HIVE_SITE_XML",
-		"value": "${hive_site_xml}"
-	  },
-	  {
-		"name": "BASTION_SSH_KEY_ARN",
-		"value": "${bastion_ssh_key_arn}"
-	  },
-	  {
-		"name": "LOG4J_FORMAT_MSG_NO_LOOKUPS",
-		"value": "true"
-	  }
-	],
-	"healthCheck": {
-	  "command": ["CMD-SHELL", "curl -f http://localhost:18000/actuator/health || exit 1"],
-	  "interval": 5,
-	  "retries": 3,
-	  "startPeriod": 60,
-	  "timeout": 5
-	},
-	"ulimits": [
-	  {
-		"name": "nofile",
-		"softLimit": 65536,
-		"hardLimit": 65536
-	  },
-	  {
-		"name": "nproc",
-		"softLimit": 65536,
-		"hardLimit": 65536
-	  }
-	],
-	"systemControls": [
-          {
-                "namespace": "net.ipv4.tcp_keepalive_time",
-                "value": "${tcp_keepalive_time}"
-          },
-          {
-                "namespace": "net.ipv4.tcp_keepalive_intvl",
-                "value": "${tcp_keepalive_intvl}"
-          },
-          {
-                "namespace": "net.ipv4.tcp_keepalive_probes",
-                "value": "${tcp_keepalive_probes}"
-          }
-        ]
-}
+[
+	{
+		"name": "waggledance",
+		"image": "${docker_image}:${docker_version}",
+		${docker_auth}
+		"essential": true,
+		"logConfiguration": {
+				"logDriver": "awslogs",
+				"options": {
+				"awslogs-group": "${loggroup}",
+				"awslogs-region": "${region}",
+				"awslogs-stream-prefix": "/"
+			}
+		},
+		"dockerLabels": {
+			"com.datadoghq.ad.instances": "[{\"openmetrics_endpoint\": \"http://%%host%%:18000/actuator/prometheus\", \"namespace\": \"waggledance\", \"metrics\": [\"metastore_status\",\"counter_*\",\"jvm_*\",\"system_*\",\"timer_*\",\"http_*\",\"process_*\"], \"tags\": { \"region\":\"${region}\", \"instance_type\":\"${wd_instance_type}\"} }]",
+			"com.datadoghq.ad.check_names": "[\"openmetrics\"]",
+			"com.datadoghq.ad.init_configs": "[{}]"
+		},
+		"portMappings": [
+			{
+				"containerPort": 48869,
+				"hostPort": 48869
+			}
+		],
+		"environment": [
+			{
+				"name": "HEAPSIZE",
+				"value": "${heapsize}"
+			},
+			{
+				"name": "LOGLEVEL",
+				"value": "${loglevel}"
+			},
+			{
+				"name": "INVOCATIONLOGLEVEL",
+				"value": "${invocationloglevel}"
+			},
+			{
+				"name": "SERVER_YAML",
+				"value": "${server_yaml}"
+			},
+			{
+				"name": "FEDERATION_YAML",
+				"value": "${federation_yaml}"
+			},
+			{
+				"name": "HIVE_SITE_XML",
+				"value": "${hive_site_xml}"
+			},
+			{
+				"name": "BASTION_SSH_KEY_ARN",
+				"value": "${bastion_ssh_key_arn}"
+			},
+			{
+				"name": "LOG4J_FORMAT_MSG_NO_LOOKUPS",
+				"value": "true"
+			}
+		],
+		"healthCheck": {
+			"command": ["CMD-SHELL", "curl -f http://localhost:18000/actuator/health || exit 1"],
+			"interval": 5,
+			"retries": 3,
+			"startPeriod": 60,
+			"timeout": 5
+		},
+		"ulimits": [
+			{
+				"name": "nofile",
+				"softLimit": 65536,
+				"hardLimit": 65536
+			},
+			{
+				"name": "nproc",
+				"softLimit": 65536,
+				"hardLimit": 65536
+			}
+		],
+		"systemControls": [
+			{
+				"namespace": "net.ipv4.tcp_keepalive_time",
+				"value": "${tcp_keepalive_time}"
+			},
+			{
+				"namespace": "net.ipv4.tcp_keepalive_intvl",
+				"value": "${tcp_keepalive_intvl}"
+			},
+			{
+				"namespace": "net.ipv4.tcp_keepalive_probes",
+				"value": "${tcp_keepalive_probes}"
+			}
+		]
+	}
+	%{ if include_datadog_agent }
+  ,{
+   "name": "datadog-agent",
+   "image": "public.ecr.aws/datadog/agent:${datadog_agent_version}",
+   "essential": true,
+   "logConfiguration": {
+     "logDriver": "awslogs",
+     "options": {
+       "awslogs-group": "${loggroup}",
+       "awslogs-region": "${region}",
+       "awslogs-stream-prefix": "/"
+     }
+   },
+   "environment": [
+     {
+       "name": "DD_API_KEY",
+       "value": "${datadog_secret_key}"
+     },
+     {
+       "name": "ECS_FARGATE",
+       "value": "true"
+     },
+     {
+      "name": "DD_TAGS",
+      "value": "${datadog_tags}"
+     }
+   ]
+ }
+%{ endif }
+]
\ No newline at end of file