Merge f3c07e7569 into 99ecdd24d1

reme/steps/evolve/auto_memory.py

@@ -147,6 +147,35 @@ def _sanitize_msg_for_save(msg: Msg) -> Msg:
     return msg.model_copy(update={"content": new_content})
 
 
+def _canonical_message_json(msg: Msg) -> str:
+    """Compare the JSON values written to disk without Python bool/int aliasing or key-order noise."""
+    persisted = json.loads(msg.model_dump_json())
+    return json.dumps(persisted, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+
+
+def _stabilize_raw_content_ids(item: dict) -> dict:
+    """Give raw content blocks stable IDs when the caller did not supply one."""
+    content = item.get("content")
+    if not isinstance(content, list):
+        return item
+
+    stable_content = []
+    changed = False
+    for index, raw_block in enumerate(content):
+        if not isinstance(raw_block, dict):
+            stable_content.append(raw_block)
+            continue
+        block = dict(raw_block)
+        if "id" not in block:
+            canonical = json.dumps(block, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)
+            block["id"] = hashlib.sha256(f"{index}:{canonical}".encode("utf-8")).hexdigest()
+        if "created_at" not in block and item.get("created_at"):
+            block["created_at"] = item["created_at"]
+        stable_content.append(block)
+        changed = changed or block != raw_block
+    return {**item, "content": stable_content} if changed else item
+
+
 def _normalize_msg_timestamp(item: dict) -> dict:
     """Map common message timestamp aliases to AgentScope's ``created_at`` field."""
     if item.get("created_at"):
@@ -320,10 +349,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))
+                _canonical_message_json(merged[i]) == _canonical_message_json(existing[i]) for i in range(len(existing))
             )
 
             path.parent.mkdir(parents=True, exist_ok=True)
@@ -333,7 +365,7 @@ class AutoMemoryStep(BaseStep):
                 if new_msgs:
                     async with aiofiles.open(path, "a", encoding="utf-8", newline="") 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 +378,7 @@ class AutoMemoryStep(BaseStep):
             else:
                 async with aiofiles.open(path, "w", encoding="utf-8", newline="") 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)}",
@@ -360,6 +392,8 @@ class AutoMemoryStep(BaseStep):
             item = _normalize_msg_timestamp(item)
         if isinstance(item, dict) and isinstance(item.get("content"), str):
             item = {**item, "content": [{"type": "text", "text": item["content"]}]}
+        if isinstance(item, dict):
+            item = _stabilize_raw_content_ids(item)
         return Msg.model_validate(item)
 
     @staticmethod

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]

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,131 @@ 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 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"]
+    assert len(message.content) == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("append_new", [False, True])
+async def test_metadata_bool_to_int_correction_is_persisted(step, append_new):
+    """JSON true and 1 differ on disk even though Python considers them equal."""
+    original = _msg("m0", "2026-10-01T10:00:00").model_copy(update={"metadata": {"value": True}})
+    await step._save_session_messages(_SESSION, [original])
+    corrected = original.model_copy(update={"metadata": {"value": 1}})
+    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 saved[0].metadata["value"] == 1
+    assert type(saved[0].metadata["value"]) is int
+    assert [message.id for message in saved] == [message.id for message in messages]
+
+
+@pytest.mark.asyncio
+async def test_identical_raw_retry_with_generated_content_ids_does_not_rewrite(step, monkeypatch):
+    """Parsing the same raw messages twice should not create false block-ID changes."""
+    raw = {
+        "id": "m0",
+        "name": "Alice",
+        "role": "user",
+        "created_at": "2026-10-01T10:00:00",
+        "content": [{"type": "text", "text": "same raw payload"}],
+    }
+    first = step._build_messages([raw])
+    await step._save_session_messages(_SESSION, first)
+    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)
+    second = step._build_messages([raw])
+    await step._save_session_messages(_SESSION, second)
+
+    assert not write_modes
+    assert step._session_path(_SESSION).read_bytes() == original
+
+
+def test_raw_message_conversion_preserves_explicit_content_block_ids(step):
+    """Caller-provided block IDs remain authoritative during normalization."""
+    message = step._to_msg(
+        {
+            "id": "m0",
+            "name": "Alice",
+            "role": "user",
+            "created_at": "2026-10-01T10:00:00",
+            "content": [{"type": "text", "text": "same raw payload", "id": "caller-block-id"}],
+        },
+    )
+
+    assert message.content[0].id == "caller-block-id"
+
+
+@pytest.mark.asyncio
+async def test_identical_metadata_with_different_key_order_does_not_rewrite(step, monkeypatch):
+    """Canonical JSON comparison ignores object key order while preserving value types."""
+    original = _msg("m0", "2026-10-01T10:00:00").model_copy(update={"metadata": {"a": 1, "b": 2}})
+    await step._save_session_messages(_SESSION, [original])
+    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)
+    reordered = original.model_copy(update={"metadata": {"b": 2, "a": 1}})
+    await step._save_session_messages(_SESSION, [reordered])
+
+    assert not write_modes