Protect tag-creation cleanup from cancellation

haiku_rag_slim/haiku/rag/store/engine.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+from collections.abc import Coroutine
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
 from enum import Enum
@@ -200,6 +201,27 @@ RESTORE_TABLE_ORDER: tuple[str, ...] = tuple(
 ) + ("documents",)
 
 
+async def _wait_protected[T](coro: Coroutine[Any, Any, T]) -> tuple[T, bool]:
+    """Await a recovery coroutine that a cancellation cannot interrupt.
+
+    Runs the coroutine as a task and keeps waiting for it even if this
+    coroutine is cancelled, so a Ctrl-C cannot leave recovery half applied.
+    Returns the result and whether a cancellation was absorbed; the caller
+    must re-deliver an absorbed cancellation.
+    """
+    task = asyncio.ensure_future(coro)
+    cancelled = False
+    while True:
+        try:
+            return await asyncio.shield(task), cancelled
+        except asyncio.CancelledError:
+            if task.done():
+                # The recovery coroutine itself ended cancelled; there is
+                # nothing left to wait for.
+                raise
+            cancelled = True
+
+
 def _safety_tag_name(existing: set[str]) -> str:
     """Collision-resistant name for the pre-restore safety tag."""
     base = f"before-restore-{datetime.now(UTC).strftime('%Y%m%dT%H%M%SZ')}"
@@ -926,26 +948,47 @@ class Store:
             )
 
         versions = await self.current_table_versions()
-        created: list[str] = []
         try:
             for table_name, table in tables.items():
                 await table.tags.create(name, versions[table_name])
-                created.append(table_name)
+        except BaseException as exc:
-        except Exception as exc:
+            # BaseException: cancellation must also trigger cleanup, and the
-            failed_cleanup: list[str] = []
+            # cleanup itself is protected from further cancellation. The
-            for table_name in created:
+            # sweep covers all tables, not only the recorded ones: a
-                try:
+            # cancellation can land after lance committed a table's tag but
-                    await tables[table_name].tags.delete(name)
+            # before this attempt recorded it, and preflight guarantees the
-                except Exception:
+            # name was unused, so any occurrence belongs to this attempt.
-                    failed_cleanup.append(table_name)
+            (_, failed_cleanup), cancelled = await _wait_protected(
+                self._delete_tag_locked(name)
+            )
             if failed_cleanup:
                 raise RuntimeError(
-                    f"Tag '{name}' creation failed ({exc}) and cleanup "
+                    f"Tag '{name}' creation failed ({exc!r}) and cleanup "
                     f"failed on: {', '.join(failed_cleanup)}. A partial "
                     "tag may remain; delete it with delete_tag."
                 ) from exc
+            if cancelled and not isinstance(exc, asyncio.CancelledError):
+                raise asyncio.CancelledError()
             raise
 
+    async def _delete_tag_locked(self, name: str) -> tuple[bool, list[str]]:
+        """Delete the tag from every table that has it; the caller must
+        hold the write lock.
+
+        Returns whether the tag was found anywhere and the tables where
+        listing or deletion failed.
+        """
+        found = False
+        failed: list[str] = []
+        for table_name, table in self._tables().items():
+            try:
+                if name in await table.tags.list():
+                    found = True
+                    await table.tags.delete(name)
+            except Exception:
+                failed.append(table_name)
+        return found, failed
+
     async def list_tags(self) -> dict[str, TagInfo]:
         """Aggregate per-table tags into database-level tags.
 
@@ -977,15 +1020,7 @@ class Store:
         self._assert_writable()
         self._assert_not_rebuilding()
         async with self._rebuild_lock, self._write_lock:
-            found = False
+            found, failed = await self._delete_tag_locked(name)
-            failed: list[str] = []
-            for table_name, table in self._tables().items():
-                try:
-                    if name in await table.tags.list():
-                        found = True
-                        await table.tags.delete(name)
-                except Exception:
-                    failed.append(table_name)
             if failed:
                 # A listing failure obscures whether the tag exists on that
                 # table, so failures take precedence over not-found.
@@ -1020,22 +1055,10 @@ class Store:
     ) -> tuple[list[tuple[str, Exception]], bool]:
         """Best-effort rollback that a cancellation cannot interrupt.
 
-        Runs the rollback as a task and keeps waiting for it even if this
+        Returns the rollback failures and whether a cancellation was
-        coroutine is cancelled, so a Ctrl-C cannot leave the rollback half
+        absorbed; the caller must re-deliver an absorbed cancellation.
-        applied. Returns the rollback failures and whether a cancellation
-        was absorbed; the caller must re-deliver an absorbed cancellation.
         """
-        task = asyncio.ensure_future(self._restore_tables(snapshot, best_effort=True))
+        return await _wait_protected(self._restore_tables(snapshot, best_effort=True))
-        cancelled = False
-        while True:
-            try:
-                return await asyncio.shield(task), cancelled
-            except asyncio.CancelledError:
-                if task.done():
-                    # The rollback coroutine itself ended cancelled; there is
-                    # nothing left to wait for.
-                    raise
-                cancelled = True
 
     async def restore_tag(self, name: str) -> str:
         """Restore every table to the versions of a complete tag.

tests/store/test_tags.py

@@ -445,3 +445,90 @@ async def test_delete_tag_reports_listing_failures(temp_db_path, monkeypatch):
 
         await store.delete_tag("release-1")
         assert await store.list_tags() == {}
+
+
+@pytest.mark.asyncio
+async def test_create_tag_cancellation_cleans_up(temp_db_path, monkeypatch):
+    """Cancellation during per-table tag creation must not leave a partial
+    tag behind: cleanup runs before the cancellation propagates."""
+    async with Store(temp_db_path, create=True) as store:
+        real_create = AsyncTags.create
+        calls = {"n": 0}
+
+        async def cancelled_create(self, name: str, version: int) -> None:
+            calls["n"] += 1
+            if calls["n"] == 4:
+                raise asyncio.CancelledError()
+            await real_create(self, name, version)
+
+        monkeypatch.setattr(AsyncTags, "create", cancelled_create)
+
+        with pytest.raises(asyncio.CancelledError):
+            await store.create_tag("broken")
+
+        monkeypatch.undo()
+        assert await store.list_tags() == {}
+
+
+@pytest.mark.asyncio
+async def test_create_tag_cleanup_survives_cancellation(temp_db_path, monkeypatch):
+    """Cancelling create_tag while it cleans up a failed creation does not
+    interrupt the cleanup: no partial tag remains and the cancellation is
+    delivered afterwards."""
+    async with Store(temp_db_path, create=True) as store:
+        real_create = AsyncTags.create
+        real_delete = AsyncTags.delete
+        create_calls = {"n": 0}
+        cleanup_started = asyncio.Event()
+        release = asyncio.Event()
+
+        async def flaky_create(self, name: str, version: int) -> None:
+            create_calls["n"] += 1
+            if create_calls["n"] == 4:
+                raise RuntimeError("create boom")
+            await real_create(self, name, version)
+
+        async def slow_delete(self, name: str) -> None:
+            cleanup_started.set()
+            await release.wait()
+            await real_delete(self, name)
+
+        monkeypatch.setattr(AsyncTags, "create", flaky_create)
+        monkeypatch.setattr(AsyncTags, "delete", slow_delete)
+
+        task = asyncio.create_task(store.create_tag("broken"))
+        await cleanup_started.wait()
+        task.cancel()
+        release.set()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        monkeypatch.undo()
+        assert await store.list_tags() == {}
+
+
+@pytest.mark.asyncio
+async def test_create_tag_cancellation_after_commit_cleans_committed_tag(
+    temp_db_path, monkeypatch
+):
+    """Cancellation arriving after lance committed a table's tag but before
+    the attempt recorded it must still clean that table: cleanup sweeps all
+    tables, relying on the preflight guarantee that the name was unused."""
+    async with Store(temp_db_path, create=True) as store:
+        real_create = AsyncTags.create
+        calls = {"n": 0}
+
+        async def committing_cancelled_create(self, name: str, version: int) -> None:
+            calls["n"] += 1
+            await real_create(self, name, version)
+            if calls["n"] == 4:
+                raise asyncio.CancelledError()
+
+        monkeypatch.setattr(AsyncTags, "create", committing_cancelled_create)
+
+        with pytest.raises(asyncio.CancelledError):
+            await store.create_tag("broken")
+
+        monkeypatch.undo()
+        assert await store.list_tags() == {}