From f3c07e7569d79d4213b80f2f104718fb5496d728 Mon Sep 17 00:00:00 2001 From: adarshh347 Date: Tue, 3 Nov 2026 10:15:41 +0530 Subject: [PATCH] Fix session transcript retry comparisons --- reme/steps/evolve/auto_memory.py | 33 +++++++- .../unit/test_auto_memory_session_messages.py | 83 +++++++++++++++++++ 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index eeec48ba..47ddcbe1 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/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"): @@ -326,7 +355,7 @@ class AutoMemoryStep(BaseStep): # 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].model_dump() == existing[i].model_dump() 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) @@ -363,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 diff --git a/tests/unit/test_auto_memory_session_messages.py b/tests/unit/test_auto_memory_session_messages.py index 778f9d58..6f5e87c7 100644 --- a/tests/unit/test_auto_memory_session_messages.py +++ b/tests/unit/test_auto_memory_session_messages.py @@ -242,3 +242,86 @@ async def test_identical_retry_with_sanitized_content_does_not_rewrite(step, mon 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