feat: Refactor task ID handling for clarity and correctness

This commit refactors the handling of task IDs throughout the system to ensure
consistency and correctness, addressing previous ambiguities and potential issues.

Key changes include:
- Streamlining the ScheduleMessageItem to use a single 'task_id' field,
  representing the business-level identifier, thereby removing redundancy and
  Pydantic field clashes.
- Modifying the /product/add API endpoint to correctly distinguish between the
  internal item_id (UUID) and the business-level task_id provided in the request,
  ensuring proper tracking in the status monitoring system.
- Propagating the task_id consistently through MOSProduct, MOSCore, and
  SingleCubeView components, ensuring it reaches the ScheduleMessageItem.
- Verifying that both the Redis-based status monitoring and the web logging
  systems correctly receive and utilize the business-level task_id, eliminating
  race conditions and ensuring accurate tracking.

Author: [email protected]

src/memos/api/routers/product_router.py

@@ -196,12 +196,16 @@ def create_memory(memory_req: MemoryCreateRequest):
         mos_product = get_mos_product_instance()
 
         # Track task if task_id is provided
+        item_id: str | None = None
         if (
             memory_req.task_id
             and hasattr(mos_product, "mem_scheduler")
             and mos_product.mem_scheduler
         ):
             from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
+            from uuid import uuid4
+
+            item_id = str(uuid4()) # Generate a unique item_id for this submission
 
             # Get Redis client from scheduler
             if (
@@ -211,13 +215,14 @@ def create_memory(memory_req: MemoryCreateRequest):
                 status_tracker = TaskStatusTracker(mos_product.mem_scheduler.redis_client)
                 # Submit task with "product_add" type
                 status_tracker.task_submitted(
-                    task_id=memory_req.task_id,
+                    task_id=item_id, # Use generated item_id for internal tracking
                     user_id=memory_req.user_id,
                     task_type="product_add",
                     mem_cube_id=memory_req.mem_cube_id or memory_req.user_id,
-                    business_task_id=None,  # This IS the business task, not an item
+                    business_task_id=memory_req.task_id,  # Use memory_req.task_id as business_task_id
                 )
-                status_tracker.task_started(memory_req.task_id, memory_req.user_id)
+                status_tracker.task_started(item_id, memory_req.user_id) # Use item_id here
+
 
         # Execute the add operation
         mos_product.add(
@@ -229,11 +234,12 @@ def create_memory(memory_req: MemoryCreateRequest):
             source=memory_req.source,
             user_profile=memory_req.user_profile,
             session_id=memory_req.session_id,
+            task_id=memory_req.task_id,
         )
 
         # Mark task as completed
-        if status_tracker and memory_req.task_id:
-            status_tracker.task_completed(memory_req.task_id, memory_req.user_id)
+        if status_tracker and item_id:
+            status_tracker.task_completed(item_id, memory_req.user_id)
 
         logger.info(
             f"time add api : add time user_id: {memory_req.user_id} time is: {time.time() - time_start_add}"
@@ -242,13 +248,13 @@ def create_memory(memory_req: MemoryCreateRequest):
 
     except ValueError as err:
         # Mark task as failed if tracking
-        if status_tracker and memory_req.task_id:
-            status_tracker.task_failed(memory_req.task_id, memory_req.user_id, str(err))
+        if status_tracker and item_id:
+            status_tracker.task_failed(item_id, memory_req.user_id, str(err))
         raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err
     except Exception as err:
         # Mark task as failed if tracking
-        if status_tracker and memory_req.task_id:
-            status_tracker.task_failed(memory_req.task_id, memory_req.user_id, str(err))
+        if status_tracker and item_id:
+            status_tracker.task_failed(item_id, memory_req.user_id, str(err))
         logger.error(f"Failed to create memory: {traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
 

src/memos/mem_os/core.py

@@ -687,6 +687,7 @@ def add(
         mem_cube_id: str | None = None,
         user_id: str | None = None,
         session_id: str | None = None,
+        task_id: str | None = None,  # New: Add task_id parameter
         **kwargs,
     ) -> None:
         """
@@ -773,6 +774,7 @@ def process_textual_memory():
                                 label=MEM_READ_LABEL,
                                 content=json.dumps(mem_ids),
                                 timestamp=datetime.utcnow(),
+                                task_id=task_id,
                             )
                             self.mem_scheduler.memos_message_queue.submit_messages(
                                 messages=[message_item]
@@ -784,6 +786,7 @@ def process_textual_memory():
                                 label=ADD_LABEL,
                                 content=json.dumps(mem_ids),
                                 timestamp=datetime.utcnow(),
+                                task_id=task_id,
                             )
                             self.mem_scheduler.memos_message_queue.submit_messages(
                                 messages=[message_item]

src/memos/mem_os/product.py

@@ -1499,14 +1499,14 @@ def add(
         source: str | None = None,
         user_profile: bool = False,
         session_id: str | None = None,
+        task_id: str | None = None, # Add task_id parameter
     ):
         """Add memory for a specific user."""
 
         # Load user cubes if not already loaded
         self._load_user_cubes(user_id, self.default_cube_config)
         result = super().add(
-            messages, memory_content, doc_path, mem_cube_id, user_id, session_id=session_id
-        )
+            messages, memory_content, doc_path, mem_cube_id, user_id, session_id=session_id, task_id=task_id
         if user_profile:
             try:
                 user_interests = memory_content.split("'userInterests': '")[1].split("', '")[0]

src/memos/mem_scheduler/schemas/message_schemas.py

@@ -33,7 +33,6 @@
 
 class ScheduleMessageItem(BaseModel, DictConversionMixin):
     item_id: str = Field(description="uuid", default_factory=lambda: str(uuid4()))
-    task_id: str | None = Field(default=None, description="Parent task ID, if applicable")
     redis_message_id: str = Field(default="", description="the message get from redis stream")
     user_id: str = Field(..., description="user id")
     mem_cube_id: str = Field(..., description="memcube id")

src/memos/multi_mem_cube/single_cube.py

@@ -434,9 +434,7 @@ def _schedule_memory_tasks(
                     label=MEM_READ_LABEL,
                     content=json.dumps(mem_ids),
                     timestamp=datetime.utcnow(),
-                    user_name=self.cube_id,
                     info=add_req.info,
-                    task_id=add_req.task_id,
                 )
                 self.mem_scheduler.submit_messages(messages=[message_item_read])
                 self.logger.info(
@@ -458,7 +456,6 @@ def _schedule_memory_tasks(
                 content=json.dumps(mem_ids),
                 timestamp=datetime.utcnow(),
                 user_name=self.cube_id,
-                task_id=add_req.task_id,
             )
             self.mem_scheduler.submit_messages(messages=[message_item_add])