fix: atomic consistency in commit method

Author: scream4ik

examples/rag_hook_demo.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Any, Dict, List
+from typing import Dict, List
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -45,11 +45,11 @@ def __init__(self, vector_db: MockVectorDB):
         # We are only interested in facts of the "knowledge_base" type.
         self.target_types = {"knowledge_base", "chat_log"}
 
-    def __call__(self, op: str, fact_id: str, data: Dict[str, Any] | None):
+    def __call__(self, op: str, fact_id: str, data: Fact | None):
         # data is the state of the fact. If DELETE, this is the state BEFORE deletion.
 
         # Type checking (do not vectorize system data)
-        if not data or data.get("type") not in self.target_types:
+        if not data or data.type not in self.target_types:
             return
 
         # Processing of deletion
@@ -58,7 +58,7 @@ def __call__(self, op: str, fact_id: str, data: Dict[str, Any] | None):
             return
 
         # Text extraction for vectorization
-        payload = data.get("payload", {})
+        payload = data.payload
         # Trying to find a text field
         text_content = payload.get("content") or payload.get("message") or payload.get("summary")
 

memstate/storage.py

@@ -68,14 +68,17 @@ def commit(
         ephemeral: bool = False,
         actor: str | None = None,
         reason: str | None = None,
-    ):
+    ) -> str:
         with self._lock:
             validated_payload = self._schema_registry.validate(fact.type, fact.payload)
             fact.payload = validated_payload
 
             if session_id:
                 fact.session_id = session_id
 
+            previous_state = None
+            op = Operation.COMMIT
+
             constraint = self._constraints.get(fact.type)
 
             if constraint and constraint.singleton_key:
@@ -86,26 +89,37 @@ def commit(
 
                     if matches:
                         existing_raw = matches[0]
-                        existing_id = existing_raw["id"]
-
                         if constraint.immutable:
                             raise ConflictError(f"Immutable constraint violation: {fact.type}:{key_val}")
 
-                        before = copy.deepcopy(existing_raw)
-                        fact.id = existing_id
-                        after = fact.model_dump()
+                        # We found a duplicate, so this is an UPDATE of an existing one
+                        previous_state = copy.deepcopy(existing_raw)
+                        fact.id = existing_raw["id"]  # We replace the ID of the new fact with the old one
+                        op = Operation.UPDATE
+
+            if op != Operation.UPDATE:
+                existing = self.storage.load(fact.id)
+                if existing:
+                    previous_state = copy.deepcopy(existing)
+                    op = Operation.UPDATE
+                else:
+                    op = Operation.COMMIT_EPHEMERAL if ephemeral else Operation.COMMIT
+
+            try:
+                new_state = fact.model_dump()
+                self.storage.save(new_state)
+                self._log_tx(op, fact.id, previous_state, new_state, actor, reason)
+                self._notify_hooks(op, fact.id, fact)
 
-                        self.storage.save(after)
-                        self._log_tx(Operation.UPDATE, existing_id, before, after, actor, reason)
-                        return existing_id
+                return fact.id
 
-            existing = self.storage.load(fact.id)
-            op = Operation.UPDATE if existing else (Operation.COMMIT_EPHEMERAL if ephemeral else Operation.COMMIT)
+            except HookError as e:
+                if op == Operation.UPDATE and previous_state:
+                    self.storage.save(previous_state)
+                else:
+                    self.storage.delete(fact.id)
 
-            self.storage.save(fact.model_dump())
-            self._log_tx(op, fact.id, existing, fact.model_dump(), actor, reason)
-            self._notify_hooks(op, fact.id, fact)
-            return fact.id
+                raise e
 
     def update(self, fact_id: str, patch: dict[str, Any], actor: str | None = None, reason: str | None = None) -> str:
         with self._lock:

tests/test_storage.py

@@ -109,8 +109,7 @@ def crashing_hook(op, fid, data):
         memory.commit(Fact(type="user", payload={"name": "Survivor", "age": 50}))
 
     facts = memory.query(filters={"payload.name": "Survivor"})
-    assert len(facts) == 1
-    assert facts[0]["payload"]["name"] == "Survivor"
+    assert len(facts) == 0
 
 
 def test_ephemeral_session_discard(memory):