From f02cde5ddf2072576c5e40ddc0813c955b31f295 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 16 Jul 2026 11:52:23 +0300 Subject: [PATCH] Fix same-tick cancellation losing recovery results in _wait_protected --- haiku_rag_slim/haiku/rag/cli.py | 3 ++ haiku_rag_slim/haiku/rag/mcp.py | 36 ++++-------------------- haiku_rag_slim/haiku/rag/store/engine.py | 6 ++-- tests/store/test_restore.py | 25 ++++++++++++++++ tests/test_cli.py | 7 +++++ 5 files changed, 45 insertions(+), 32 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 3dddbfe1..20c7318a 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/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}") diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 4ee94339..e2e9d1bd 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index a5a26c3f..1c5cd926 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/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 diff --git a/tests/store/test_restore.py b/tests/store/test_restore.py index 9f71e6b5..5202bb97 100644 --- a/tests/store/test_restore.py +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index de404775..d16e1ff0 100644 --- a/tests/test_cli.py +++ b/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