feat(scheduler): Implement conditional cloud status updates

Adds functionality to send task status updates (success/failure) to RabbitMQ via , specifically for the cloud service platform. This includes:
- Adding a  field to .
- Passing  to .
- Modifying  to conditionally send  with  or  (along with  and ) based on the  environment variable.

Author: [email protected]

poetry.lock

@@ -46,7 +46,6 @@ dependencies = [
     "scikit-learn (>=1.7.0,<2.0.0)",  # Machine learning
     "fastmcp (>=2.10.5,<3.0.0)",
     "python-dateutil (>=2.9.0.post0,<3.0.0)",
-    "prometheus-client (>=0.23.1,<0.24.0)",
 ]
 
 [project.urls]
@@ -76,6 +75,7 @@ tree-mem = [
 mem-scheduler = [
     "redis (>=6.2.0,<7.0.0)",  # Key-value store
     "pika (>=1.3.2,<2.0.0)",  # RabbitMQ client
+    "prometheus-client (>=0.23.1,<0.24.0)", # For metrics
 ]
 
 # MemUser (MySQL support)

pyproject.toml

@@ -14,17 +14,18 @@
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
 
-from memos.log import get_logger
-
 # Imports for new implementation
-from memos.api.product_models import StatusRequest, StatusResponse, StatusResponseItem
+from memos.api.product_models import StatusResponse, StatusResponseItem
+from memos.log import get_logger
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
 
 
 logger = get_logger(__name__)
 
 
-def handle_scheduler_status(user_id: str, status_tracker: TaskStatusTracker, task_id: str | None = None) -> StatusResponse:
+def handle_scheduler_status(
+    user_id: str, status_tracker: TaskStatusTracker, task_id: str | None = None
+) -> StatusResponse:
     """
     Get scheduler running status for one or all tasks of a user.
 
@@ -41,20 +42,24 @@ def handle_scheduler_status(user_id: str, status_tracker: TaskStatusTracker, tas
     Raises:
         HTTPException: If a specific task is not found.
     """
-    request = StatusRequest(user_id=user_id, task_id=task_id) # Construct StatusRequest internally
     response_data: list[StatusResponseItem] = []
 
     try:
         if task_id:
             task_data = status_tracker.get_task_status(task_id, user_id)
             if not task_data:
-                raise HTTPException(status_code=404, detail=f"Task {task_id} not found for user {user_id}")
+                raise HTTPException(
+                    status_code=404, detail=f"Task {task_id} not found for user {user_id}"
+                )
             response_data.append(StatusResponseItem(task_id=task_id, status=task_data["status"]))
         else:
             all_tasks = status_tracker.get_all_tasks_for_user(user_id)
             # The plan returns an empty list, which is good.
             # No need to check "if not all_tasks" explicitly before the list comprehension
-            response_data = [StatusResponseItem(task_id=tid, status=t_data["status"]) for tid, t_data in all_tasks.items()]
+            response_data = [
+                StatusResponseItem(task_id=tid, status=t_data["status"])
+                for tid, t_data in all_tasks.items()
+            ]
 
         return StatusResponse(data=response_data)
     except HTTPException:
@@ -94,21 +99,19 @@ def handle_scheduler_wait(
         while time.time() - start_time < timeout_seconds:
             # Directly call the new, reliable status logic
             status_response = handle_scheduler_status(
-                user_id=user_name,
-                status_tracker=status_tracker
+                user_id=user_name, status_tracker=status_tracker
             )
 
             # System is idle if the data list is empty or no tasks are active
             is_idle = not status_response.data or all(
-                task.status in ["completed", "failed", "cancelled"]
-                for task in status_response.data
+                task.status in ["completed", "failed", "cancelled"] for task in status_response.data
             )
 
             if is_idle:
                 return {
                     "message": "idle",
                     "data": {
-                        "running_tasks": 0, # Kept for compatibility
+                        "running_tasks": 0,  # Kept for compatibility
                         "waited_seconds": round(time.time() - start_time, 3),
                         "timed_out": False,
                         "user_name": user_name,
@@ -124,7 +127,7 @@ def handle_scheduler_wait(
         return {
             "message": "timeout",
             "data": {
-                "running_tasks": len(active_tasks), # A more accurate count of active tasks
+                "running_tasks": len(active_tasks),  # A more accurate count of active tasks
                 "waited_seconds": round(time.time() - start_time, 3),
                 "timed_out": True,
                 "user_name": user_name,
@@ -134,7 +137,9 @@ def handle_scheduler_wait(
         # Re-raise HTTPException directly to preserve its status code
         raise
     except Exception as err:
-        logger.error(f"Failed while waiting for scheduler for user {user_name}: {traceback.format_exc()}")
+        logger.error(
+            f"Failed while waiting for scheduler for user {user_name}: {traceback.format_exc()}"
+        )
         raise HTTPException(status_code=500, detail="Failed while waiting for scheduler") from err
 
 
@@ -169,8 +174,12 @@ def event_generator():
                 elapsed = time.time() - start_time
                 if elapsed > timeout_seconds:
                     # Send timeout message and break
-                    final_status = handle_scheduler_status(user_id=user_name, status_tracker=status_tracker)
-                    active_tasks = [t for t in final_status.data if t.status in ["waiting", "in_progress"]]
+                    final_status = handle_scheduler_status(
+                        user_id=user_name, status_tracker=status_tracker
+                    )
+                    active_tasks = [
+                        t for t in final_status.data if t.status in ["waiting", "in_progress"]
+                    ]
                     payload = {
                         "user_name": user_name,
                         "active_tasks": len(active_tasks),
@@ -184,10 +193,11 @@ def event_generator():
 
                 # Get status
                 status_response = handle_scheduler_status(
-                    user_id=user_name,
-                    status_tracker=status_tracker
+                    user_id=user_name, status_tracker=status_tracker
                 )
-                active_tasks = [t for t in status_response.data if t.status in ["waiting", "in_progress"]]
+                active_tasks = [
+                    t for t in status_response.data if t.status in ["waiting", "in_progress"]
+                ]
                 num_active = len(active_tasks)
 
                 payload = {
@@ -200,7 +210,7 @@ def event_generator():
                 yield "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"
 
                 if num_active == 0:
-                    break # Exit loop if idle
+                    break  # Exit loop if idle
 
                 time.sleep(poll_interval)
 
@@ -215,4 +225,3 @@ def event_generator():
             yield "data: " + json.dumps(err_payload, ensure_ascii=False) + "\n\n"
 
     return StreamingResponse(event_generator(), media_type="text/event-stream")
-

src/memos/api/handlers/scheduler_handler.py

@@ -141,6 +141,7 @@ def __init__(self, config: BaseSchedulerConfig):
             enable_parallel_dispatch=self.enable_parallel_dispatch,
             status_tracker=self.status_tracker,
             metrics=self.metrics,
+            submit_web_logs=self._submit_web_logs,
         )
 
         # other attributes

src/memos/mem_scheduler/base_scheduler.py

@@ -141,6 +141,7 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
     )
     memcube_name: str | None = Field(default=None, description="Display name for memcube")
     memory_len: int | None = Field(default=None, description="Count of items involved in the event")
+    status: str | None = Field(default=None, description="Completion status of the task (e.g., 'completed', 'failed')")
 
     def debug_info(self) -> dict[str, Any]:
         """Return structured debug information for logging purposes."""

src/memos/mem_scheduler/schemas/message_schemas.py

@@ -12,7 +12,7 @@
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.general_modules.task_threads import ThreadManager
 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 ScheduleMessageItem, ScheduleLogForWebItem
 from memos.mem_scheduler.schemas.task_schemas import RunningTaskItem
 from memos.mem_scheduler.utils.misc_utils import group_messages_by_user_and_mem_cube
 from memos.mem_scheduler.utils import metrics
@@ -44,6 +44,7 @@ def __init__(
         config=None,
         status_tracker: TaskStatusTracker | None = None,
         metrics: Any | None = None,
+        submit_web_logs: Callable | None = None, # ADDED
     ):
         super().__init__()
         self.config = config
@@ -95,6 +96,7 @@ def __init__(
 
         self.metrics = metrics
         self.status_tracker = status_tracker
+        self.submit_web_logs = submit_web_logs # ADDED
 
     def on_messages_enqueued(self, msgs: list[ScheduleMessageItem]) -> None:
         if not msgs:
@@ -153,6 +155,18 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                     self.status_tracker.task_completed(task_id=task_item.item_id, user_id=task_item.user_id)
                 self.metrics.task_completed(user_id=m.user_id, task_type=m.label)
 
+                is_cloud_env = os.getenv('MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME') == 'memos-memory-change'
+                if self.submit_web_logs and is_cloud_env:
+                    status_log = ScheduleLogForWebItem(
+                        user_id=task_item.user_id,
+                        mem_cube_id=task_item.mem_cube_id,
+                        item_id=task_item.item_id,
+                        label=m.label,
+                        log_content=f"Task {task_item.item_id} completed successfully for user {task_item.user_id}.",
+                        status="completed"
+                    )
+                    self.submit_web_logs([status_log])
+
                 # acknowledge redis messages
                 if self.use_redis_queue and self.memos_message_queue is not None:
                     for msg in messages:
@@ -188,6 +202,19 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                         if len(self._completed_tasks) > self.completed_tasks_max_show_size:
                             self._completed_tasks.pop(0)
                 logger.error(f"Task failed: {task_item.get_execution_info()}, Error: {e}")
+
+                is_cloud_env = os.getenv('MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME') == 'memos-memory-change'
+                if self.submit_web_logs and is_cloud_env:
+                    status_log = ScheduleLogForWebItem(
+                        user_id=task_item.user_id,
+                        mem_cube_id=task_item.mem_cube_id,
+                        item_id=task_item.item_id,
+                        label=m.label,
+                        log_content=f"Task {task_item.item_id} failed for user {task_item.user_id} with error: {str(e)}.",
+                        status="failed",
+                        exception=str(e)
+                    )
+                    self.submit_web_logs([status_log])
                 raise
 
         return wrapped_handler

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -1,10 +1,16 @@
 # src/memos/mem_scheduler/utils/status_tracker.py
 import json
-from datetime import datetime, timezone
+
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
 
 from memos.dependency import require_python_package
 
 
+if TYPE_CHECKING:
+    import redis
+
+
 class TaskStatusTracker:
     @require_python_package(import_name="redis", install_command="pip install redis")
     def __init__(self, redis_client: "redis.Redis"):
@@ -22,42 +28,54 @@ def task_submitted(self, task_id: str, user_id: str, task_type: str, mem_cube_id
             "submitted_at": datetime.now(timezone.utc).isoformat(),
         }
         self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
 
     def task_started(self, task_id: str, user_id: str):
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
         if not existing_data_json:
-            # 容错处理：如果任务不存在，也创建一个
-            payload = {"status": "in_progress", "started_at": datetime.now(timezone.utc).isoformat()}
+            # 容错处理: 如果任务不存在, 也创建一个
+            payload = {
+                "status": "in_progress",
+                "started_at": datetime.now(timezone.utc).isoformat(),
+            }
         else:
             payload = json.loads(existing_data_json)
             payload["status"] = "in_progress"
             payload["started_at"] = datetime.now(timezone.utc).isoformat()
         self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
 
     def task_completed(self, task_id: str, user_id: str):
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
-        if not existing_data_json: return
+        if not existing_data_json:
+            return
         payload = json.loads(existing_data_json)
         payload["status"] = "completed"
         payload["completed_at"] = datetime.now(timezone.utc).isoformat()
-        # 设置该任务条目的过期时间，例如 24 小时
-        # 注意：Redis Hash 不能为单个 field 设置 TTL，这里我们可以 通过后台任务清理或在获取时判断时间戳
-        # 简单起见，我们暂时依赖一个后台清理任务
+        # 设置该任务条目的过期时间, 例如 24 小时
+        # 注意: Redis Hash 不能为单个 field 设置 TTL, 这里我们可以 通过后台任务清理或在获取时判断时间戳
+        # 简单起见, 我们暂时依赖一个后台清理任务
         self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
 
     def task_failed(self, task_id: str, user_id: str, error_message: str):
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
         if not existing_data_json:
-            payload = {"status": "failed", "error": error_message, "failed_at": datetime.now(timezone.utc).isoformat()}
+            payload = {
+                "status": "failed",
+                "error": error_message,
+                "failed_at": datetime.now(timezone.utc).isoformat(),
+            }
         else:
             payload = json.loads(existing_data_json)
             payload["status"] = "failed"
             payload["error"] = error_message
             payload["failed_at"] = datetime.now(timezone.utc).isoformat()
         self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
 
     def get_task_status(self, task_id: str, user_id: str) -> dict | None:
         key = self._get_key(user_id)