mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-23 05:35:30 +00:00
fix(auto-memory): serialize session transcript writes per session (#633)
* fix(auto-memory): serialize session transcript writes per session _save_session_messages ran its read-merge-write on the per-session JSONL transcript without any lock. Concurrent HTTP requests saving the same session could interleave: two writers both appending against a stale snapshot write overlapping history twice, and an append racing a full rewrite can be overwritten by the rewriter's stale merge, losing messages. Hold get_path_lock(path) across the whole read, merge, and write, consistent with the existing session image handling in the same file. Merge, ordering, and can_append logic are unchanged. Add deterministic regression tests that gate one writer's read between snapshot capture and write: concurrent appends must merge disjoint messages in created_at order, and an appended message must survive a concurrent full rewrite. Both tests fail without the lock. Closes #622 * test(auto-memory): make session-lock regression tests deterministic Replace fixed sleeps with explicit synchronization: a shared operation log records lock/read/write progress per writer, and writer A holds its post-read pause until writer B is queued on the transcript lock. The tests now assert the serialization order from the log (B reads only after A's write completed) in addition to transcript content, so they no longer depend on IO timing. Also fix pylint findings from CI: use from-import for the auto_memory module and drop unused tmp_path parameters.
This commit is contained in:
celanwang
GitHub
parent
47b1498615
commit
681cc85fb5
2 changed files with 240 additions and 40 deletions
|
|
@ -299,54 +299,55 @@ class AutoMemoryStep(BaseStep):
|
|||
f"[{self.name}] save session start session_id={session_id!r} messages={len(messages)} path={path}",
|
||||
)
|
||||
|
||||
existing: list[Msg] = []
|
||||
if path.exists():
|
||||
async with aiofiles.open(path, encoding="utf-8") as f:
|
||||
content = await f.read()
|
||||
for line in content.splitlines():
|
||||
line = line.strip()
|
||||
if line:
|
||||
try:
|
||||
existing.append(Msg.model_validate_json(line))
|
||||
except Exception:
|
||||
pass
|
||||
async with await get_path_lock(path):
|
||||
existing: list[Msg] = []
|
||||
if path.exists():
|
||||
async with aiofiles.open(path, encoding="utf-8") as f:
|
||||
content = await f.read()
|
||||
for line in content.splitlines():
|
||||
line = line.strip()
|
||||
if line:
|
||||
try:
|
||||
existing.append(Msg.model_validate_json(line))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
by_id: dict[str, Msg] = {}
|
||||
for msg in existing:
|
||||
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)
|
||||
by_id: dict[str, Msg] = {}
|
||||
for msg in existing:
|
||||
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)
|
||||
|
||||
can_append = 0 < len(existing) <= len(merged) and all(
|
||||
merged[i].id == existing[i].id for i in range(len(existing))
|
||||
)
|
||||
can_append = 0 < len(existing) <= len(merged) and all(
|
||||
merged[i].id == existing[i].id for i in range(len(existing))
|
||||
)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if can_append:
|
||||
new_msgs = merged[len(existing) :]
|
||||
if new_msgs:
|
||||
async with aiofiles.open(path, "a", encoding="utf-8") as f:
|
||||
for msg in new_msgs:
|
||||
if can_append:
|
||||
new_msgs = merged[len(existing) :]
|
||||
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")
|
||||
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)}",
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
f"[{self.name}] save session unchanged session_id={session_id!r} "
|
||||
f"existing={len(existing)} total={len(merged)}",
|
||||
)
|
||||
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")
|
||||
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)}",
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
f"[{self.name}] save session unchanged session_id={session_id!r} "
|
||||
f"[{self.name}] save session rewrote session_id={session_id!r} "
|
||||
f"existing={len(existing)} total={len(merged)}",
|
||||
)
|
||||
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")
|
||||
self.logger.info(
|
||||
f"[{self.name}] save session rewrote session_id={session_id!r} "
|
||||
f"existing={len(existing)} total={len(merged)}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _to_msg(item) -> Msg:
|
||||
|
|
|
|||
199
tests/unit/test_auto_memory_session_messages.py
Normal file
199
tests/unit/test_auto_memory_session_messages.py
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
"""Session transcript saves stay consistent under concurrent writers."""
|
||||
|
||||
# pylint: disable=protected-access,missing-function-docstring
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
from agentscope.message import Msg, TextBlock
|
||||
import pytest
|
||||
|
||||
from reme.steps.evolve import auto_memory
|
||||
from reme.steps.evolve.auto_memory import AutoMemoryStep
|
||||
|
||||
_SESSION = "race-session"
|
||||
|
||||
|
||||
def _msg(message_id: str, timestamp: str) -> Msg:
|
||||
return Msg.model_validate(
|
||||
{
|
||||
"id": message_id,
|
||||
"name": "Alice",
|
||||
"role": "user",
|
||||
"created_at": timestamp,
|
||||
"content": [TextBlock(text=f"text of {message_id}")],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(name="step")
|
||||
def memory_step(tmp_path):
|
||||
step = AutoMemoryStep()
|
||||
step.file_store = SimpleNamespace(workspace_path=tmp_path)
|
||||
return step
|
||||
|
||||
|
||||
class _LoggedLock:
|
||||
"""Record lock contention; writer B signals once it is queued on the lock."""
|
||||
|
||||
def __init__(self, lock, tag, interleaving):
|
||||
self._lock = lock
|
||||
self._tag = tag
|
||||
self._interleaving = interleaving
|
||||
|
||||
async def __aenter__(self):
|
||||
self._interleaving.log.append((self._tag, "lock-wait"))
|
||||
if self._tag == "B":
|
||||
self._interleaving.b_lock_waiting.set()
|
||||
await self._lock.acquire()
|
||||
self._interleaving.log.append((self._tag, "lock-acquired"))
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
self._lock.release()
|
||||
return None
|
||||
|
||||
|
||||
class _LoggedRead:
|
||||
"""Record read progress; writer A pauses after reading until B awaits the lock."""
|
||||
|
||||
def __init__(self, context_manager, tag, interleaving):
|
||||
self._context_manager = context_manager
|
||||
self._tag = tag
|
||||
self._interleaving = interleaving
|
||||
|
||||
async def __aenter__(self):
|
||||
handle = await self._context_manager.__aenter__()
|
||||
interleaving, tag = self._interleaving, self._tag
|
||||
interleaving.log.append((tag, "read-start"))
|
||||
|
||||
async def read():
|
||||
content = await handle.read()
|
||||
interleaving.log.append((tag, "read-done"))
|
||||
if tag == "A":
|
||||
await interleaving.b_lock_waiting.wait()
|
||||
return content
|
||||
|
||||
return SimpleNamespace(read=read)
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
return await self._context_manager.__aexit__(*exc_info)
|
||||
|
||||
|
||||
class _LoggedWrite:
|
||||
"""Record write completion once the file handle is closed and flushed."""
|
||||
|
||||
def __init__(self, context_manager, tag, interleaving):
|
||||
self._context_manager = context_manager
|
||||
self._tag = tag
|
||||
self._interleaving = interleaving
|
||||
|
||||
async def __aenter__(self):
|
||||
self._interleaving.log.append((self._tag, "write-start"))
|
||||
return await self._context_manager.__aenter__()
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
result = await self._context_manager.__aexit__(*exc_info)
|
||||
self._interleaving.log.append((self._tag, "write-done"))
|
||||
return result
|
||||
|
||||
|
||||
class _Interleaving:
|
||||
"""Deterministic two-writer interleaving driven by events, not sleeps.
|
||||
|
||||
Writer A snapshots the transcript, then pauses until writer B is queued on
|
||||
the transcript lock. Without mutual exclusion B would read and write against
|
||||
A's stale snapshot in between; with the lock B cannot start reading until A
|
||||
finished writing.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.log: list[tuple[str, str]] = []
|
||||
self.b_lock_waiting = asyncio.Event()
|
||||
self.task_a: asyncio.Task | None = None
|
||||
self.task_b: asyncio.Task | None = None
|
||||
|
||||
def patch(self, monkeypatch):
|
||||
real_get_path_lock = auto_memory.get_path_lock
|
||||
real_open = auto_memory.aiofiles.open
|
||||
|
||||
async def get_path_lock(target):
|
||||
return _LoggedLock(await real_get_path_lock(target), self._tag(), self)
|
||||
|
||||
def open_(path, mode="r", **kwargs):
|
||||
context_manager = real_open(path, mode, **kwargs)
|
||||
tag = self._tag()
|
||||
if mode == "r":
|
||||
return _LoggedRead(context_manager, tag, self)
|
||||
if mode in ("a", "w"):
|
||||
return _LoggedWrite(context_manager, tag, self)
|
||||
return context_manager
|
||||
|
||||
monkeypatch.setattr(auto_memory, "get_path_lock", get_path_lock)
|
||||
monkeypatch.setattr(auto_memory.aiofiles, "open", open_)
|
||||
|
||||
def _tag(self):
|
||||
task = asyncio.current_task()
|
||||
if task is self.task_a:
|
||||
return "A"
|
||||
if task is self.task_b:
|
||||
return "B"
|
||||
return None
|
||||
|
||||
async def run(self, step, a_messages, b_messages):
|
||||
self.task_a = asyncio.create_task(step._save_session_messages(_SESSION, a_messages))
|
||||
await asyncio.sleep(0) # Let writer A acquire the transcript lock before B starts.
|
||||
self.task_b = asyncio.create_task(step._save_session_messages(_SESSION, b_messages))
|
||||
await asyncio.wait_for(asyncio.gather(self.task_a, self.task_b), timeout=5)
|
||||
|
||||
|
||||
def _assert_transcript(path, expected_ids):
|
||||
lines = [line for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
|
||||
messages = [Msg.model_validate_json(line) for line in lines]
|
||||
assert [message.id for message in messages] == expected_ids
|
||||
assert len({message.id for message in messages}) == len(messages)
|
||||
assert [message.created_at for message in messages] == sorted(message.created_at for message in messages)
|
||||
|
||||
|
||||
def _assert_serialized(interleaving):
|
||||
log = interleaving.log
|
||||
# Writer B queued on the transcript lock while writer A held it: A's
|
||||
# post-read pause only ends once B attempted the lock, so B's lock-wait
|
||||
# must precede A's write.
|
||||
assert log.index(("B", "lock-wait")) < log.index(("A", "write-start"))
|
||||
# Mutual exclusion: B acquires the lock and re-reads only after A's write
|
||||
# completed, so B always merges against A's fresh transcript state.
|
||||
assert log.index(("A", "write-done")) < log.index(("B", "lock-acquired"))
|
||||
assert log.index(("B", "lock-acquired")) < log.index(("B", "read-start"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_appends_merge_disjoint_messages_in_order(step, monkeypatch):
|
||||
await step._save_session_messages(_SESSION, [_msg("m0", "2026-10-01T10:00:00")])
|
||||
|
||||
interleaving = _Interleaving()
|
||||
interleaving.patch(monkeypatch)
|
||||
# Writer A snapshots [m0]; writer B appends m2 while A is paused.
|
||||
await interleaving.run(step, [_msg("m1", "2026-10-01T10:01:00")], [_msg("m2", "2026-10-01T10:02:00")])
|
||||
|
||||
_assert_transcript(step._session_path(_SESSION), ["m0", "m1", "m2"])
|
||||
_assert_serialized(interleaving)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_append_survives_rewrite_race(step, monkeypatch):
|
||||
await step._save_session_messages(_SESSION, [_msg("m0", "2026-10-01T10:00:00")])
|
||||
await step._save_session_messages(_SESSION, [_msg("m1", "2026-10-01T10:01:00")])
|
||||
|
||||
interleaving = _Interleaving()
|
||||
interleaving.patch(monkeypatch)
|
||||
# Writer A inserts m_mid before m1, forcing the full-rewrite branch against
|
||||
# its stale snapshot; writer B appends m2 while A is paused.
|
||||
await interleaving.run(
|
||||
step,
|
||||
[_msg("m_mid", "2026-10-01T10:00:30")],
|
||||
[_msg("m2", "2026-10-01T10:02:00")],
|
||||
)
|
||||
|
||||
_assert_transcript(step._session_path(_SESSION), ["m0", "m_mid", "m1", "m2"])
|
||||
_assert_serialized(interleaving)
|
||||
Loading…
Add table
Reference in a new issue