Fix same-tick cancellation losing recovery results in _wait_protected

haiku_rag_slim/haiku/rag/cli.py

@@ -686,6 +686,9 @@ def tag_restore(  # pragma: no cover
     ),
 ):
     app = create_app(db)
+    if app._is_local and not app.db_path.exists():
+        typer.echo(f"Error: Database path does not exist: {app.db_path}", err=True)
+        raise typer.Exit(1)
     if not yes:
         typer.echo(f"Database: {app.db_path}")
         typer.echo(f"Tag: {name}")

haiku_rag_slim/haiku/rag/mcp.py

@@ -100,11 +100,7 @@ def create_mcp_server(
         response (smaller JSON payload for plain-text consumers).
         """
         try:
-            async with HaikuRAG(
-                db_path,
-                config=config,
-                read_only=read_only,
-            ) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 return await rag.search(
                     query, limit=limit, include_images=include_images
                 )
@@ -139,11 +135,7 @@ def create_mcp_server(
             except Exception:
                 return []
             try:
-                async with HaikuRAG(
-                    db_path,
-                    config=config,
-                    read_only=read_only,
-                ) as rag:
+                async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                     return await rag.search(
                         raw, limit=limit, include_images=include_images
                     )
@@ -154,11 +146,7 @@ def create_mcp_server(
     async def get_document(document_id: str) -> Document | None:
         """Get a document by its ID."""
         try:
-            async with HaikuRAG(
-                db_path,
-                config=config,
-                read_only=read_only,
-            ) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 return await rag.get_document_by_id(document_id)
         except Exception:
             return None
@@ -177,11 +165,7 @@ def create_mcp_server(
             filter: Optional SQL WHERE clause to filter documents.
         """
         try:
-            async with HaikuRAG(
-                db_path,
-                config=config,
-                read_only=read_only,
-            ) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 documents = await rag.list_documents(limit, offset, filter)
 
                 return [
@@ -211,11 +195,7 @@ def create_mcp_server(
             The answer as a string.
         """
         try:
-            async with HaikuRAG(
-                db_path,
-                config=config,
-                read_only=read_only,
-            ) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 answer, citations = await rag.ask(question)
                 if cite and citations:
                     answer += "\n\n" + format_citations(citations)
@@ -242,11 +222,7 @@ def create_mcp_server(
             The answer as a string.
         """
         try:
-            async with HaikuRAG(
-                db_path,
-                config=config,
-                read_only=read_only,
-            ) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 result = await rag.analyze(question, filter=filter)
                 return result.answer
         except Exception as e:

haiku_rag_slim/haiku/rag/store/engine.py

@@ -215,9 +215,11 @@ async def _wait_protected[T](coro: Coroutine[Any, Any, T]) -> tuple[T, bool]:
         try:
             return await asyncio.shield(task), cancelled
         except asyncio.CancelledError:
-            if task.done():
+            if task.cancelled():
                 # The recovery coroutine itself ended cancelled; there is
-                # nothing left to wait for.
+                # nothing left to wait for. A task that completed (even in
+                # the same tick as the cancellation) still returns its
+                # result on the next pass.
                 raise
             cancelled = True
 

tests/store/test_restore.py

@@ -381,3 +381,28 @@ async def test_restore_failure_rollback_survives_cancellation(
         # 3 forward calls (2 ok, 1 failed) + all 5 rollback calls ran.
         assert calls["n"] == 8
         assert await _doc_contents(store) == {"First document", "Second document"}
+
+
+@pytest.mark.asyncio
+async def test_wait_protected_returns_result_on_same_tick_cancellation():
+    """A cancellation landing after the recovery task completed but before
+    the waiter resumed must not discard the recovery result."""
+    import asyncio
+
+    from haiku.rag.store.engine import _wait_protected
+
+    async def recovery() -> str:
+        return "done"
+
+    outer = asyncio.create_task(_wait_protected(recovery()))
+    # First pass: outer starts, spawns the recovery task, suspends on shield.
+    await asyncio.sleep(0)
+    # Second pass: the recovery task completes; outer is scheduled to resume.
+    await asyncio.sleep(0)
+    # Cancellation beats the resumption: delivered at the shield await even
+    # though the recovery already finished.
+    outer.cancel()
+
+    result, cancelled = await outer
+    assert result == "done"
+    assert cancelled is True

tests/test_cli.py

@@ -256,6 +256,13 @@ class TestTagRestore:
         assert result.exit_code == 1
         assert "does not exist" in result.output
 
+        # Without --yes the missing database is reported before the
+        # confirmation prompt, not after the user confirms.
+        result = runner.invoke(cli, ["tag", "restore", "r1", "--db", str(missing)])
+        assert result.exit_code == 1
+        assert "does not exist" in result.output
+        assert "Continue?" not in result.output
+
     def test_tag_help_includes_restore(self):
         result = runner.invoke(cli, ["tag", "--help"])
         assert result.exit_code == 0