From 38a8f45d55a20a9e3d5f15a7e53e6c414e90398b Mon Sep 17 00:00:00 2001 From: adarshh347 Date: Sat, 24 Oct 2026 12:01:43 +0530 Subject: [PATCH 1/2] fix(auto-memory): persist same-ID transcript corrections --- reme/steps/evolve/auto_memory.py | 11 +++-- tests/unit/test_auto_memory_direct_images.py | 2 +- .../unit/test_auto_memory_session_messages.py | 47 ++++++++++++++++++- 3 files changed, 54 insertions(+), 6 deletions(-) diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 5e513fab..eeec48ba 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -320,10 +320,13 @@ class AutoMemoryStep(BaseStep): by_id[msg.id] = msg for msg in messages: by_id[msg.id] = msg - merged = sorted(by_id.values(), key=lambda m: m.created_at) + merged = [_sanitize_msg_for_save(msg) for msg in sorted(by_id.values(), key=lambda m: m.created_at)] + # An unchanged ID does not imply unchanged persisted content. Compare + # sanitized records so identical retries with omitted image/tool data + # retain the append fast path, while corrections force a rewrite. can_append = 0 < len(existing) <= len(merged) and all( - merged[i].id == existing[i].id for i in range(len(existing)) + merged[i].model_dump() == existing[i].model_dump() for i in range(len(existing)) ) path.parent.mkdir(parents=True, exist_ok=True) @@ -333,7 +336,7 @@ class AutoMemoryStep(BaseStep): if new_msgs: async with aiofiles.open(path, "a", encoding="utf-8") as f: for msg in new_msgs: - await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n") + await f.write(msg.model_dump_json() + "\n") self.logger.info( f"[{self.name}] save session appended session_id={session_id!r} " f"existing={len(existing)} appended={len(new_msgs)} total={len(merged)}", @@ -346,7 +349,7 @@ class AutoMemoryStep(BaseStep): else: async with aiofiles.open(path, "w", encoding="utf-8") as f: for msg in merged: - await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n") + await f.write(msg.model_dump_json() + "\n") self.logger.info( f"[{self.name}] save session rewrote session_id={session_id!r} " f"existing={len(existing)} total={len(merged)}", diff --git a/tests/unit/test_auto_memory_direct_images.py b/tests/unit/test_auto_memory_direct_images.py index 2b63af16..e53b483f 100644 --- a/tests/unit/test_auto_memory_direct_images.py +++ b/tests/unit/test_auto_memory_direct_images.py @@ -414,7 +414,7 @@ async def test_history_merge_keeps_main_source_contract(setup, enabled, history) elif history == "backfill": messages, expected = [older, original], _saved_line(older) + initial elif history == "same-id": - messages, expected = [replacement], initial + messages, expected = [replacement], _saved_line(replacement) else: messages, expected = [older, replacement], _saved_line(older) + _saved_line(replacement) before = [message.model_dump() for message in messages] diff --git a/tests/unit/test_auto_memory_session_messages.py b/tests/unit/test_auto_memory_session_messages.py index ac4b6398..fd1bb432 100644 --- a/tests/unit/test_auto_memory_session_messages.py +++ b/tests/unit/test_auto_memory_session_messages.py @@ -5,7 +5,7 @@ import asyncio from types import SimpleNamespace -from agentscope.message import Msg, TextBlock +from agentscope.message import Base64Source, DataBlock, Msg, TextBlock import pytest from reme.steps.evolve import auto_memory @@ -197,3 +197,48 @@ async def test_concurrent_append_survives_rewrite_race(step, monkeypatch): _assert_transcript(step._session_path(_SESSION), ["m0", "m_mid", "m1", "m2"]) _assert_serialized(interleaving) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("append_new", [False, True]) +@pytest.mark.parametrize("field", ["content", "metadata", "created_at"]) +async def test_existing_message_updates_are_persisted(step, append_new, field): + original = _msg("m0", "2026-10-01T10:00:00") + await step._save_session_messages(_SESSION, [original]) + updates = { + "content": [TextBlock(text="corrected observation")], + "metadata": {"corrected": True}, + "created_at": "2026-10-01T10:00:30", + } + corrected = original.model_copy(update={field: updates[field]}) + messages = [corrected] + if append_new: + messages.append(_msg("m1", "2026-10-01T10:01:00")) + + await step._save_session_messages(_SESSION, messages) + + saved = [Msg.model_validate_json(line) for line in step._session_path(_SESSION).read_text().splitlines()] + assert [message.model_dump() for message in saved] == [message.model_dump() for message in messages] + + +@pytest.mark.asyncio +async def test_identical_retry_with_sanitized_content_does_not_rewrite(step, monkeypatch): + message = _msg("m0", "2026-10-01T10:00:00") + message.content.append(DataBlock(source=Base64Source(media_type="image/png", data="aW1hZ2U="))) + await step._save_session_messages(_SESSION, [message]) + original = step._session_path(_SESSION).read_bytes() + real_open = auto_memory.aiofiles.open + write_modes = [] + + def record_open(path, mode="r", **kwargs): + if mode in ("a", "w"): + write_modes.append(mode) + return real_open(path, mode, **kwargs) + + monkeypatch.setattr(auto_memory.aiofiles, "open", record_open) + await step._save_session_messages(_SESSION, [message]) + assert write_modes == [] + assert step._session_path(_SESSION).read_bytes() == original + await step._save_session_messages(_SESSION, [message, _msg("m1", "2026-10-01T10:01:00")]) + assert write_modes == ["a"] + assert len(message.content) == 2 From 06e32cb4dec5b670bac53e83e30e9b5043a93ee0 Mon Sep 17 00:00:00 2001 From: adarshh347 Date: Sun, 25 Oct 2026 10:44:24 +0530 Subject: [PATCH 2/2] test(auto-memory): satisfy sequence emptiness lint --- tests/unit/test_auto_memory_session_messages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_auto_memory_session_messages.py b/tests/unit/test_auto_memory_session_messages.py index fd1bb432..778f9d58 100644 --- a/tests/unit/test_auto_memory_session_messages.py +++ b/tests/unit/test_auto_memory_session_messages.py @@ -237,7 +237,7 @@ async def test_identical_retry_with_sanitized_content_does_not_rewrite(step, mon monkeypatch.setattr(auto_memory.aiofiles, "open", record_open) await step._save_session_messages(_SESSION, [message]) - assert write_modes == [] + assert not write_modes assert step._session_path(_SESSION).read_bytes() == original await step._save_session_messages(_SESSION, [message, _msg("m1", "2026-10-01T10:01:00")]) assert write_modes == ["a"]