Merge pull request #44 from beersoccer/fix/checkpoint-lock-reliability

Fix/checkpoint lock reliability

Author: beersoccer

tests/unit/tools/test_forget_memories.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import json
+from datetime import UTC, datetime, timedelta
 from unittest.mock import MagicMock, patch
 
-from tools.forget_memories import ForgetMemoriesTool
+from tools.forget_memories import ForgetMemoriesTool, _clean_expired_locks
 
 
 def _extract_json_payload(message: object) -> dict:
@@ -145,3 +146,124 @@ def test_forget_memories_deletes_old_checkpoints_in_real_run() -> None:
     payload = _extract_json_payload(messages[0])
     assert payload.get("status") == "SUCCESS"
     assert payload.get("results", {}).get("checkpoints_cleaned") == 2
+
+
+def _make_lock_memory(
+    mem_id: str,
+    holder_id: str,
+    acquired_at: str,
+    ttl_seconds: int,
+) -> dict:
+    """Build a fake lock memory entry for testing."""
+    lock_data = {
+        "lock_id": "lock:extraction:u1:*",
+        "holder_id": holder_id,
+        "acquired_at": acquired_at,
+        "ttl_seconds": ttl_seconds,
+    }
+    return {
+        "id": mem_id,
+        "memory": json.dumps(lock_data, ensure_ascii=False),
+        "metadata": {
+            "__internal": True,
+            "internal_type": "distributed_lock",
+        },
+    }
+
+
+def test_clean_expired_locks_deletes_expired_only() -> None:
+    """_clean_expired_locks should delete expired locks and keep active ones."""
+    now = datetime.now(UTC)
+    expired_time = (now - timedelta(seconds=7200)).isoformat()
+    active_time = now.isoformat()
+
+    mem = MagicMock()
+    mem.get_all.return_value = {
+        "results": [
+            _make_lock_memory("lock-1", "run-old", expired_time, 3600),
+            _make_lock_memory("lock-2", "run-active", active_time, 3600),
+        ]
+    }
+
+    count = _clean_expired_locks(
+        mem, user_id="u1", app_id=None, dry_run=False, request_id="req-1"
+    )
+
+    assert count == 1
+    mem.delete.assert_called_once_with("lock-1")
+
+
+def test_clean_expired_locks_dry_run_does_not_delete() -> None:
+    """Dry run should count but not delete."""
+    expired_time = (datetime.now(UTC) - timedelta(seconds=7200)).isoformat()
+
+    mem = MagicMock()
+    mem.get_all.return_value = {
+        "results": [
+            _make_lock_memory("lock-1", "run-old", expired_time, 3600),
+        ]
+    }
+
+    count = _clean_expired_locks(
+        mem, user_id="u1", app_id=None, dry_run=True, request_id="req-1"
+    )
+
+    assert count == 1
+    mem.delete.assert_not_called()
+
+
+def test_clean_expired_locks_corrupted_record() -> None:
+    """Corrupted lock records (unparseable JSON) should be cleaned up."""
+    mem = MagicMock()
+    mem.get_all.return_value = {
+        "results": [
+            {
+                "id": "lock-corrupt",
+                "memory": "not-valid-json",
+                "metadata": {
+                    "__internal": True,
+                    "internal_type": "distributed_lock",
+                },
+            }
+        ]
+    }
+
+    count = _clean_expired_locks(
+        mem, user_id="u1", app_id=None, dry_run=False, request_id="req-1"
+    )
+
+    assert count == 1
+    mem.delete.assert_called_once_with("lock-corrupt")
+
+
+def test_forget_memories_includes_locks_cleaned_in_result() -> None:
+    """Integration: ForgetMemoriesTool result should include locks_cleaned."""
+    tool = _build_tool()
+    mem = MagicMock()
+
+    # Call 1: regular memories, Call 2: checkpoints, Call 3: locks
+    expired_time = (datetime.now(UTC) - timedelta(seconds=7200)).isoformat()
+    mem.get_all.side_effect = [
+        {"results": []},  # no regular memories
+        {"results": []},  # no checkpoints
+        {
+            "results": [
+                _make_lock_memory("lock-1", "old-run", expired_time, 3600),
+            ]
+        },  # one expired lock
+    ]
+    client = MagicMock(memory=mem)
+    mgr = MagicMock()
+    mgr.load.return_value = ("log-1", {})
+
+    with (
+        patch("tools.forget_memories.init_request_context", return_value=("req-1", 0.0)),
+        patch("tools.forget_memories.validate_user_id", return_value="u1"),
+        patch("tools.forget_memories.get_sync_client", return_value=client),
+        patch("tools.forget_memories.SyncAccessLogManager", return_value=mgr),
+    ):
+        messages = list(tool._invoke({"user_id": "u1", "dry_run": False}))
+
+    payload = _extract_json_payload(messages[0])
+    assert payload.get("status") == "SUCCESS"
+    assert payload.get("results", {}).get("locks_cleaned") == 1

tests/unit/utils/test_checkpoint.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
 
+import asyncio
+import threading
 from typing import Any
 
 import pytest
 
 from utils.checkpoint import (
     CHECKPOINT_VERSION,
+    AsyncCheckpointManager,
     SyncCheckpointManager,
     checkpoint_filters,
     checkpoint_metadata,
 )
-from utils.extraction import UserCheckpoint
+from utils.extraction import ConversationCheckpoint, UserCheckpoint
 
 
 class FakeMemory:
@@ -129,18 +132,232 @@ def test_save_checkpoint_add_and_update(monkeypatch: pytest.MonkeyPatch) -> None
     assert mem.added[0]["user_id"] == "u1"
     assert mem.added[0]["agent_id"] is None
 
-    # update existing (uses delete+add, not update)
-    ok2, same_id = mgr.save(
+    # update existing (uses add-first-then-delete)
+    ok2, new_id2 = mgr.save(
         checkpoint_id=new_id,
         user_id="u1",
         app_id=None,
         checkpoint=UserCheckpoint(),
     )
     assert ok2 is True
-    assert same_id == new_id
-    # save uses delete+add instead of update to avoid embedding
+    assert new_id2 is not None
+    # save adds new first, then deletes old
     assert new_id in mem.deleted  # Old checkpoint deleted
     assert len(mem.added) == 2  # New checkpoint added
 
 
-    
+def test_save_add_before_delete_order() -> None:
+    """Verify that save() adds the new checkpoint BEFORE deleting the old one."""
+    mem = FakeMemory()
+    mgr = SyncCheckpointManager(mem)
+
+    # Create initial checkpoint
+    ok, initial_id = mgr.save(
+        checkpoint_id=None,
+        user_id="u1",
+        app_id=None,
+        checkpoint=UserCheckpoint(),
+    )
+    assert ok and initial_id
+
+    # Track operation order
+    operations: list[str] = []
+    original_add = mem.add
+    original_delete = mem.delete
+
+    def tracking_add(text: str, **kwargs: Any) -> dict[str, Any]:
+        operations.append("add")
+        return original_add(text, **kwargs)
+
+    def tracking_delete(memory_id: str) -> dict[str, Any]:
+        operations.append("delete")
+        return original_delete(memory_id)
+
+    mem.add = tracking_add  # type: ignore[assignment]
+    mem.delete = tracking_delete  # type: ignore[assignment]
+
+    # Update checkpoint
+    ok2, new_id = mgr.save(
+        checkpoint_id=initial_id,
+        user_id="u1",
+        app_id=None,
+        checkpoint=UserCheckpoint(),
+    )
+    assert ok2 is True
+    # add must come before delete
+    assert operations == ["add", "delete"]
+
+
+def test_save_keeps_old_checkpoint_on_add_failure() -> None:
+    """If add() fails, old checkpoint should NOT be deleted."""
+    mem = FakeMemory()
+    mgr = SyncCheckpointManager(mem)
+
+    # Create initial checkpoint
+    ok, initial_id = mgr.save(
+        checkpoint_id=None,
+        user_id="u1",
+        app_id=None,
+        checkpoint=UserCheckpoint(),
+    )
+    assert ok and initial_id
+
+    # Make add return no ID (simulating failure)
+    mem.add = lambda text, **kwargs: {"results": []}  # type: ignore[assignment]
+
+    ok2, new_id = mgr.save(
+        checkpoint_id=initial_id,
+        user_id="u1",
+        app_id=None,
+        checkpoint=UserCheckpoint(),
+    )
+    # add returned no ID, so old checkpoint must NOT be deleted
+    assert initial_id not in mem.deleted
+
+
+class AsyncFakeMemory:
+    """Async-compatible fake Mem0 memory for testing."""
+
+    def __init__(self) -> None:
+        self._store: list[dict[str, Any]] = []
+        self.added: list[dict[str, Any]] = []
+        self.deleted: list[str] = []
+
+    def _match(self, md: dict[str, Any], filt: Any) -> bool:
+        if not filt or not isinstance(filt, dict):
+            return True
+        # Leaf filter: {"key": value} or {"key": {"eq": value}}
+        for k, v in filt.items():
+            if isinstance(v, dict) and "eq" in v:
+                if md.get(k) != v.get("eq"):
+                    return False
+            else:
+                if md.get(k) != v:
+                    return False
+        return True
+
+    async def get_all(self, **kwargs: Any) -> dict[str, Any]:
+        filt = kwargs.get("filters")
+        out: list[dict[str, Any]] = []
+        for item in self._store:
+            md = item.get("metadata") or {}
+            if self._match(md, filt):
+                out.append(item)
+        return {"results": out}
+
+    async def add(self, text: str, **kwargs: Any) -> dict[str, Any]:
+        md = kwargs.get("metadata") or {}
+        new_id = f"cp_{len(self._store)+1}"
+        self._store.append({"id": new_id, "memory": text, "metadata": md})
+        self.added.append({"id": new_id, "text": text})
+        return {"results": [{"id": new_id, "event": "ADD"}]}
+
+    async def delete(self, memory_id: str) -> dict[str, Any]:
+        self.deleted.append(memory_id)
+        self._store = [x for x in self._store if x.get("id") != memory_id]
+        return {"message": "deleted"}
+
+
+def _run_async(coro: Any) -> Any:
+    """Run a coroutine in a dedicated thread to avoid event loop conflicts with pytest-asyncio."""
+    result: list[Any] = []
+    exc: list[BaseException] = []
+
+    def _thread_target() -> None:
+        # Clear any inherited running-loop state from the parent thread.
+        # In pytest-asyncio AUTO mode the main thread runs inside an event
+        # loop, and CPython's C-level _get_running_loop() can inherit that
+        # state in child threads, causing run_until_complete() to raise
+        # "Cannot run the event loop while another loop is running".
+        asyncio.events._set_running_loop(None)  # type: ignore[attr-defined]
+        asyncio.set_event_loop(None)
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            result.append(loop.run_until_complete(coro))
+        except BaseException as e:  # noqa: BLE001
+            exc.append(e)
+        finally:
+            loop.close()
+            asyncio.set_event_loop(None)
+
+    t = threading.Thread(target=_thread_target)
+    t.start()
+    t.join()
+    if exc:
+        raise exc[0]
+    return result[0] if result else None
+
+
+def test_async_load_restores_resume_fields() -> None:
+    """AsyncCheckpointManager.load() must restore resume_* fields."""
+
+    async def _run() -> None:
+        mem = AsyncFakeMemory()
+        mgr = AsyncCheckpointManager(mem)
+
+        cp = UserCheckpoint(
+            conversations={
+                "conv1": ConversationCheckpoint(
+                    last_processed_message_id="msg-100",
+                    processed_range_start="2026-01-01T00:00:00Z",
+                    processed_range_end="2026-01-02T00:00:00Z",
+                ),
+            },
+            resume_conversation_cursor="conv1",
+            resume_run_at="2026-04-01T12:00:00Z",
+            resume_start_time="2026-03-01T00:00:00Z",
+        )
+
+        ok, cp_id = await mgr.save(
+            checkpoint_id=None, user_id="u1", app_id=None, checkpoint=cp
+        )
+        assert ok and cp_id
+
+        loaded_id, loaded_cp = await mgr.load(user_id="u1", app_id=None)
+        assert loaded_id == cp_id
+        assert loaded_cp is not None
+        assert loaded_cp.resume_conversation_cursor == "conv1"
+        assert loaded_cp.resume_run_at == "2026-04-01T12:00:00Z"
+        assert loaded_cp.resume_start_time == "2026-03-01T00:00:00Z"
+        assert "conv1" in loaded_cp.conversations
+        assert loaded_cp.conversations["conv1"].last_processed_message_id == "msg-100"
+
+    _run_async(_run())
+
+
+def test_async_save_add_before_delete() -> None:
+    """Async save() must add new checkpoint before deleting old."""
+
+    async def _run() -> None:
+        mem = AsyncFakeMemory()
+        mgr = AsyncCheckpointManager(mem)
+
+        ok, initial_id = await mgr.save(
+            checkpoint_id=None, user_id="u1", app_id=None, checkpoint=UserCheckpoint()
+        )
+        assert ok and initial_id
+
+        operations: list[str] = []
+        original_add = mem.add
+        original_delete = mem.delete
+
+        async def tracking_add(text: str, **kwargs: Any) -> dict[str, Any]:
+            operations.append("add")
+            return await original_add(text, **kwargs)
+
+        async def tracking_delete(memory_id: str) -> dict[str, Any]:
+            operations.append("delete")
+            return await original_delete(memory_id)
+
+        mem.add = tracking_add  # type: ignore[assignment]
+        mem.delete = tracking_delete  # type: ignore[assignment]
+
+        ok2, new_id = await mgr.save(
+            checkpoint_id=initial_id, user_id="u1", app_id=None, checkpoint=UserCheckpoint()
+        )
+        assert ok2 is True
+        assert operations == ["add", "delete"]
+
+    _run_async(_run())
+