diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 431785c3..630c626a 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -147,6 +147,47 @@ 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_content_block(raw_block: dict, path: tuple[int, ...], message_timestamp: str | None) -> dict: + """Stabilize defaults generated while parsing a raw content block.""" + block = dict(raw_block) + if isinstance(block.get("hint"), list): + block["hint"] = [ + _stabilize_content_block(child, (*path, index), message_timestamp) if isinstance(child, dict) else child + for index, child in enumerate(block["hint"]) + ] + + if "id" not in block: + canonical = json.dumps(raw_block, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str) + position = str(path[0]) if len(path) == 1 else ".".join(map(str, path)) + block["id"] = hashlib.sha256(f"{position}:{canonical}".encode("utf-8")).hexdigest() + if "created_at" not in block and message_timestamp: + block["created_at"] = message_timestamp + if block.get("type") == "hint" and "finished_at" not in block and message_timestamp: + block["finished_at"] = message_timestamp + return block + + +def _stabilize_raw_content_ids(item: dict) -> dict: + """Stabilize generated IDs and timestamps in raw (including nested hint) blocks.""" + content = item.get("content") + if not isinstance(content, list): + return item + + timestamp = item.get("created_at") + stable_content = [ + _stabilize_content_block(block, (index,), timestamp) if isinstance(block, dict) else block + for index, block in enumerate(content) + ] + changed = stable_content != content + 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 +361,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 +377,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 +390,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 +404,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 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..6d6244b6 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,183 @@ 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 +@pytest.mark.parametrize( + "content", + [ + [{"type": "text", "text": "same raw payload"}], + [{"type": "hint", "hint": "notification"}], + [ + { + "type": "hint", + "hint": [{"type": "text", "text": "notification"}], + "finished_at": "2026-10-01T10:00:00", + }, + ], + ], + ids=["text", "hint", "nested hint"], +) +async def test_identical_raw_retry_with_generated_content_defaults_does_not_rewrite(step, monkeypatch, content): + """Parsing identical raw messages twice should not churn generated block defaults.""" + raw = { + "id": "m0", + "name": "Assistant", + "role": "assistant", + "created_at": "2026-10-01T10:00:00", + "content": content, + } + 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_hint_conversion_preserves_explicit_nested_ids_and_timestamps(step): + """Caller-provided IDs and timestamps remain authoritative at every hint depth.""" + message = step._to_msg( + { + "id": "m0", + "name": "Alice", + "role": "assistant", + "created_at": "2026-10-01T10:00:00", + "content": [ + { + "type": "hint", + "id": "outer-id", + "created_at": "2026-10-01T10:00:01", + "finished_at": "2026-10-01T10:00:02", + "hint": [ + { + "type": "text", + "text": "notification", + "id": "nested-id", + "created_at": "2026-10-01T10:00:03", + }, + ], + }, + ], + }, + ) + + outer = message.content[0] + nested = outer.hint[0] + assert (outer.id, outer.created_at, outer.finished_at) == ( + "outer-id", + "2026-10-01T10:00:01", + "2026-10-01T10:00:02", + ) + assert (nested.id, nested.created_at) == ("nested-id", "2026-10-01T10:00:03") + + +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