From 22d8e9804474351b940e10f7c30f88bc38fb81fc Mon Sep 17 00:00:00 2001 From: MemOS AutoDev Date: Thu, 1 Oct 2026 00:14:19 +0800 Subject: [PATCH 1/2] fix: harden prepare_reference_data against dict entries with missing/non-string id (#2448) The dict branch of prepare_reference_data used to assume every entry carried a string "id" key, so cached search results / MCP payloads that had already been serialized would raise KeyError('id') or AttributeError on .split('-'). Both errors escaped the streaming pipeline. Normalize dict entries instead: create the metadata dict if missing, skip ref_id derivation when id is absent (keep id slot as None), and coerce non-string ids (int, uuid.UUID, ...) with str() before the prefix split. The original id value is preserved in metadata["id"] unchanged so downstream consumers can still round-trip it. Regression tests in tests/mem_os/utils/test_reference_utils.py cover both crash repros from the issue plus the TextualMemoryItem branch (baseline). --- src/memos/mem_os/utils/reference_utils.py | 36 ++++- tests/mem_os/utils/test_reference_utils.py | 149 +++++++++++++++++++++ 2 files changed, 178 insertions(+), 7 deletions(-) create mode 100644 tests/mem_os/utils/test_reference_utils.py diff --git a/src/memos/mem_os/utils/reference_utils.py b/src/memos/mem_os/utils/reference_utils.py index 09b812207..4c225ad8f 100644 --- a/src/memos/mem_os/utils/reference_utils.py +++ b/src/memos/mem_os/utils/reference_utils.py @@ -139,7 +139,14 @@ def process_streaming_references_complete(text_buffer: str) -> tuple[str, str]: def prepare_reference_data(memories_list: list[TextualMemoryItem]) -> list[dict]: - # Prepare reference data + """Normalize a mixed list of memory entries into reference dicts. + + Accepts both :class:`TextualMemoryItem` objects and pre-serialized dicts + (e.g. cached search results or MCP payloads). Dict entries are handled + defensively: a missing ``id`` skips ref_id derivation instead of raising + ``KeyError``, and a non-string ``id`` (int, uuid.UUID, ...) is stringified + before ``.split("-")`` instead of raising ``AttributeError``. See #2448. + """ reference = [] for memories in memories_list: if isinstance(memories, TextualMemoryItem): @@ -152,11 +159,26 @@ def prepare_reference_data(memories_list: list[TextualMemoryItem]) -> list[dict] reference.append({"metadata": memories_json["metadata"]}) else: memories_json = memories - memories_json["metadata"]["ref_id"] = f"{memories_json['id'].split('-')[0]}" - memories_json["metadata"]["embedding"] = [] - memories_json["metadata"]["sources"] = [] - memories_json["metadata"]["memory"] = memories_json["memory"] - memories_json["metadata"]["id"] = memories_json["id"] - reference.append({"metadata": memories_json["metadata"]}) + metadata = memories_json.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + memories_json["metadata"] = metadata + + raw_id = memories_json.get("id") + if raw_id is None: + # Missing id: skip ref_id derivation, keep id slot explicit. + metadata["ref_id"] = "" + metadata["id"] = None + else: + # Non-string id (int, uuid.UUID, ...) is coerced for the prefix + # split only; the original value is preserved in metadata["id"]. + str_id = str(raw_id) + metadata["ref_id"] = f"{str_id.split('-')[0]}" + metadata["id"] = raw_id + + metadata["embedding"] = [] + metadata["sources"] = [] + metadata["memory"] = memories_json.get("memory") + reference.append({"metadata": metadata}) return reference diff --git a/tests/mem_os/utils/test_reference_utils.py b/tests/mem_os/utils/test_reference_utils.py new file mode 100644 index 000000000..b625a5124 --- /dev/null +++ b/tests/mem_os/utils/test_reference_utils.py @@ -0,0 +1,149 @@ +""" +Regression tests for ``prepare_reference_data``. + +Related issue: https://github.com/MemTensor/MemOS/issues/2448 + +The dict branch of ``prepare_reference_data`` used to assume that every dict +carried an ``id`` key and that the ``id`` was always a string. Both +assumptions are false for already-serialized memory payloads (cached search +results, MCP messages, etc.) and led to ``KeyError`` / ``AttributeError`` +escaping the streaming pipeline. These tests pin the normalized behavior in +place so the regression cannot come back. +""" + +import uuid + +import pytest + +from memos.mem_os.utils.reference_utils import prepare_reference_data +from memos.memories.textual.item import TextualMemoryItem + + +class TestPrepareReferenceDataDictEntries: + """Dict entries should be normalized safely, never raise.""" + + def test_missing_id_does_not_raise_keyerror(self): + """Repro from #2448: dict entry without ``id`` used to raise KeyError.""" + entry = {"metadata": {"memory": "m"}} + + result = prepare_reference_data([entry]) + + assert len(result) == 1 + metadata = result[0]["metadata"] + # Missing id → ref_id derivation skipped; id slot preserved as None. + assert metadata["ref_id"] == "" + assert metadata["id"] is None + # Metadata bookkeeping still populated. + assert metadata["embedding"] == [] + assert metadata["sources"] == [] + + def test_non_string_id_is_stringified(self): + """Repro from #2448: int id used to raise AttributeError on .split.""" + entry = {"id": 12345, "memory": "m", "metadata": {}} + + result = prepare_reference_data([entry]) + + assert len(result) == 1 + metadata = result[0]["metadata"] + # ``.split("-")[0]`` on ``str(12345)`` == "12345". + assert metadata["ref_id"] == "12345" + # Original id value round-trips unchanged into metadata["id"]. + assert metadata["id"] == 12345 + assert metadata["memory"] == "m" + + def test_uuid_object_id_is_stringified(self): + """UUID objects (another non-string id shape) must also be tolerated.""" + raw_id = uuid.UUID("12345678-1234-5678-1234-567812345678") + entry = {"id": raw_id, "memory": "m", "metadata": {}} + + result = prepare_reference_data([entry]) + + metadata = result[0]["metadata"] + assert metadata["ref_id"] == "12345678" + assert metadata["id"] == raw_id + + def test_string_id_still_derives_prefix(self): + """Baseline: a well-formed dict entry keeps the pre-fix behavior.""" + entry = { + "id": "abcdef12-3456-7890-abcd-ef1234567890", + "memory": "hello", + "metadata": {"source": "x"}, + } + + result = prepare_reference_data([entry]) + + metadata = result[0]["metadata"] + assert metadata["ref_id"] == "abcdef12" + assert metadata["id"] == "abcdef12-3456-7890-abcd-ef1234567890" + assert metadata["memory"] == "hello" + assert metadata["source"] == "x" + assert metadata["embedding"] == [] + assert metadata["sources"] == [] + + def test_missing_metadata_dict_is_auto_created(self): + """A dict entry lacking metadata must not blow up on ["metadata"]["ref_id"] = ...""" + entry = {"id": "abcdef12-1111-2222-3333-444455556666", "memory": "m"} + + result = prepare_reference_data([entry]) + + metadata = result[0]["metadata"] + assert metadata["ref_id"] == "abcdef12" + assert metadata["id"] == "abcdef12-1111-2222-3333-444455556666" + assert metadata["memory"] == "m" + assert metadata["embedding"] == [] + assert metadata["sources"] == [] + + def test_missing_memory_falls_back_to_none(self): + """Missing memory key should not raise, just fall through as None.""" + entry = {"id": "abcdef12-aaaa-bbbb-cccc-ddddeeeeffff", "metadata": {}} + + result = prepare_reference_data([entry]) + + metadata = result[0]["metadata"] + assert metadata["memory"] is None + assert metadata["ref_id"] == "abcdef12" + + +class TestPrepareReferenceDataTextualMemoryItem: + """The TextualMemoryItem branch keeps its original contract.""" + + def test_textual_memory_item_branch_unchanged(self): + item = TextualMemoryItem( + id="abcdef12-3456-7890-abcd-ef1234567890", + memory="hello world", + ) + + result = prepare_reference_data([item]) + + metadata = result[0]["metadata"] + assert metadata["ref_id"] == "abcdef12" + assert metadata["id"] == "abcdef12-3456-7890-abcd-ef1234567890" + assert metadata["memory"] == "hello world" + assert metadata["embedding"] == [] + assert metadata["sources"] == [] + + +class TestPrepareReferenceDataOldBehaviorDemonstration: + """ + Pre-fix demonstration tests kept for auditability. + + These document the exact error signatures reported in #2448 and are + negated by the assertions above. They are xfailed (strict=False) so they + do NOT gate CI — but if the fix is ever reverted, they'll flip to xpassed + (strict=False → not a failure) while the tests above will hard-fail, + which is the desired signal. + """ + + @pytest.mark.parametrize( + "entry", + [ + {"metadata": {"memory": "m"}}, # missing id + {"id": 12345, "memory": "m", "metadata": {}}, # int id + ], + ) + def test_pre_fix_would_have_raised(self, entry): + """Before the fix, both cases raised. After the fix, both succeed.""" + # After the fix this returns normally; the test simply verifies the + # call no longer raises the two exceptions from the issue. + result = prepare_reference_data([entry]) + assert isinstance(result, list) and len(result) == 1 From eaa5d00ba8cb601335169b0129814d153dbb12cf Mon Sep 17 00:00:00 2001 From: AutoDev Bot Date: Thu, 1 Oct 2026 00:30:59 +0800 Subject: [PATCH 2/2] fix: address OCR review findings on prepare_reference_data (#2448) OCR findings resolved: 1. src/memos/mem_os/utils/reference_utils.py: the dict-entry branch aliased ``memories_json = memories`` and then wrote ``ref_id``, ``embedding``, ``sources``, ``memory`` and ``id`` into the caller's ``metadata`` dict in place. Cached payloads (search cache, MCP messages reused across calls) would be silently corrupted. Now we shallow-copy the outer entry and the metadata dict before mutating. 2. tests/mem_os/utils/test_reference_utils.py: the class docstring documented ``test_pre_fix_would_have_raised`` as ``xfail(strict=False)`` but the decorator was missing, so the test ran as a plain passing test and the "pre-fix demonstration" gating story was silently lost. The decorator is now applied and the docstring rewritten to reflect the actual XPASS/XFAIL flip semantics. Also adds ``test_caller_dict_is_not_mutated`` as a regression guard for finding 1 (the pre-existing tests could not catch the aliasing bug). Verification: - pytest tests/mem_os/utils/test_reference_utils.py -v 8 passed, 2 xpassed (documented XPASS from the pre-fix demo cases) - ruff check src/memos/mem_os/utils/reference_utils.py tests/mem_os/utils/test_reference_utils.py All checks passed --- src/memos/mem_os/utils/reference_utils.py | 12 +++--- tests/mem_os/utils/test_reference_utils.py | 46 ++++++++++++++++++++-- 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/src/memos/mem_os/utils/reference_utils.py b/src/memos/mem_os/utils/reference_utils.py index 4c225ad8f..44ff97fdd 100644 --- a/src/memos/mem_os/utils/reference_utils.py +++ b/src/memos/mem_os/utils/reference_utils.py @@ -158,11 +158,13 @@ def prepare_reference_data(memories_list: list[TextualMemoryItem]) -> list[dict] memories_json["metadata"]["id"] = memories.id reference.append({"metadata": memories_json["metadata"]}) else: - memories_json = memories - metadata = memories_json.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - memories_json["metadata"] = metadata + # Shallow-copy the outer dict AND the metadata dict so that our + # normalization does not mutate the caller's payload (e.g. a cache + # entry or an MCP message being reused across calls). See #2448. + memories_json = dict(memories) + orig_metadata = memories_json.get("metadata") + metadata = dict(orig_metadata) if isinstance(orig_metadata, dict) else {} + memories_json["metadata"] = metadata raw_id = memories_json.get("id") if raw_id is None: diff --git a/tests/mem_os/utils/test_reference_utils.py b/tests/mem_os/utils/test_reference_utils.py index b625a5124..fae8983cf 100644 --- a/tests/mem_os/utils/test_reference_utils.py +++ b/tests/mem_os/utils/test_reference_utils.py @@ -103,6 +103,35 @@ def test_missing_memory_falls_back_to_none(self): assert metadata["memory"] is None assert metadata["ref_id"] == "abcdef12" + def test_caller_dict_is_not_mutated(self): + """The dict branch must not mutate the caller's entry or its metadata. + + Regression guard: earlier versions did + ``memories_json = memories`` (alias, not copy) and then wrote + ``metadata["ref_id"] = ...``. If a caller cached the payload or reused + it across MCP calls, those side-effects would silently corrupt it. + """ + original_metadata = {"source": "cache"} + entry = { + "id": "abcdef12-3456-7890-abcd-ef1234567890", + "memory": "hello", + "metadata": original_metadata, + } + # Snapshot the caller-visible shape so we can compare after the call. + entry_snapshot = { + "id": entry["id"], + "memory": entry["memory"], + "metadata": dict(original_metadata), + } + + prepare_reference_data([entry]) + + # Caller's outer dict unchanged (no injected ref_id/embedding/sources/id). + assert entry == entry_snapshot + # Caller's inner metadata dict unchanged (same object, same keys). + assert entry["metadata"] is original_metadata + assert original_metadata == {"source": "cache"} + class TestPrepareReferenceDataTextualMemoryItem: """The TextualMemoryItem branch keeps its original contract.""" @@ -128,10 +157,11 @@ class TestPrepareReferenceDataOldBehaviorDemonstration: Pre-fix demonstration tests kept for auditability. These document the exact error signatures reported in #2448 and are - negated by the assertions above. They are xfailed (strict=False) so they - do NOT gate CI — but if the fix is ever reverted, they'll flip to xpassed - (strict=False → not a failure) while the tests above will hard-fail, - which is the desired signal. + xfailed (strict=False) so they do NOT gate CI — after the fix the call + returns normally and pytest reports XPASS (non-fatal); if the fix is ever + reverted the call raises again and pytest reports XFAIL, again non-fatal. + Either way the ``TestPrepareReferenceDataDictEntries`` cases above are + the hard gate: they will fail loudly if the fix is reverted. """ @pytest.mark.parametrize( @@ -141,6 +171,14 @@ class TestPrepareReferenceDataOldBehaviorDemonstration: {"id": 12345, "memory": "m", "metadata": {}}, # int id ], ) + @pytest.mark.xfail( + strict=False, + reason=( + "Documents pre-fix error signatures from #2448. Passes after the " + "fix (XPASS, strict=False so it does not gate CI); if the fix is " + "ever reverted the call will raise again and this test will XFAIL." + ), + ) def test_pre_fix_would_have_raised(self, entry): """Before the fix, both cases raised. After the fix, both succeed.""" # After the fix this returns normally; the test simply verifies the