Fix(rabbitmq): Handle exchange and routing key based on cloud env for… (#667)

Fix(rabbitmq): Handle exchange and routing key based on cloud env for specific message types

Refactor rabbitmq_publish_message to correctly handle exchange and routing key
determination based on the presence of MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME
environment variable and the message label.

- Default exchange and routing key are used for most messages.
- 'knowledgeBaseUpdate' messages always have an empty routing key.
- If MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME is set, and the message label is
  'taskStatus' or 'knowledgeBaseUpdate', the exchange is overridden with
  the environment variable's value, and the routing key is set to empty.
  This fixes the 'taskStatus' routing issue in cloud environments.
- Logging for cloud-affected messages is now specific and separate from local
  'knowledgeBaseUpdate' logging.

Co-authored-by: [email protected] <>

Author: glin93

src/memos/mem_scheduler/webservice_modules/rabbitmq_service.py

@@ -283,18 +283,28 @@ def rabbitmq_publish_message(self, message: dict):
 
         exchange_name = self.rabbitmq_exchange_name
         routing_key = self.rabbit_queue_name
+        label = message.get("label")
 
-        if message.get("label") == "knowledgeBaseUpdate":
-            kb_specific_exchange_name = os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME")
+        # Special handling for knowledgeBaseUpdate in local environment: always empty routing key
+        if label == "knowledgeBaseUpdate":
+            routing_key = ""
 
-            if kb_specific_exchange_name:
-                exchange_name = kb_specific_exchange_name
-
-            routing_key = ""  # User specified empty routing key for KB updates
+        # Cloud environment override: applies to specific message types if MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME is set
+        env_exchange_name = os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME")
+        if env_exchange_name and label in ["taskStatus", "knowledgeBaseUpdate"]:
+            exchange_name = env_exchange_name
+            routing_key = ""  # Routing key is always empty in cloud environment for these types
 
+            # Specific diagnostic logging for messages affected by cloud environment settings
+            logger.info(
+                f"[DIAGNOSTIC] Publishing {label} message in Cloud Env. "
+                f"Exchange: {exchange_name}, Routing Key: '{routing_key}'."
+            )
+            logger.info(f"  - Message Content: {json.dumps(message, indent=2)}")
+        elif label == "knowledgeBaseUpdate":
+            # Original diagnostic logging for knowledgeBaseUpdate if NOT in cloud env
             logger.info(
-                f"[DIAGNOSTIC] Publishing KB Update message. "
-                f"ENV_EXCHANGE_NAME_USED: {kb_specific_exchange_name is not None}. "
+                f"[DIAGNOSTIC] Publishing knowledgeBaseUpdate message (Local Env). "
                 f"Current configured Exchange: {exchange_name}, Routing Key: '{routing_key}'."
             )
             logger.info(f"  - Message Content: {json.dumps(message, indent=2)}")