Serialize tag operations, vacuum, and metadata refresh with writes

haiku_rag_slim/haiku/rag/client/documents.py

@@ -395,6 +395,7 @@ async def _refresh_doc_metadata(
         updated = True
 
     if updated:
+        async with client.store._write_lock:
             result = await client.document_repository.update_meta(doc)
         # Reclaim the document_meta churn from rolling source_revision sweeps.
         # The vacuum is debounced, and document_meta is tiny, so this is cheap.

haiku_rag_slim/haiku/rag/store/engine.py

@@ -520,7 +520,7 @@ class Store:
         if self._vacuum_lock.locked():
             return
 
-        async with self._vacuum_lock:
+        async with self._vacuum_lock, self._write_lock:
             try:
                 # Evaluate config at runtime to allow dynamic changes
                 if retention_seconds is None:
@@ -866,6 +866,9 @@ class Store:
     async def create_tag(self, name: str) -> None:
         """Tag the current version of every table with the given name.
 
+        Serializes with client writes via the write lock so a write cannot
+        land between the version snapshot and the per-table tag creation.
+
         Raises:
             ReadOnlyError: If the store is in read-only mode.
             ValueError: If the tag already exists on any table. A partial tag
@@ -875,6 +878,7 @@ class Store:
         self._assert_writable()
         tables = self._tables()
 
+        async with self._write_lock:
             existing = [
                 table_name
                 for table_name, table in tables.items()
@@ -920,11 +924,14 @@ class Store:
     async def delete_tag(self, name: str) -> None:
         """Delete the tag from every table that has it.
 
+        Serializes with create_tag and client writes via the write lock.
+
         Raises:
             ReadOnlyError: If the store is in read-only mode.
             ValueError: If no table has the tag.
         """
         self._assert_writable()
+        async with self._write_lock:
             found = False
             for table in self._tables().values():
                 if name in await table.tags.list():

tests/store/test_tags.py

@@ -85,6 +85,37 @@ async def test_create_tag_rolls_back_own_tags_on_failure(temp_db_path, monkeypat
         assert tags["keep"].complete is True
 
 
+@pytest.mark.asyncio
+async def test_create_tag_waits_for_write_lock(temp_db_path):
+    """create_tag serializes with client writes so a write cannot land
+    between the version snapshot and the per-table tag creation."""
+    async with Store(temp_db_path, create=True) as store:
+        async with store._write_lock:
+            task = asyncio.create_task(store.create_tag("release-1"))
+            await asyncio.sleep(0.1)
+            assert not task.done()
+        await task
+
+        tags = await store.list_tags()
+        assert tags["release-1"].complete is True
+
+
+@pytest.mark.asyncio
+async def test_delete_tag_waits_for_write_lock(temp_db_path):
+    """delete_tag serializes with create_tag and client writes so it cannot
+    remove tags out from under a concurrent create_tag."""
+    async with Store(temp_db_path, create=True) as store:
+        await store.create_tag("release-1")
+
+        async with store._write_lock:
+            task = asyncio.create_task(store.delete_tag("release-1"))
+            await asyncio.sleep(0.1)
+            assert not task.done()
+        await task
+
+        assert await store.list_tags() == {}
+
+
 @pytest.mark.asyncio
 async def test_delete_tag(temp_db_path):
     async with Store(temp_db_path, create=True) as store:
@@ -144,6 +175,18 @@ async def test_vacuum_cleans_untagged_versions_and_keeps_tagged(temp_db_path):
         assert rows == 2
 
 
+@pytest.mark.asyncio
+async def test_vacuum_waits_for_write_lock(temp_db_path):
+    """Vacuum serializes with writers and tag operations so a tag cannot be
+    created between _tag_safe_retention's read and the optimize call."""
+    async with Store(temp_db_path, create=True) as store:
+        async with store._write_lock:
+            task = asyncio.create_task(store.vacuum(retention_seconds=0))
+            await asyncio.sleep(0.1)
+            assert not task.done()
+        await task
+
+
 @pytest.mark.asyncio
 async def test_tag_writes_raise_when_read_only(temp_db_path):
     async with Store(temp_db_path, create=True) as store:

tests/test_vacuum_debounce.py

@@ -92,3 +92,33 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path):
             source_metadata={"source_revision": "r2", "md5": "same"},
         )
         assert client._vacuum_dirty is True
+
+
+@pytest.mark.asyncio
+async def test_metadata_refresh_waits_for_write_lock(temp_db_path):
+    """The revision/MD5 short-circuit write serializes with other writers so
+    it cannot land inside another writer's critical section (e.g. between
+    create_tag's version snapshot and its per-table tag creation)."""
+    dim = Config.embeddings.model.vector_dim
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(
+            _docling_doc("d", "body"),
+            [Chunk(content="body", embedding=[0.1] * dim, order=0)],
+            uri="mem://sweep",
+            metadata={"source_revision": "r1"},
+        )
+
+        async with client.store._write_lock:
+            task = asyncio.create_task(
+                _refresh_doc_metadata(
+                    client,
+                    doc,
+                    title=None,
+                    user_metadata={},
+                    source_metadata={"source_revision": "r2", "md5": "same"},
+                )
+            )
+            await asyncio.sleep(0.1)
+            assert not task.done()
+        refreshed = await task
+        assert refreshed.metadata["source_revision"] == "r2"