Add task completion log event for cloud tasks

Author: [email protected]

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -1,4 +1,5 @@
 import concurrent
+import os
 import threading
 import time
 
@@ -19,7 +20,7 @@
 from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_STOP_WAIT,
 )
-from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.message_schemas import ScheduleLogForWebItem, ScheduleMessageItem
 from memos.mem_scheduler.schemas.task_schemas import RunningTaskItem
 from memos.mem_scheduler.task_schedule_modules.orchestrator import SchedulerOrchestrator
 from memos.mem_scheduler.task_schedule_modules.redis_queue import SchedulerRedisQueue
@@ -200,6 +201,7 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                 if self.status_tracker:
                     for msg in messages:
                         self.status_tracker.task_completed(task_id=msg.item_id, user_id=msg.user_id)
+                    self._maybe_emit_task_completion(messages)
                 self.metrics.task_completed(user_id=m.user_id, task_type=m.label)
 
                 emit_monitor_event(
@@ -284,6 +286,56 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
 
         return wrapped_handler
 
+    def _maybe_emit_task_completion(self, messages: list[ScheduleMessageItem]) -> None:
+        """If all item_ids under a business task are completed, emit a single completion log."""
+        if not self.submit_web_logs or not self.status_tracker:
+            return
+
+        # messages in one batch can belong to different business task_ids; check each
+        task_ids = {getattr(msg, "task_id", None) for msg in messages}
+        task_ids.discard(None)
+        if not task_ids:
+            return
+
+        # Use the first message only for shared fields; mem_cube_id is same within a batch
+        first = messages[0]
+        user_id = first.user_id
+        mem_cube_id = first.mem_cube_id
+
+        try:
+            is_cloud_env = (
+                os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME") == "memos-memory-change"
+            )
+            if not is_cloud_env:
+                return
+
+            for task_id in task_ids:
+                status_data = self.status_tracker.get_task_status_by_business_id(
+                    business_task_id=task_id, user_id=user_id
+                )
+                if not status_data or status_data.get("status") != "completed":
+                    continue
+
+                event = ScheduleLogForWebItem(
+                    task_id=task_id,
+                    user_id=user_id,
+                    mem_cube_id=mem_cube_id,
+                    label="taskStatus",
+                    from_memory_type="status",
+                    to_memory_type="status",
+                    log_content=f"Task {task_id} completed",
+                    status="completed",
+                )
+                self.submit_web_logs(event)
+        except Exception:
+            logger.warning(
+                "Failed to emit task completion log. user_id=%s mem_cube_id=%s task_ids=%s",
+                user_id,
+                mem_cube_id,
+                list(task_ids),
+                exc_info=True,
+            )
+
     def get_running_tasks(
         self, filter_func: Callable[[RunningTaskItem], bool] | None = None
     ) -> dict[str, RunningTaskItem]: