From c5814b31bf176ad1553e0a658a40f0f3f0e3ff39 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Mon, 1 Jun 2026 17:29:53 +0300
Subject: [PATCH] Test source clients close after the worker pool stops

Workers share the pollers' Source instances for fetch(), so the httpx
clients must be closed only after the pool has stopped. Guards both
run_batch() and serve() against reintroducing the shutdown-order bug.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/ingester/test_run_batch.py | 58 ++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py
index c646d5d6..8edd461d 100644
--- a/tests/ingester/test_run_batch.py
+++ b/tests/ingester/test_run_batch.py
@@ -23,6 +23,7 @@ from haiku.rag.config import (
     WorkerConfig,
 )
 from haiku.rag.ingester.app import IngesterApp
+from haiku.rag.ingester.pollers.manager import PollerManager
 from haiku.rag.ingester.workers.pool import WorkerPool
 from haiku.rag.store.models.document import Document
 
@@ -305,3 +306,60 @@ async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog):
 
     assert pool.released == 1
     assert "Shutdown grace" in caplog.text
+
+
+def _record_close_order(monkeypatch) -> list[str]:
+    """Record the order of _stop_pool and PollerManager.close_sources.
+
+    Workers share the pollers' Source instances for fetch(), so the source
+    clients must be closed only after the pool has stopped — otherwise an
+    in-flight fetch during the shutdown grace hits a closed client.
+    """
+    order: list[str] = []
+    orig_stop_pool = IngesterApp._stop_pool
+    orig_close = PollerManager.close_sources
+
+    async def rec_stop(self):
+        order.append("stop_pool")
+        await orig_stop_pool(self)
+
+    async def rec_close(self):
+        order.append("close_sources")
+        await orig_close(self)
+
+    monkeypatch.setattr(IngesterApp, "_stop_pool", rec_stop)
+    monkeypatch.setattr(PollerManager, "close_sources", rec_close)
+    return order
+
+
+@pytest.mark.asyncio
+async def test_run_batch_closes_sources_after_pool_stops(
+    tmp_path, use_client, monkeypatch
+):
+    (tmp_path / "a.md").write_text("hello")
+    use_client(_mock_client())
+    app = IngesterApp(config=_config(tmp_path), db_path=tmp_path / "db.lancedb")
+    order = _record_close_order(monkeypatch)
+
+    await app.run_batch()
+
+    assert order == ["stop_pool", "close_sources"]
+
+
+@pytest.mark.asyncio
+async def test_serve_closes_sources_after_pool_stops(tmp_path, use_client, monkeypatch):
+    use_client(_mock_client())
+    config = _config(tmp_path)
+    config.ingester.api = APIConfig(enabled=False)
+    app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
+    order = _record_close_order(monkeypatch)
+
+    task = asyncio.create_task(app.serve(api=False))
+    try:
+        await _wait_until(lambda: app._pool is not None and app._pool.live_workers > 0)
+    finally:
+        task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(task, timeout=5.0)
+
+    assert order == ["stop_pool", "close_sources"]