fix bugs: completely addressed bugs about redis queue

Author: tangg555

examples/mem_scheduler/api_w_scheduler.py

@@ -25,9 +25,6 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
     print(f"My test handler received {len(messages)} messages:")
     for msg in messages:
         print(f" my_test_handler - {msg.item_id}: {msg.content}")
-        print(
-            f"{queue._redis_conn.xinfo_groups(queue.stream_key_prefix)} qsize: {queue.qsize()} messages:{messages}"
-        )
 
 
 # 2. Register the handler

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -150,7 +150,6 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                     self.metrics.on_done(label=m.label, mem_cube_id=m.mem_cube_id, now=time.time())
 
                 # acknowledge redis messages
-
                 if self.use_redis_queue and self.memos_message_queue is not None:
                     for msg in messages:
                         redis_message_id = msg.redis_message_id

src/memos/mem_scheduler/task_schedule_modules/local_queue.py

@@ -79,8 +79,7 @@ def put(
 
     def get(
         self,
-        user_id: str,
-        mem_cube_id: str,
+        stream_key: str,
         block: bool = True,
         timeout: float | None = None,
         batch_size: int | None = None,
@@ -91,8 +90,6 @@ def get(
             )
             return []
 
-        stream_key = self.get_stream_key(user_id=user_id, mem_cube_id=mem_cube_id)
-
         # Return empty list if queue does not exist
         if stream_key not in self.queue_streams:
             logger.error(f"Stream {stream_key} does not exist when trying to get messages.")

src/memos/mem_scheduler/task_schedule_modules/redis_queue.py

@@ -5,6 +5,7 @@
 the local memos_message_queue functionality in BaseScheduler.
 """
 
+import re
 import time
 
 from collections.abc import Callable
@@ -165,8 +166,7 @@ def ack_message(self, user_id, mem_cube_id, redis_message_id) -> None:
 
     def get(
         self,
-        user_id: str,
-        mem_cube_id: str,
+        stream_key: str,
         block: bool = True,
         timeout: float | None = None,
         batch_size: int | None = None,
@@ -175,8 +175,6 @@ def get(
             raise ConnectionError("Not connected to Redis. Redis connection not available.")
 
         try:
-            stream_key = self.get_stream_key(user_id=user_id, mem_cube_id=mem_cube_id)
-
             # Calculate timeout for Redis
             redis_timeout = None
             if block and timeout is not None:
@@ -295,17 +293,21 @@ def get_stream_keys(self) -> list[str]:
         if not self._redis_conn:
             return []
 
-        try:
-            # Use match parameter and decode byte strings to regular strings
-            stream_keys = [
-                key.decode("utf-8") if isinstance(key, bytes) else key
-                for key in self._redis_conn.scan_iter(match=f"{self.stream_key_prefix}:*")
-            ]
-            logger.debug(f"get stream_keys from redis: {stream_keys}")
-            return stream_keys
-        except Exception as e:
-            logger.error(f"Failed to list Redis stream keys: {e}")
-            return []
+        # First, get all keys that might match (using Redis pattern matching)
+        redis_pattern = f"{self.stream_key_prefix}:*"
+        raw_keys = [
+            key.decode("utf-8") if isinstance(key, bytes) else key
+            for key in self._redis_conn.scan_iter(match=redis_pattern)
+        ]
+
+        # Second, filter using Python regex to ensure exact prefix match
+        # Escape special regex characters in the prefix, then add :.*
+        escaped_prefix = re.escape(self.stream_key_prefix)
+        regex_pattern = f"^{escaped_prefix}:"
+        stream_keys = [key for key in raw_keys if re.match(regex_pattern, key)]
+
+        logger.debug(f"get stream_keys from redis: {stream_keys}")
+        return stream_keys
 
     def size(self) -> int:
         """

src/memos/mem_scheduler/task_schedule_modules/task_queue.py

@@ -5,8 +5,6 @@
 the local memos_message_queue functionality in BaseScheduler.
 """
 
-from collections import defaultdict
-
 from memos.log import get_logger
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_scheduler.task_schedule_modules.local_queue import SchedulerLocalQueue
@@ -58,9 +56,10 @@ def debug_mode_on(self):
 
     def get_stream_keys(self) -> list[str]:
         if isinstance(self.memos_message_queue, SchedulerRedisQueue):
-            return self.memos_message_queue.get_stream_keys()
+            stream_keys = self.memos_message_queue.get_stream_keys()
         else:
-            return list(self.memos_message_queue.queue_streams.keys())
+            stream_keys = list(self.memos_message_queue.queue_streams.keys())
+        return stream_keys
 
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit messages to the message queue (either local queue or Redis)."""
@@ -98,46 +97,21 @@ def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageIt
                         )
 
     def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
-        # Discover all active streams via queue API
-        streams: list[tuple[str, str]] = []
-
         stream_keys = self.get_stream_keys()
-        for stream_key in stream_keys:
-            try:
-                parts = stream_key.split(":")
-                if len(parts) >= 3:
-                    user_id = parts[-2]
-                    mem_cube_id = parts[-1]
-                    streams.append((user_id, mem_cube_id))
-            except Exception as e:
-                logger.debug(f"Failed to parse stream key {stream_key}: {e}")
-
-        if not streams:
+
+        if len(stream_keys) == 0:
             return []
 
         messages: list[ScheduleMessageItem] = []
 
-        # Group by user: {user_id: [mem_cube_id, ...]}
-
-        streams_by_user: dict[str, list[str]] = defaultdict(list)
-        for user_id, mem_cube_id in streams:
-            streams_by_user[user_id].append(mem_cube_id)
-
-        # For each user, fairly consume up to batch_size across their streams
-        for user_id, mem_cube_ids in streams_by_user.items():
-            if not mem_cube_ids:
-                continue
-
-            # First pass: give each stream an equal share for this user
-            for mem_cube_id in mem_cube_ids:
-                fetched = self.memos_message_queue.get(
-                    user_id=user_id,
-                    mem_cube_id=mem_cube_id,
-                    block=False,
-                    batch_size=batch_size,
-                )
+        for stream_key in stream_keys:
+            fetched = self.memos_message_queue.get(
+                stream_key=stream_key,
+                block=False,
+                batch_size=batch_size,
+            )
 
-                messages.extend(fetched)
+            messages.extend(fetched)
 
         logger.info(
             f"Fetched {len(messages)} messages across users with per-user batch_size={batch_size}"