From b144e620de8ad9b860b2aba31753e0a836410fd0 Mon Sep 17 00:00:00 2001
From: Chris McDonough <chrism@plope.com>
Date: Mon, 1 Jun 2026 09:47:10 -0400
Subject: [PATCH] Fix dead-worker condition and test for run_batch abort

The condition only checked claimed jobs, but queued jobs with no
live workers also hang forever. Check live_workers == 0 regardless
of whether outstanding work is queued or claimed.

Rewrite the test to actually crash workers: patch _process to raise
a bare Exception (which _worker_loop doesn't catch), use
worker_count=1 so the single crash leaves live_workers == 0, and
assert the abort log message fires.
---
 haiku_rag_slim/haiku/rag/ingester/app.py | 11 ++---
 tests/ingester/test_run_batch.py         | 51 ++++++++++--------------
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py
index 23234582..2d51dedd 100644
--- a/haiku_rag_slim/haiku/rag/ingester/app.py
+++ b/haiku_rag_slim/haiku/rag/ingester/app.py
@@ -204,12 +204,13 @@ class IngesterApp:
                     counts = await self._jobs.counts_by_status()
                     if not counts.get("queued") and not counts.get("claimed"):
                         break
-                    if counts.get("claimed") and self._pool.live_workers == 0:
+                    if self._pool.live_workers == 0:
+                        outstanding = counts.get("queued", 0) + counts.get("claimed", 0)
                         logger.error(
-                            "All workers have died with %d claimed job(s) — "
-                            "aborting batch; stranded jobs will be reaped on "
-                            "next start",
-                            counts["claimed"],
+                            "All workers have died with %d outstanding job(s) "
+                            "— aborting batch; stranded jobs will be reaped "
+                            "on next start",
+                            outstanding,
                         )
                         break
                     await asyncio.sleep(0.1)
diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py
index 60a96d8e..dbb616c6 100644
--- a/tests/ingester/test_run_batch.py
+++ b/tests/ingester/test_run_batch.py
@@ -235,45 +235,36 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
 
 
 @pytest.mark.asyncio
-async def test_run_batch_aborts_when_all_workers_die(tmp_path, use_client, monkeypatch):
-    """If all workers crash with claimed jobs still outstanding, run_batch
-    should break out of the drain loop instead of spinning forever."""
+async def test_run_batch_aborts_when_all_workers_die(
+    tmp_path, use_client, monkeypatch, caplog
+):
+    """If all workers crash with outstanding jobs, run_batch should break
+    out of the drain loop instead of spinning forever."""
     (tmp_path / "a.md").write_text("hello")
 
     client = _mock_client()
-    # Block the worker forever so the job stays claimed until we kill it.
-    stall = asyncio.Event()
-    client.create_document_from_source.side_effect = lambda *a, **k: stall.wait()
-
     use_client(client)
     config = _config(tmp_path)
-    app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
 
-    async def _kill_workers_after_claim():
-        """Wait until at least one job is claimed, then kill all workers."""
-        pool = app._pool
-        assert pool is not None
-        for _ in range(200):
-            counts = await app._jobs.counts_by_status()
-            if counts.get("claimed"):
-                break
-            await asyncio.sleep(0.05)
-        for task in pool._workers:
-            task.cancel()
-        await asyncio.gather(*pool._workers, return_exceptions=True)
+    config.ingester.workers.worker_count = 1
 
-    # Run the killer concurrently with run_batch.
-    batch_task = asyncio.create_task(app.run_batch())
-    # Give run_batch a moment to start, then schedule the killer.
-    await asyncio.sleep(0.1)
-    killer_task = asyncio.create_task(_kill_workers_after_claim())
+    # Patch _process to raise an unhandled exception, simulating a hard
+    # worker crash.  _process only catches CancelledError, PermanentError,
+    # and TransientError — anything else propagates and kills the task.
+    monkeypatch.setattr(
+        WorkerPool,
+        "_process",
+        AsyncMock(side_effect=Exception("worker crash")),
+    )
+
+    with caplog.at_level("ERROR", logger="haiku.rag.ingester.app"):
+        report = await asyncio.wait_for(
+            IngesterApp(config=config, db_path=tmp_path / "db.lancedb").run_batch(),
+            timeout=10.0,
+        )
 
-    report = await asyncio.wait_for(batch_task, timeout=10.0)
-    # Killer may still be running against a closed DB — suppress errors.
-    killer_task.cancel()
-    await asyncio.gather(killer_task, return_exceptions=True)
-    # The batch should have exited without hanging.
     assert report is not None
+    assert "All workers have died" in caplog.text
 
 
 async def _wait_until(predicate, *, timeout: float = 5.0):