✨(tasks) switch from celery to dramatiq

Makefile

@@ -133,6 +133,19 @@ logs: ## display all services logs (follow mode)
 	@$(COMPOSE) logs -f
 .PHONY: logs
 
+logs-worker: ## display worker logs (follow mode)
+	@$(COMPOSE) logs worker-dev -f
+.PHONY: logs-worker
+
+logs-back: ## display backend logs (follow mode)
+	@$(COMPOSE) logs backend-dev -f
+.PHONY: logs-back
+
+logs-front: ## display frontend logs (follow mode)
+	@$(COMPOSE) logs frontend-dev -f
+.PHONY: logs-front
+
+
 start: ## start all development services
 	@$(COMPOSE) up --force-recreate --build -d frontend-dev backend-dev worker-dev mta-in --wait
 .PHONY: start

compose.yaml

@@ -154,7 +154,7 @@ services:
       args:
         DOCKER_USER: ${DOCKER_USER:-1000}
     user: ${DOCKER_USER:-1000}
-    command: ["python", "worker.py", "--loglevel=DEBUG"]
+    command: ["python", "worker.py", "-v", "2"]
     environment:
       - DJANGO_CONFIGURATION=Development
     env_file:

src/backend/core/admin.py

@@ -17,7 +17,6 @@
 from sentry_sdk import capture_exception
 
 from core.api.utils import get_file_key
-from core.api.viewsets.task import register_task_owner
 from core.mda.outbound_tasks import retry_messages_task
 from core.services.dns.provisioning import provision_domain_dns
 from core.services.exporter.tasks import export_mailbox_task
@@ -327,7 +326,7 @@ def export_messages_view(self, request, object_id):
         # Start the export task
         try:
             task = export_mailbox_task.delay(str(mailbox_obj.id), str(request.user.id))
-            register_task_owner(task.id, request.user.id)
+            task.track_owner(request.user.id)
         except Exception:  # pylint: disable=broad-exception-caught
             logging.exception(
                 "Failed to queue export task for mailbox %s", mailbox_obj.id

src/backend/core/api/viewsets/send.py

@@ -14,7 +14,6 @@
 from rest_framework.views import APIView
 
 from core import models
-from core.api.viewsets.task import register_task_owner
 from core.mda.outbound import prepare_outbound_message
 from core.mda.outbound_tasks import send_message_task
 
@@ -122,7 +121,7 @@ def post(self, request):
 
         # Launch async task for sending the message
         task = send_message_task.delay(str(message.id), must_archive=must_archive)
-        register_task_owner(task.id, request.user.id)
+        task.track_owner(request.user.id)
 
         # --- Finalize ---
         # Message state should be updated by prepare_outbound_message/send_message

src/backend/core/api/viewsets/task.py

@@ -1,11 +1,10 @@
-"""API ViewSet for Celery task status."""
+"""API ViewSet for asynchronous task statuses."""
 
 import logging
 
-from django.core.cache import cache
+import dramatiq
+from dramatiq.results import ResultFailure, ResultMissing, Results
 
-from celery import states as celery_states
-from celery.result import AsyncResult
 from drf_spectacular.utils import (
     OpenApiExample,
     extend_schema,
@@ -17,16 +16,12 @@
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
-from messages.celery_app import app as celery_app
+from core.utils import get_task_progress, get_task_tracking
 
 logger = logging.getLogger(__name__)
 
-TASK_OWNER_CACHE_TTL = 86400  # 24 hours
 
-
-def register_task_owner(task_id, user_id):
-    """Register the owner of a task for permission checks."""
-    cache.set(f"task_owner:{task_id}", str(user_id), timeout=TASK_OWNER_CACHE_TTL)
+TASK_STATES = ["PENDING", "SUCCESS", "FAILURE", "PROGRESS"]
 
 
 @extend_schema(
@@ -44,11 +39,13 @@
         200: inline_serializer(
             name="TaskStatusResponse",
             fields={
-                "status": drf_serializers.ChoiceField(
-                    choices=sorted({*celery_states.ALL_STATES, "PROGRESS"})
-                ),
+                "status": drf_serializers.ChoiceField(choices=sorted(TASK_STATES)),
                 "result": drf_serializers.JSONField(allow_null=True),
                 "error": drf_serializers.CharField(allow_null=True),
+                # Present when status == "PROGRESS"
+                "progress": drf_serializers.IntegerField(required=False),
+                "message": drf_serializers.CharField(required=False, allow_blank=True),
+                "timestamp": drf_serializers.FloatField(required=False),
             },
         )
     },
@@ -69,43 +66,69 @@
     ],
 )
 class TaskDetailView(APIView):
-    """View to retrieve the status of a Celery task."""
+    """View to retrieve the status of a task."""
 
     permission_classes = [permissions.IsAuthenticated]
 
     def get(self, request, task_id):
-        """Get the status of a Celery task."""
-        owner_id = cache.get(f"task_owner:{task_id}")
-        if owner_id is None:
+        """Get the status of a task."""
+        tracking = get_task_tracking(task_id)
+        if tracking is None:
             raise PermissionDenied("Task not found or access expired.")
-        if str(request.user.id) != owner_id:
+        if str(request.user.id) != tracking["owner"]:
             raise PermissionDenied("You do not have access to this task.")
 
-        task_result = AsyncResult(task_id, app=celery_app)
-
-        # By default unknown tasks will be in PENDING. There is no reliable
-        # way to check if a task exists or not with Celery.
-        # https://github.com/celery/celery/issues/3596#issuecomment-262102185
-
-        # Prepare the response data
-        result_data = {
-            "status": task_result.status,
-            "result": None,
-            "error": None,
-        }
-
-        # If the result is a dict with status/result/error, unpack and propagate status
-        if isinstance(task_result.result, dict) and set(task_result.result.keys()) >= {
-            "status",
-            "result",
-            "error",
-        }:
-            result_data["status"] = task_result.result["status"]
-            result_data["result"] = task_result.result["result"]
-            result_data["error"] = task_result.result["error"]
-        else:
-            result_data["result"] = task_result.result
-        if task_result.state == "PROGRESS" and task_result.info:
-            result_data.update(task_result.info)
-
-        return Response(result_data)
+        # Try to fetch the result from dramatiq's native result backend
+        message = dramatiq.Message(
+            queue_name=tracking["queue_name"],
+            actor_name=tracking["actor_name"],
+            args=(),
+            kwargs={},
+            options={},
+            message_id=task_id,
+        )
+        try:
+            result_data = message.get_result(block=False)
+        except ResultMissing:
+            result_data = None
+        except ResultFailure as exc:
+            return Response({
+                "status": "FAILURE",
+                "result": None,
+                "error": str(exc),
+            })
+
+        if result_data is not None:
+            response = {"status": "SUCCESS", "result": result_data, "error": None}
+            # If the result follows the {status, result, error} convention, unpack it
+            if (
+                isinstance(result_data, dict)
+                and {"status", "result", "error"} <= result_data.keys()
+            ):
+                response["status"] = result_data["status"]
+                response["result"] = result_data["result"]
+                response["error"] = result_data["error"]
+            return Response(response)
+
+        # Check if we have progress data for this task
+        progress_data = get_task_progress(task_id)
+        if progress_data:
+            return Response(
+                {
+                    "status": "PROGRESS",
+                    "result": None,
+                    "error": None,
+                    "progress": progress_data.get("progress"),
+                    "message": progress_data.get("metadata", {}).get("message"),
+                    "timestamp": progress_data.get("timestamp"),
+                }
+            )
+
+        # Default to pending when no result and no progress
+        return Response(
+            {
+                "status": "PENDING",
+                "result": None,
+                "error": None,
+            }
+        )

src/backend/core/management/commands/run_task.py

@@ -76,7 +76,7 @@ def handle(self, *args, **options):
 
         try:
             # Execute task synchronously
-            result = task_func.apply(args=task_args, kwargs=kwargs)
+            result = task_func(*task_args, **kwargs)
 
             # Output result
             if options["json"]:

src/backend/core/mda/inbound_tasks.py

@@ -8,16 +8,15 @@
 from django.core.cache import cache
 from django.utils import timezone
 
-import requests
-from celery.utils.log import get_task_logger
+import logging
 
 from core import models
 from core.mda.inbound_create import _create_message_from_inbound
 from core.mda.rfc5322 import parse_email_message
+import requests
+from core.utils import cron_task, register_task
 
-from messages.celery_app import app as celery_app
-
-logger = get_task_logger(__name__)
+logger = logging.getLogger(__name__)
 
 
 def _check_spam_with_hardcoded_rules(
@@ -167,8 +166,8 @@ def _check_spam_with_rspamd(
         return False, str(e)
 
 
-@celery_app.task(bind=True)
-def process_inbound_message_task(self, inbound_message_id: str):
+@register_task(queue="inbound")
+def process_inbound_message_task(inbound_message_id: str):
     """Process an inbound message from the queue: check spam and create message.
 
     Args:
@@ -285,12 +284,13 @@ def process_inbound_message_task(self, inbound_message_id: str):
         cache.delete(lock_key)
 
 
-@celery_app.task(bind=True)
-def process_inbound_messages_queue_task(self, batch_size: int = 10):
+@cron_task(interval=300)
+@register_task(queue="inbound")
+def process_inbound_messages_queue_task(batch_size: int = 10):
     """Retry processing of inbound messages that are older than 5 minutes.
 
     This task only handles retries for messages that may have failed or gotten stuck.
-    Regular messages are processed immediately when created via process_inbound_message_task.delay().
+    Regular messages are processed immediately when created via process_inbound_message_task.send().
 
     Args:
         batch_size: Number of messages to process in this batch

src/backend/core/mda/outbound.py

@@ -364,7 +364,7 @@ def _validate_attachments_size(total_size: int) -> None:
 def send_message(message: models.Message, force_mta_out: bool = False):
     """Send an existing Message, internally or externally.
 
-    This part is called asynchronously from the celery worker.
+    This part is called asynchronously from the background worker.
     """
 
     # Refuse to send messages that are draft or not senders