feat: a range of new feats to make a better redis scheduler

examples/mem_scheduler/task_stop_rerun.py

@@ -28,6 +28,7 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
         try:
             print(f"writing {file_path}...")
             file_path.write_text(f"Task {task_id} processed.\n")
+            sleep(5)
         except Exception as e:
             print(f"Failed to write {file_path}: {e}")
 
@@ -57,6 +58,8 @@ def submit_tasks():
 TEST_HANDLER_LABEL = "test_handler"
 mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})
 
+# 10s to restart
+mem_scheduler.orchestrator.tasks_min_idle_ms[TEST_HANDLER_LABEL] = 10_000
 
 tmp_dir = Path("./tmp")
 tmp_dir.mkdir(exist_ok=True)
@@ -69,10 +72,15 @@ def submit_tasks():
     submit_tasks()
 
 # 6. Wait until tmp has 100 files or timeout
-poll_interval = 0.01
+poll_interval = 1
 expected = 100
 tmp_dir = Path("tmp")
-while mem_scheduler.get_tasks_status()["remaining"] != 0:
+tasks_status = mem_scheduler.get_tasks_status()
+mem_scheduler.print_tasks_status(tasks_status=tasks_status)
+while (
+    mem_scheduler.get_tasks_status()["remaining"] != 0
+    or mem_scheduler.get_tasks_status()["running"] != 0
+):
     count = len(list(tmp_dir.glob("*.txt"))) if tmp_dir.exists() else 0
     tasks_status = mem_scheduler.get_tasks_status()
     mem_scheduler.print_tasks_status(tasks_status=tasks_status)

src/memos/api/handlers/scheduler_handler.py

@@ -22,10 +22,13 @@
     AllStatusResponseData,
     StatusResponse,
     StatusResponseItem,
+    TaskQueueData,
+    TaskQueueResponse,
     TaskSummary,
 )
 from memos.log import get_logger
 from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
 
 
@@ -243,6 +246,96 @@ def handle_scheduler_status(
         raise HTTPException(status_code=500, detail="Failed to get scheduler status") from err
 
 
+def handle_task_queue_status(
+    user_id: str, mem_scheduler: OptimizedScheduler, task_id: str | None = None
+) -> TaskQueueResponse:
+    try:
+        queue = getattr(mem_scheduler, "memos_message_queue", None)
+        if queue is None:
+            raise HTTPException(status_code=503, detail="Scheduler queue is not available")
+
+        # Only support Redis-backed queue for now; try lazy init if not connected
+        redis_conn = getattr(queue, "_redis_conn", None)
+        if redis_conn is None:
+            try:
+                if hasattr(queue, "auto_initialize_redis"):
+                    queue.auto_initialize_redis()
+                    redis_conn = getattr(queue, "_redis_conn", None)
+                if redis_conn and hasattr(queue, "connect"):
+                    queue.connect()
+            except Exception:
+                redis_conn = None
+
+        if redis_conn is None:
+            raise HTTPException(status_code=503, detail="Scheduler queue not connected to Redis")
+
+        stream_keys = queue.get_stream_keys()
+        # Filter by user_id; stream key format: {prefix}:{user_id}:{mem_cube_id}:{task_label}
+        user_stream_keys = [sk for sk in stream_keys if f":{user_id}:" in sk]
+
+        if not user_stream_keys:
+            raise HTTPException(
+                status_code=404, detail=f"No scheduler streams found for user {user_id}"
+            )
+
+        def _parse_user_id_from_stream(stream_key: str) -> str | None:
+            try:
+                parts = stream_key.split(":")
+                if len(parts) < 3:
+                    return None
+                # prefix may contain multiple segments; user_id is the 2nd segment from the end - 1
+                return parts[-3]
+            except Exception:
+                return None
+
+        user_ids_present = {
+            uid for uid in (_parse_user_id_from_stream(sk) for sk in stream_keys) if uid
+        }
+
+        pending_total = 0
+        pending_detail: list[str] = []
+        remaining_total = 0
+        remaining_detail: list[str] = []
+
+        consumer_group = getattr(queue, "consumer_group", None) or "scheduler_group"
+        for sk in user_stream_keys:
+            try:
+                pending_info = redis_conn.xpending(sk, consumer_group)
+                pending_count = pending_info[0] if pending_info else 0
+            except Exception:
+                pending_count = 0
+            pending_total += pending_count
+            pending_detail.append(f"{sk}:{pending_count}")
+
+            try:
+                remaining_count = redis_conn.xlen(sk)
+            except Exception:
+                remaining_count = 0
+            remaining_total += remaining_count
+            remaining_detail.append(f"{sk}:{remaining_count}")
+
+        data = TaskQueueData(
+            user_id=user_id,
+            user_name=None,
+            mem_cube_id=None,
+            stream_keys=user_stream_keys,
+            users_count=len(user_ids_present),
+            pending_tasks_count=pending_total,
+            remaining_tasks_count=remaining_total,
+            pending_tasks_detail=pending_detail,
+            remaining_tasks_detail=remaining_detail,
+        )
+        return TaskQueueResponse(data=data)
+    except HTTPException:
+        # Re-raise HTTPException directly to preserve its status code (e.g., 404)
+        raise
+    except Exception as err:
+        logger.error(
+            f"Failed to get task queue status for user {user_id}: {traceback.format_exc()}"
+        )
+        raise HTTPException(status_code=500, detail="Failed to get scheduler status") from err
+
+
 def handle_scheduler_wait(
     user_name: str,
     status_tracker: TaskStatusTracker,

src/memos/api/product_models.py

@@ -884,6 +884,34 @@ class StatusResponse(BaseResponse[list[StatusResponseItem]]):
     message: str = "Memory get status successfully"
 
 
+class TaskQueueData(BaseModel):
+    """Queue-level metrics for scheduler tasks."""
+
+    user_id: str = Field(..., description="User ID the query is scoped to")
+    user_name: str | None = Field(None, description="User name if available")
+    mem_cube_id: str | None = Field(
+        None, description="MemCube ID if a single cube is targeted; otherwise None"
+    )
+    stream_keys: list[str] = Field(..., description="Matched Redis stream keys for this user")
+    users_count: int = Field(..., description="Distinct users currently present in queue streams")
+    pending_tasks_count: int = Field(
+        ..., description="Count of pending (delivered, not acked) tasks"
+    )
+    remaining_tasks_count: int = Field(..., description="Count of enqueued tasks (xlen)")
+    pending_tasks_detail: list[str] = Field(
+        ..., description="Per-stream pending counts, formatted as '{stream_key}:{count}'"
+    )
+    remaining_tasks_detail: list[str] = Field(
+        ..., description="Per-stream remaining counts, formatted as '{stream_key}:{count}'"
+    )
+
+
+class TaskQueueResponse(BaseResponse[TaskQueueData]):
+    """Response model for scheduler task queue status."""
+
+    message: str = "Scheduler task queue status retrieved successfully"
+
+
 class TaskSummary(BaseModel):
     """Aggregated counts of tasks by status."""
 

src/memos/api/routers/server_router.py

@@ -41,6 +41,7 @@
     StatusResponse,
     SuggestionRequest,
     SuggestionResponse,
+    TaskQueueResponse,
 )
 from memos.log import get_logger
 from memos.mem_scheduler.base_scheduler import BaseScheduler
@@ -143,6 +144,20 @@ def scheduler_status(
     )
 
 
+@router.get(  # Changed from post to get
+    "/scheduler/task_queue_status",
+    summary="Get scheduler task queue status",
+    response_model=TaskQueueResponse,
+)
+def scheduler_task_queue_status(
+    user_id: str = Query(..., description="User ID whose queue status is requested"),
+):
+    """Get scheduler task queue backlog/pending status for a user."""
+    return handlers.scheduler_handler.handle_task_queue_status(
+        user_id=user_id, mem_scheduler=mem_scheduler
+    )
+
+
 @router.post("/scheduler/wait", summary="Wait until scheduler is idle for a specific user")
 def scheduler_wait(
     user_name: str,

src/memos/configs/mem_reader.py

@@ -44,6 +44,7 @@ def parse_datetime(cls, value):
 class SimpleStructMemReaderConfig(BaseMemReaderConfig):
     """SimpleStruct MemReader configuration class."""
 
+    # Allow passing additional fields without raising validation errors
     model_config = ConfigDict(extra="allow", strict=True)
 
 

src/memos/mem_scheduler/base_scheduler.py

@@ -4,6 +4,7 @@
 import time
 
 from collections.abc import Callable
+from contextlib import suppress
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING, Union
@@ -47,6 +48,15 @@
     ScheduleMessageItem,
 )
 from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ADD_TASK_LABEL,
+    ANSWER_TASK_LABEL,
+    MEM_ARCHIVE_TASK_LABEL,
+    MEM_ORGANIZE_TASK_LABEL,
+    MEM_UPDATE_TASK_LABEL,
+    QUERY_TASK_LABEL,
+    TaskPriorityLevel,
+)
 from memos.mem_scheduler.task_schedule_modules.dispatcher import SchedulerDispatcher
 from memos.mem_scheduler.task_schedule_modules.orchestrator import SchedulerOrchestrator
 from memos.mem_scheduler.task_schedule_modules.task_queue import ScheduleTaskQueue
@@ -55,6 +65,7 @@
 from memos.mem_scheduler.utils.filter_utils import (
     transform_name_to_key,
 )
+from memos.mem_scheduler.utils.misc_utils import group_messages_by_user_and_mem_cube
 from memos.mem_scheduler.utils.monitor_event_utils import emit_monitor_event, to_iso
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
 from memos.mem_scheduler.webservice_modules.rabbitmq_service import RabbitMQSchedulerModule
@@ -642,19 +653,115 @@ def update_activation_memory_periodically(
             logger.error(f"Error in update_activation_memory_periodically: {e}", exc_info=True)
 
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
+        """Submit messages for processing, with priority-aware dispatch.
+
+        - LEVEL_1 tasks dispatch immediately to the appropriate handler.
+        - Lower-priority tasks are enqueued via the configured message queue.
+        """
         if isinstance(messages, ScheduleMessageItem):
             messages = [messages]
-        for message in messages:
-            self.metrics.task_enqueued(user_id=message.user_id, task_type=message.label)
+
+        if not messages:
+            return
+
+        immediate_msgs: list[ScheduleMessageItem] = []
+        queued_msgs: list[ScheduleMessageItem] = []
+
+        for msg in messages:
+            # basic metrics and status tracking
+            with suppress(Exception):
+                self.metrics.task_enqueued(user_id=msg.user_id, task_type=msg.label)
+
+            # ensure timestamp exists for monitoring
+            if getattr(msg, "timestamp", None) is None:
+                msg.timestamp = get_utc_now()
+
             if self.status_tracker:
-                self.status_tracker.task_submitted(
-                    task_id=message.item_id,
-                    user_id=message.user_id,
-                    task_type=message.label,
-                    mem_cube_id=message.mem_cube_id,
-                    business_task_id=message.task_id,  # Pass business task_id if provided
+                try:
+                    self.status_tracker.task_submitted(
+                        task_id=msg.item_id,
+                        user_id=msg.user_id,
+                        task_type=msg.label,
+                        mem_cube_id=msg.mem_cube_id,
+                        business_task_id=msg.task_id,
+                    )
+                except Exception:
+                    logger.warning("status_tracker.task_submitted failed", exc_info=True)
+
+            # honor disabled handlers
+            if self.disabled_handlers and msg.label in self.disabled_handlers:
+                logger.info(f"Skipping disabled handler: {msg.label} - {msg.content}")
+                continue
+
+            # decide priority path
+            task_priority = self.orchestrator.get_task_priority(task_label=msg.label)
+            if task_priority == TaskPriorityLevel.LEVEL_1:
+                immediate_msgs.append(msg)
+            else:
+                queued_msgs.append(msg)
+
+        # Dispatch high-priority tasks immediately
+        if immediate_msgs:
+            # emit enqueue events for consistency
+            for m in immediate_msgs:
+                emit_monitor_event(
+                    "enqueue", m, {"enqueue_ts": to_iso(getattr(m, "timestamp", None))}
                 )
-        self.memos_message_queue.submit_messages(messages=messages)
+
+            # simulate dequeue for immediately dispatched messages so monitor logs stay complete
+            for m in immediate_msgs:
+                try:
+                    now = time.time()
+                    enqueue_ts_obj = getattr(m, "timestamp", None)
+                    enqueue_epoch = None
+                    if isinstance(enqueue_ts_obj, int | float):
+                        enqueue_epoch = float(enqueue_ts_obj)
+                    elif hasattr(enqueue_ts_obj, "timestamp"):
+                        dt = enqueue_ts_obj
+                        if dt.tzinfo is None:
+                            dt = dt.replace(tzinfo=timezone.utc)
+                        enqueue_epoch = dt.timestamp()
+
+                    queue_wait_ms = None
+                    if enqueue_epoch is not None:
+                        queue_wait_ms = max(0.0, now - enqueue_epoch) * 1000
+
+                    object.__setattr__(m, "_dequeue_ts", now)
+                    emit_monitor_event(
+                        "dequeue",
+                        m,
+                        {
+                            "enqueue_ts": to_iso(enqueue_ts_obj),
+                            "dequeue_ts": datetime.fromtimestamp(now, tz=timezone.utc).isoformat(),
+                            "queue_wait_ms": queue_wait_ms,
+                        },
+                    )
+                    self.metrics.task_dequeued(user_id=m.user_id, task_type=m.label)
+                except Exception:
+                    logger.debug("Failed to emit dequeue for immediate task", exc_info=True)
+
+            user_cube_groups = group_messages_by_user_and_mem_cube(immediate_msgs)
+            for user_id, cube_groups in user_cube_groups.items():
+                for mem_cube_id, user_cube_msgs in cube_groups.items():
+                    label_groups: dict[str, list[ScheduleMessageItem]] = {}
+                    for m in user_cube_msgs:
+                        label_groups.setdefault(m.label, []).append(m)
+
+                    for label, msgs_by_label in label_groups.items():
+                        handler = self.dispatcher.handlers.get(
+                            label, self.dispatcher._default_message_handler
+                        )
+                        self.dispatcher.execute_task(
+                            user_id=user_id,
+                            mem_cube_id=mem_cube_id,
+                            task_label=label,
+                            msgs=msgs_by_label,
+                            handler_call_back=handler,
+                        )
+
+        # Enqueue lower-priority tasks
+        if queued_msgs:
+            self.memos_message_queue.submit_messages(messages=queued_msgs)
 
     def _submit_web_logs(
         self,
@@ -706,15 +813,6 @@ def get_web_log_messages(self) -> list[dict]:
                 break
 
         def _map_label(label: str) -> str:
-            from memos.mem_scheduler.schemas.task_schemas import (
-                ADD_TASK_LABEL,
-                ANSWER_TASK_LABEL,
-                MEM_ARCHIVE_TASK_LABEL,
-                MEM_ORGANIZE_TASK_LABEL,
-                MEM_UPDATE_TASK_LABEL,
-                QUERY_TASK_LABEL,
-            )
-
             mapping = {
                 QUERY_TASK_LABEL: "addMessage",
                 ANSWER_TASK_LABEL: "addMessage",