From 6015f23e79ea049f2e128ff2052d87c190714e8d Mon Sep 17 00:00:00 2001
From: Ponder <ruan.lj@foxmail.com>
Date: Tue, 10 Feb 2026 17:55:24 +0800
Subject: [PATCH] feat: enhancement celery configuration (#32145)

---
 api/configs/middleware/__init__.py | 9 +++++++++
 api/extensions/ext_celery.py       | 6 ++++++
 docker/.env.example                | 2 ++
 docker/docker-compose.yaml         | 1 +
 4 files changed, 18 insertions(+)

diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py
index a15e42babf..0532a42371 100644
--- a/api/configs/middleware/__init__.py
+++ b/api/configs/middleware/__init__.py
@@ -259,11 +259,20 @@ class CeleryConfig(DatabaseConfig):
         description="Password of the Redis Sentinel master.",
         default=None,
     )
+
     CELERY_SENTINEL_SOCKET_TIMEOUT: PositiveFloat | None = Field(
         description="Timeout for Redis Sentinel socket operations in seconds.",
         default=0.1,
     )
 
+    CELERY_TASK_ANNOTATIONS: dict[str, Any] | None = Field(
+        description=(
+            "Annotations for Celery tasks as a JSON mapping of task name -> options "
+            "(for example, rate limits or other task-specific settings)."
+        ),
+        default=None,
+    )
+
     @computed_field
     def CELERY_RESULT_BACKEND(self) -> str | None:
         if self.CELERY_BACKEND in ("database", "rabbitmq"):
diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py
index 9917b4c88a..7b6a73af52 100644
--- a/api/extensions/ext_celery.py
+++ b/api/extensions/ext_celery.py
@@ -80,8 +80,14 @@ def init_app(app: DifyApp) -> Celery:
         worker_hijack_root_logger=False,
         timezone=pytz.timezone(dify_config.LOG_TZ or "UTC"),
         task_ignore_result=True,
+        task_annotations=dify_config.CELERY_TASK_ANNOTATIONS,
     )
 
+    if dify_config.CELERY_BACKEND == "redis":
+        celery_app.conf.update(
+            result_backend_transport_options=broker_transport_options,
+        )
+
     # Apply SSL configuration if enabled
     ssl_options = _get_celery_ssl_options()
     if ssl_options:
diff --git a/docker/.env.example b/docker/.env.example
index 8edd0f203e..c8db23b9ed 100644
--- a/docker/.env.example
+++ b/docker/.env.example
@@ -387,6 +387,8 @@ CELERY_USE_SENTINEL=false
 CELERY_SENTINEL_MASTER_NAME=
 CELERY_SENTINEL_PASSWORD=
 CELERY_SENTINEL_SOCKET_TIMEOUT=0.1
+# e.g. {"tasks.add": {"rate_limit": "10/s"}}
+CELERY_TASK_ANNOTATIONS=null
 
 # ------------------------------
 # CORS Configuration
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 5c09d203cd..afd64963c4 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -106,6 +106,7 @@ x-shared-env: &shared-api-worker-env
   CELERY_SENTINEL_MASTER_NAME: ${CELERY_SENTINEL_MASTER_NAME:-}
   CELERY_SENTINEL_PASSWORD: ${CELERY_SENTINEL_PASSWORD:-}
   CELERY_SENTINEL_SOCKET_TIMEOUT: ${CELERY_SENTINEL_SOCKET_TIMEOUT:-0.1}
+  CELERY_TASK_ANNOTATIONS: ${CELERY_TASK_ANNOTATIONS:-null}
   WEB_API_CORS_ALLOW_ORIGINS: ${WEB_API_CORS_ALLOW_ORIGINS:-*}
   CONSOLE_CORS_ALLOW_ORIGINS: ${CONSOLE_CORS_ALLOW_ORIGINS:-*}
   COOKIE_DOMAIN: ${COOKIE_DOMAIN:-}