Hotfix/task id loss

examples/data/config/mem_scheduler/general_scheduler_config.yaml

@@ -4,7 +4,7 @@ config:
   act_mem_update_interval: 30
   context_window_size: 10
   thread_pool_max_workers: 5
-  consume_interval_seconds: 1
+  consume_interval_seconds: 0.01
   working_mem_monitor_capacity: 20
   activation_mem_monitor_capacity: 5
   enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_optimized_scheduler.yaml

@@ -38,7 +38,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
-    consume_interval_seconds: 1
+    consume_interval_seconds: 0.01
     working_mem_monitor_capacity: 20
     activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -38,7 +38,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
-    consume_interval_seconds: 1
+    consume_interval_seconds: 0.01
     working_mem_monitor_capacity: 20
     activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true

examples/mem_scheduler/task_fair_schedule.py

@@ -0,0 +1,88 @@
+import sys
+
+from collections import defaultdict
+from pathlib import Path
+
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent
+sys.path.insert(0, str(BASE_DIR))
+
+
+def make_message(user_id: str, mem_cube_id: str, label: str, idx: int | str) -> ScheduleMessageItem:
+    return ScheduleMessageItem(
+        item_id=f"{user_id}:{mem_cube_id}:{label}:{idx}",
+        user_id=user_id,
+        mem_cube_id=mem_cube_id,
+        label=label,
+        content=f"msg-{idx} for {user_id}/{mem_cube_id}/{label}",
+    )
+
+
+def seed_messages_for_test_fairness(queue, combos, per_stream):
+    # send overwhelm message by one user
+    (u, c, label) = combos[0]
+    task_target = 100
+    print(f"{u}:{c}:{label} submit {task_target} messages")
+    for i in range(task_target):
+        msg = make_message(u, c, label, f"overwhelm_{i}")
+        queue.submit_messages(msg)
+
+    for u, c, label in combos:
+        print(f"{u}:{c}:{label} submit {per_stream} messages")
+        for i in range(per_stream):
+            msg = make_message(u, c, label, i)
+            queue.submit_messages(msg)
+    print("======= seed_messages Done ===========")
+
+
+def count_by_stream(messages):
+    counts = defaultdict(int)
+    for m in messages:
+        key = f"{m.user_id}:{m.mem_cube_id}:{m.label}"
+        counts[key] += 1
+    return counts
+
+
+def run_fair_redis_schedule(batch_size: int = 3):
+    print("=== Redis Fairness Demo ===")
+    print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
+    mem_scheduler.consume_batch = batch_size
+    queue = mem_scheduler.memos_message_queue
+
+    # Isolate and clear queue
+    queue.debug_mode_on(debug_stream_prefix="fair_redis_schedule")
+    queue.clear()
+
+    # Define multiple streams: (user_id, mem_cube_id, task_label)
+    combos = [
+        ("u1", "u1", "labelX"),
+        ("u1", "u1", "labelY"),
+        ("u2", "u2", "labelX"),
+        ("u2", "u2", "labelY"),
+    ]
+    per_stream = 5
+
+    # Seed messages evenly across streams
+    seed_messages_for_test_fairness(queue, combos, per_stream)
+
+    # Compute target batch size (fair split across streams)
+    print(f"Request batch_size={batch_size} for {len(combos)} streams")
+
+    for _ in range(len(combos)):
+        # Fetch one brokered pack
+        msgs = queue.get_messages(batch_size=batch_size)
+        print(f"Fetched {len(msgs)} messages in first pack")
+
+        # Check fairness: counts per stream
+        counts = count_by_stream(msgs)
+        for k in sorted(counts):
+            print(f"{k}: {counts[k]}")
+
+
+if __name__ == "__main__":
+    # task 1 fair redis schedule
+    run_fair_redis_schedule()

examples/mem_scheduler/task_stop_rerun.py

@@ -0,0 +1,86 @@
+from pathlib import Path
+from time import sleep
+
+# Note: we skip API handler status/wait utilities in this demo
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+# Debug: Print scheduler configuration
+print("=== Scheduler Configuration Debug ===")
+print(f"Scheduler type: {type(mem_scheduler).__name__}")
+print(f"Config: {mem_scheduler.config}")
+print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
+print(f"Queue type: {type(mem_scheduler.memos_message_queue).__name__}")
+print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
+print("=====================================\n")
+
+queue = mem_scheduler.memos_message_queue
+queue.debug_mode_on(debug_stream_prefix="task_stop_rerun")
+
+
+# Define a handler function
+def my_test_handler(messages: list[ScheduleMessageItem]):
+    print(f"My test handler received {len(messages)} messages: {[one.item_id for one in messages]}")
+    for msg in messages:
+        # Create a file named by task_id (use item_id as numeric id 0..99)
+        task_id = str(msg.item_id)
+        file_path = tmp_dir / f"{task_id}.txt"
+        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}")
+
+
+def submit_tasks():
+    mem_scheduler.memos_message_queue.clear()
+
+    # Create 100 messages (task_id 0..99)
+    users = ["user_A", "user_B"]
+    messages_to_send = [
+        ScheduleMessageItem(
+            item_id=str(i),
+            user_id=users[i % 2],
+            mem_cube_id="test_mem_cube",
+            label=TEST_HANDLER_LABEL,
+            content=f"Create file for task {i}",
+        )
+        for i in range(100)
+    ]
+    # Submit messages in batch and print completion
+    print(f"Submitting {len(messages_to_send)} messages to the scheduler...")
+    mem_scheduler.memos_message_queue.submit_messages(messages_to_send)
+    print(f"Task submission done! tasks in queue: {mem_scheduler.get_tasks_status()}")
+
+
+# Register the handler
+TEST_HANDLER_LABEL = "test_handler"
+mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})
+
+
+tmp_dir = Path("./tmp")
+tmp_dir.mkdir(exist_ok=True)
+
+# Test stop-and-restart: if tmp already has >1 files, skip submission and print info
+existing_count = len(list(Path("tmp").glob("*.txt"))) if Path("tmp").exists() else 0
+if existing_count > 1:
+    print(f"Skip submission: found {existing_count} files in tmp (>1), continue processing")
+else:
+    submit_tasks()
+
+# 6. Wait until tmp has 100 files or timeout
+poll_interval = 1
+expected = 100
+tmp_dir = Path("tmp")
+while mem_scheduler.get_tasks_status()["remaining"] != 0:
+    count = len(list(tmp_dir.glob("*.txt"))) if tmp_dir.exists() else 0
+    user_status_running = mem_scheduler.get_tasks_status()
+    print(f"[Monitor] user_status_running: {user_status_running}; Files in tmp: {count}/{expected}")
+    sleep(poll_interval)
+print(f"[Result] Final files in tmp: {len(list(tmp_dir.glob('*.txt')))})")
+
+# 7. Stop the scheduler
+print("Stopping the scheduler...")
+mem_scheduler.stop()

src/memos/api/handlers/add_handler.py

@@ -45,7 +45,9 @@ def handle_add_memories(self, add_req: APIADDRequest) -> MemoryResponse:
         Returns:
             MemoryResponse with added memory information
         """
-        self.logger.info(f"[AddHandler] Add Req is: {add_req}")
+        self.logger.info(
+            f"[DIAGNOSTIC] server_router -> add_handler.handle_add_memories called (Modified at 2025-11-29 18:46). Full request: {add_req.model_dump_json(indent=2)}"
+        )
 
         if add_req.info:
             exclude_fields = list_all_fields()

src/memos/api/handlers/base_handler.py

@@ -8,7 +8,7 @@
 from typing import Any
 
 from memos.log import get_logger
-from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
 from memos.memories.textual.tree_text_memory.retrieve.advanced_searcher import AdvancedSearcher
 
 
@@ -127,7 +127,7 @@ def mem_reader(self):
         return self.deps.mem_reader
 
     @property
-    def mem_scheduler(self) -> BaseScheduler:
+    def mem_scheduler(self) -> OptimizedScheduler:
         """Get scheduler instance."""
         return self.deps.mem_scheduler
 

src/memos/api/routers/product_router.py

@@ -188,6 +188,7 @@ def get_all_memories(memory_req: GetMemoryPlaygroundRequest):
 @router.post("/add", summary="add a new memory", response_model=SimpleResponse)
 def create_memory(memory_req: MemoryCreateRequest):
     """Create a new memory for a specific user."""
+    logger.info("DIAGNOSTIC: /product/add endpoint called. This confirms the new code is deployed.")
     # Initialize status_tracker outside try block to avoid NameError in except blocks
     status_tracker = None
 

src/memos/mem_os/core.py

@@ -788,6 +788,9 @@ def process_textual_memory():
                                 timestamp=datetime.utcnow(),
                                 task_id=task_id,
                             )
+                            logger.info(
+                                f"[DIAGNOSTIC] core.add: Submitting message to scheduler: {message_item.model_dump_json(indent=2)}"
+                            )
                             self.mem_scheduler.memos_message_queue.submit_messages(
                                 messages=[message_item]
                             )

src/memos/mem_os/utils/default_config.py

@@ -110,7 +110,7 @@ def get_default_config(
                 "act_mem_update_interval": kwargs.get("scheduler_act_mem_update_interval", 300),
                 "context_window_size": kwargs.get("scheduler_context_window_size", 5),
                 "thread_pool_max_workers": kwargs.get("scheduler_thread_pool_max_workers", 10),
-                "consume_interval_seconds": kwargs.get("scheduler_consume_interval_seconds", 3),
+                "consume_interval_seconds": kwargs.get("scheduler_consume_interval_seconds", 0.01),
                 "enable_parallel_dispatch": kwargs.get("scheduler_enable_parallel_dispatch", True),
                 "enable_activation_memory": True,
             },

src/memos/mem_scheduler/base_scheduler.py

@@ -137,7 +137,6 @@ def __init__(self, config: BaseSchedulerConfig):
         self.dispatcher = SchedulerDispatcher(
             config=self.config,
             memos_message_queue=self.memos_message_queue,
-            use_redis_queue=self.use_redis_queue,
             max_workers=self.thread_pool_max_workers,
             enable_parallel_dispatch=self.enable_parallel_dispatch,
             status_tracker=self.status_tracker,
@@ -232,8 +231,8 @@ def initialize_modules(
 
         # start queue monitor if enabled and a bot is set later
 
-    def debug_mode_on(self):
-        self.memos_message_queue.debug_mode_on()
+    def debug_mode_on(self, debug_stream_prefix="debug_mode"):
+        self.memos_message_queue.debug_mode_on(debug_stream_prefix=debug_stream_prefix)
 
     def _cleanup_on_init_failure(self):
         """Clean up resources if initialization fails."""
@@ -594,6 +593,11 @@ def _submit_web_logs(
         Args:
             messages: Single log message or list of log messages
         """
+        messages_list = [messages] if isinstance(messages, ScheduleLogForWebItem) else messages
+        for message in messages_list:
+            logger.info(
+                f"[DIAGNOSTIC] base_scheduler._submit_web_logs called. Message to publish: {message.model_dump_json(indent=2)}"
+            )
         if self.rabbitmq_config is None:
             return
 

src/memos/mem_scheduler/general_modules/scheduler_logger.py

@@ -113,9 +113,10 @@ def create_event_log(
         metadata: list[dict],
         memory_len: int,
         memcube_name: str | None = None,
+        log_content: str | None = None,
     ) -> ScheduleLogForWebItem:
         item = self.create_autofilled_log_item(
-            log_content="",
+            log_content=log_content or "",
             label=label,
             from_memory_type=from_memory_type,
             to_memory_type=to_memory_type,

src/memos/mem_scheduler/general_scheduler.py

@@ -367,16 +367,19 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                             if kb_log_content:
                                 event = self.create_event_log(
                                     label="knowledgeBaseUpdate",
+                                    # 1. 移除 log_content 参数
+                                    # 2. 补充 memory_type
                                     from_memory_type=USER_INPUT_TYPE,
                                     to_memory_type=LONG_TERM_MEMORY_TYPE,
                                     user_id=msg.user_id,
                                     mem_cube_id=msg.mem_cube_id,
                                     mem_cube=self.current_mem_cube,
                                     memcube_log_content=kb_log_content,
-                                    metadata=None,  # Per design doc for KB logs
+                                    metadata=None,
                                     memory_len=len(kb_log_content),
                                     memcube_name=self._map_memcube_name(msg.mem_cube_id),
                                 )
+                                # 3. 后置赋值 log_content
                                 event.log_content = (
                                     f"Knowledge Base Memory Update: {len(kb_log_content)} changes."
                                 )
@@ -474,6 +477,9 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
             logger.error(f"Error: {e}", exc_info=True)
 
     def _mem_read_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
+        logger.info(
+            f"[DIAGNOSTIC] general_scheduler._mem_read_message_consumer called. Received messages: {[msg.model_dump_json(indent=2) for msg in messages]}"
+        )
         logger.info(f"Messages {messages} assigned to {MEM_READ_LABEL} handler.")
 
         def process_message(message: ScheduleMessageItem):
@@ -538,6 +544,9 @@ def _process_memories_with_reader(
         task_id: str | None = None,
         info: dict | None = None,
     ) -> None:
+        logger.info(
+            f"[DIAGNOSTIC] general_scheduler._process_memories_with_reader called. mem_ids: {mem_ids}, user_id: {user_id}, mem_cube_id: {mem_cube_id}, task_id: {task_id}"
+        )
         """
         Process memories using mem_reader for enhanced memory processing.
 
@@ -635,6 +644,9 @@ def _process_memories_with_reader(
                                 }
                             )
                         if kb_log_content:
+                            logger.info(
+                                f"[DIAGNOSTIC] general_scheduler._process_memories_with_reader: Creating event log for KB update. Label: knowledgeBaseUpdate, user_id: {user_id}, mem_cube_id: {mem_cube_id}, task_id: {task_id}. KB content: {json.dumps(kb_log_content, indent=2)}"
+                            )
                             event = self.create_event_log(
                                 label="knowledgeBaseUpdate",
                                 from_memory_type=USER_INPUT_TYPE,

src/memos/mem_scheduler/memory_manage_modules/retriever.py

@@ -209,10 +209,9 @@ def _split_batches(
     def recall_for_missing_memories(
         self,
         query: str,
-        memories: list[TextualMemoryItem],
+        memories: list[str],
     ) -> tuple[str, bool]:
-        text_memories = [one.memory for one in memories] if memories else []
-        text_memories = "\n".join([f"- {mem}" for i, mem in enumerate(text_memories)])
+        text_memories = "\n".join([f"- {mem}" for i, mem in enumerate(memories)])
 
         prompt = self.build_prompt(
             template_name="enlarge_recall",