Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ def _process_memories_with_reader(

info = dict(info or {})
is_upload_skill = info.pop("is_upload_skill", False)
transfer_succeeded = False

try:
processed_memories = mem_reader.fine_transfer_simple_mem(
Expand All @@ -178,6 +179,7 @@ def _process_memories_with_reader(
user_context=user_context,
is_upload_skill=is_upload_skill,
)
transfer_succeeded = True
except Exception as e:
logger.warning("%s: Fail to transfer mem: %s", e, memory_items)
processed_memories = []
Expand Down Expand Up @@ -402,7 +404,7 @@ def _process_memories_with_reader(
if bindings_to_delete:
delete_ids.extend(list(bindings_to_delete))
delete_ids = list(dict.fromkeys(delete_ids))
if delete_ids:
if delete_ids and transfer_succeeded:
try:
if getattr(mem_reader, "memory_version_switch", "off") != "on":
text_mem.delete(delete_ids, user_name=user_name)
Expand Down Expand Up @@ -434,6 +436,11 @@ def _process_memories_with_reader(
)
except Exception as e:
logger.warning("Failed to delete some mem_ids %s: %s", delete_ids, e)
elif delete_ids:
logger.info(
"Skip deleting raw/working mem_ids after failed transfer: %s",
delete_ids,
)
else:
logger.info("No mem_ids to delete (nothing to cleanup)")

Expand Down
38 changes: 38 additions & 0 deletions tests/mem_scheduler/test_mem_read_handler.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
from types import SimpleNamespace
from unittest.mock import MagicMock

from memos.mem_scheduler.task_schedule_modules.handlers.mem_read_handler import (
MemReadMessageHandler,
)


def test_mem_read_handler_keeps_fast_memories_when_transfer_fails():
mem_reader = MagicMock()
mem_reader.fine_transfer_simple_mem.side_effect = RuntimeError("extract failed")
mem_reader.memory_version_switch = "off"

raw_memory = SimpleNamespace(
memory="raw user message",
metadata=SimpleNamespace(background="", memory_type="LongTermMemory"),
)

text_mem = MagicMock()
text_mem.get.return_value = raw_memory
text_mem.memory_manager = SimpleNamespace(
remove_and_refresh_memory=MagicMock(),
)

scheduler_context = MagicMock()
scheduler_context.get_mem_reader.return_value = mem_reader

handler = MemReadMessageHandler(scheduler_context)
handler._process_memories_with_reader(
mem_ids=["fast-1"],
user_id="user-1",
mem_cube_id="cube-1",
text_mem=text_mem,
user_name="alice",
)

text_mem.delete.assert_not_called()
text_mem.soft_delete.assert_not_called()