diff --git a/CHANGELOG.md b/CHANGELOG.md index b12d7385..189b5fc8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- `haiku-ingester run-batch`: one discover sweep across every configured source, drains the queue, then exits. Orphan deletion requires a persisted `ingester.db`. + +### Removed + +- `haiku-ingester run-once`. Use `run-batch` to ingest configured sources, or `serve` for continuous operation. + ## [0.51.0] - 2026-05-29 ### Added diff --git a/docs/ingester.md b/docs/ingester.md index 327363ad..891f0667 100644 --- a/docs/ingester.md +++ b/docs/ingester.md @@ -307,19 +307,24 @@ ingester: ## Operating -### Smoke-test a single URI +### One-shot batch build -`run-once` bypasses the queue and runs a single Job through the -pipeline. Useful for sanity-checking a source before starting the -service. +`run-batch` runs a single discover sweep across every configured source, +drains the queue, then exits. New and changed resources are ingested, +resources that vanished from a source are deleted. The periodic poller +loops never start, so the run is deterministic and finishes as soon as the +queue is empty. This is the mode for building a database in CI or on a +schedule rather than running the service continuously. ```bash -haiku-ingester run-once /path/to/test.pdf -haiku-ingester run-once https://example.com/spec.pdf -haiku-ingester run-once s3://my-bucket/key.pdf +haiku-ingester run-batch +haiku-ingester run-batch --db rag.lancedb ``` -Exit codes: `0` success, `1` transient error, `2` permanent error. +Orphan deletion compares each source against `sync_state` in the queue DB, +so persist `ingester.db` between runs for deletions to be detected. It exits +non-zero if any job dead-letters or a source's discovery sweep does not +complete. ### The queue diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index ecc0d358..7b64556f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -1,9 +1,13 @@ import asyncio import logging import signal +from contextlib import asynccontextmanager +from datetime import UTC, datetime from pathlib import Path from typing import TYPE_CHECKING +from pydantic import BaseModel + from haiku.rag.config import AppConfig from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.queue.migrations import open_queue @@ -17,6 +21,15 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +class BatchReport(BaseModel): + """Outcome of a one-shot batch run: terminal job counts after the queue + drained, plus any sources whose discovery sweep did not complete.""" + + succeeded: int = 0 + dead: int = 0 + failed_sweeps: list[str] = [] + + class IngesterApp: """Top-level lifecycle for the production ingester. @@ -32,11 +45,13 @@ class IngesterApp: self._sync: SyncStateRepo | None = None self._pool: WorkerPool | None = None self._pollers: PollerManager | None = None - self._client = None - async def serve(self, *, api: bool = True) -> None: - """Run pollers + workers (and the HTTP API when enabled) until a - SIGINT/SIGTERM is received. Drains in-flight work on shutdown.""" + @asynccontextmanager + async def _resources(self): + """Open the queue connection and construct the repos, client, pollers + and worker pool. Yields with everything built but nothing started — + callers own the start/stop lifecycle. Closes the client and queue + connection on exit.""" from haiku.rag.client import HaikuRAG from haiku.rag.converters import get_converter @@ -58,13 +73,12 @@ class IngesterApp: jitter=ingester_cfg.workers.retry.jitter, ) - # The ingester is the sole writer for its LanceDB target; - # create on first start so docker-compose / fresh deployments - # don't require a manual `haiku-rag init`. + # The ingester is the sole writer for its LanceDB target; create on + # first start so docker-compose / fresh deployments don't require a + # manual `haiku-rag init`. async with HaikuRAG( self._db_path, config=self._config, create=True ) as client: - self._client = client self._pollers = PollerManager( configs=ingester_cfg.sources, job_repo=self._jobs, @@ -86,73 +100,7 @@ class IngesterApp: # HTTP / WebDAV / S3 fetches reuse credentials. sources=self._pollers.sources, ) - - stop_event = asyncio.Event() - loop = asyncio.get_running_loop() - for sig in (signal.SIGINT, signal.SIGTERM): - try: - loop.add_signal_handler(sig, stop_event.set) - except NotImplementedError: - # Windows; signal handlers unavailable in asyncio. - pass - - await self._pollers.start() - await self._pool.start() - # Log the docling-serve fleet size when relevant so the - # operator can eyeball the worker/instance ratio. The convert - # phase is usually the throughput ceiling. - proc = self._config.processing - uses_docling_serve = ( - proc.converter == "docling-serve" or proc.chunker == "docling-serve" - ) - if uses_docling_serve: - logger.info( - "Ingester running: %d worker(s), %d source(s), " - "%d docling-serve instance(s)", - ingester_cfg.workers.worker_count, - len(ingester_cfg.sources), - len(self._config.providers.docling_serve.base_urls), - ) - else: - logger.info( - "Ingester running: %d worker(s), %d source(s)", - ingester_cfg.workers.worker_count, - len(ingester_cfg.sources), - ) - - api_task, api_server = await self._maybe_start_api(api) - - try: - await stop_event.wait() - finally: - logger.info("Shutting down ingester") - if api_server is not None: - api_server.should_exit = True - if api_task is not None: - await asyncio.gather(api_task, return_exceptions=True) - await self._pollers.stop() - grace_s = ingester_cfg.workers.shutdown_grace_s - try: - await asyncio.wait_for(self._pool.stop(), timeout=grace_s) - except TimeoutError: - # In-flight jobs stay 'claimed'; the reaper resets - # them after claim_timeout_s on the next start. - logger.warning( - "Shutdown grace of %.1fs elapsed with jobs still " - "in flight; cancelling — they'll be reclaimed after " - "claim_timeout_s on next start", - grace_s, - ) - # Drain any cancel-cleanup release Tasks the worker pool - # spawned but didn't get to await (timeout path). They - # need the queue connection that the outer finally is - # about to close. - landed = await self._pool.drain_pending_releases(timeout=2.0) - if landed: - logger.info( - "Drained %d cancel-cleanup release(s) before close", - landed, - ) + yield finally: # Close the queue connection unconditionally. aiosqlite runs the # underlying sqlite3 in a background thread; leaving it open holds @@ -162,6 +110,110 @@ class IngesterApp: await self._queue_conn.close() self._queue_conn = None + async def _stop_pool(self) -> None: + """Stop the worker pool, honouring the shutdown grace, then drain any + cancel-cleanup release tasks before the queue connection closes.""" + assert self._pool is not None + grace_s = self._config.ingester.workers.shutdown_grace_s + try: + await asyncio.wait_for(self._pool.stop(), timeout=grace_s) + except TimeoutError: + # In-flight jobs stay 'claimed'; the reaper resets them after + # claim_timeout_s on the next start. + logger.warning( + "Shutdown grace of %.1fs elapsed with jobs still in flight; " + "cancelling — they'll be reclaimed after claim_timeout_s on " + "next start", + grace_s, + ) + landed = await self._pool.drain_pending_releases(timeout=2.0) + if landed: + logger.info("Drained %d cancel-cleanup release(s) before close", landed) + + async def serve(self, *, api: bool = True) -> None: + """Run pollers + workers (and the HTTP API when enabled) until a + SIGINT/SIGTERM is received. Drains in-flight work on shutdown.""" + ingester_cfg = self._config.ingester + async with self._resources(): + assert self._pollers is not None and self._pool is not None + await self._pollers.start() + await self._pool.start() + # Log the docling-serve fleet size when relevant so the operator + # can eyeball the worker/instance ratio. The convert phase is + # usually the throughput ceiling. + proc = self._config.processing + uses_docling_serve = ( + proc.converter == "docling-serve" or proc.chunker == "docling-serve" + ) + if uses_docling_serve: + logger.info( + "Ingester running: %d worker(s), %d source(s), " + "%d docling-serve instance(s)", + ingester_cfg.workers.worker_count, + len(ingester_cfg.sources), + len(self._config.providers.docling_serve.base_urls), + ) + else: + logger.info( + "Ingester running: %d worker(s), %d source(s)", + ingester_cfg.workers.worker_count, + len(ingester_cfg.sources), + ) + + api_task, api_server = await self._maybe_start_api(api) + + stop_event = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, stop_event.set) + except NotImplementedError: # pragma: no cover - Windows only + # Windows; signal handlers unavailable in asyncio. + pass + + try: + await stop_event.wait() + finally: + logger.info("Shutting down ingester") + if api_server is not None: + api_server.should_exit = True + if api_task is not None: + await asyncio.gather(api_task, return_exceptions=True) + await self._pollers.stop() + await self._stop_pool() + + async def run_batch(self) -> BatchReport: + """Run one discover() sweep across every configured source, drain the + queue to completion, then stop. Unlike `serve`, the periodic poller + loops never start — discovery is driven explicitly, so the run is + deterministic and exits as soon as the queue is empty.""" + async with self._resources(): + assert ( + self._pollers is not None + and self._pool is not None + and self._jobs is not None + ) + # A persisted queue carries terminal rows from previous runs, and + # a recovered URI's dead row gets pruned mid-run, so the report + # counts only jobs that completed at or after this start instant. + started_at = datetime.now(UTC) + await self._pool.start() + try: + failed_sweeps = await self._pollers.sweep_all() + while True: + counts = await self._jobs.counts_by_status() + if not counts.get("queued") and not counts.get("claimed"): + break + await asyncio.sleep(0.1) + completed = await self._jobs.counts_by_status_since(started_at) + return BatchReport( + succeeded=completed.get("succeeded", 0), + dead=completed.get("dead", 0), + failed_sweeps=failed_sweeps, + ) + finally: + await self._stop_pool() + async def _maybe_start_api(self, api: bool): """Spin up the FastAPI control plane on an asyncio task. Returns (task, server) or (None, None) when the API is disabled.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index ab112015..b4a1fff4 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -1,7 +1,5 @@ import asyncio import sys -import uuid -from datetime import UTC, datetime from pathlib import Path import typer @@ -9,7 +7,6 @@ from dotenv import find_dotenv, load_dotenv load_dotenv(find_dotenv(usecwd=True)) -from haiku.rag.client import HaikuRAG # noqa: E402 from haiku.rag.config import ( # noqa: E402 AppConfig, find_config_file, @@ -18,10 +15,7 @@ from haiku.rag.config import ( # noqa: E402 set_config, ) from haiku.rag.ingester.app import IngesterApp # noqa: E402 -from haiku.rag.ingester.exceptions import PermanentError, TransientError # noqa: E402 from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402 -from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus # noqa: E402 -from haiku.rag.ingester.workers.pipeline import run_job # noqa: E402 from haiku.rag.logging import configure_cli_logging # noqa: E402 from haiku.rag.store.exceptions import ( # noqa: E402 MigrationRequiredError, @@ -163,59 +157,30 @@ def serve( asyncio.run(app.serve(api=not no_api)) -@_cli.command("run-once") -def run_once( - uri: str = typer.Argument(..., help="URI to ingest (file://, http(s)://, s3://)."), +@_cli.command("run-batch") +def run_batch( db: Path | None = typer.Option( None, "--db", help="LanceDB path (overrides config.storage.data_dir).", ), - delete: bool = typer.Option( - False, "--delete", help="Run a delete op instead of upsert." - ), ) -> None: - """Build an ad-hoc Job and run it through the worker pipeline once. - - Bypasses the queue — does NOT enqueue. Useful for smoke-testing the - Source adapter + pipeline path without spinning up the full pool. - """ - asyncio.run(_run_once(get_config(), uri, db, delete)) + """Run one discover sweep across every configured source, drain the queue, + then exit. New and changed resources are ingested, resources that vanished + from a source are deleted. Exits non-zero if any job dead-letters or a + source's sweep does not complete.""" + asyncio.run(_run_batch(get_config(), db)) -async def _run_once( - app_config: AppConfig, uri: str, db_path: Path | None, delete: bool -) -> None: +async def _run_batch(app_config: AppConfig, db_path: Path | None) -> None: db = _resolve_db_path(app_config, db_path) - now = datetime.now(UTC) - job = Job( - id=f"adhoc-{uuid.uuid4()}", - source_id="adhoc", - uri=uri, - op=JobOp.DELETE if delete else JobOp.UPSERT, - status=JobStatus.CLAIMED, - attempts=1, - max_attempts=1, - enqueued_at=now, - scheduled_at=now, - claimed_at=now, - claimed_by="run-once", - ) - - async with HaikuRAG(db, config=app_config) as client: - try: - result = await run_job(client, job) - except PermanentError as e: - typer.echo(f"PERMANENT: {e}", err=True) - raise typer.Exit(2) from e - except TransientError as e: - typer.echo(f"TRANSIENT: {e}", err=True) - raise typer.Exit(1) from e - - if result.deleted: - typer.echo(f"Deleted document at {uri}") - else: + app = IngesterApp(config=app_config, db_path=db) + report = await app.run_batch() + typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead") + if report.failed_sweeps: typer.echo( - f"Ingested {uri}: document_id={result.document_id} " - f"revision={result.revision} md5={result.content_hash}" + f"Sources that failed to sweep: {', '.join(report.failed_sweeps)}", + err=True, ) + if report.dead or report.failed_sweeps: + raise typer.Exit(1) diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index c44793a4..6339a5c8 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -77,6 +77,18 @@ class PollerManager: poller._stop.clear() self._tasks.append(asyncio.create_task(poller.run())) + async def sweep_all(self) -> list[str]: + """Run one discover() sweep on every poller, sequentially. Used by + one-shot batch runs that drive discovery explicitly rather than + through the periodic loop. Returns the source ids whose sweep did not + complete (discovery failed, circuit open, or pending work already + queued) so callers can treat a one-shot run as failed.""" + failed: list[str] = [] + for poller in self._pollers: + if not await poller._sweep_once(): + failed.append(poller.source_id) + return failed + async def stop(self) -> None: for poller in self._pollers: await poller.stop() diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 1f1b2b31..c2801e53 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -294,6 +294,20 @@ class JobRepo: rows = await cursor.fetchall() return {row["status"]: row["n"] for row in rows} + async def counts_by_status_since(self, since: datetime) -> dict[str, int]: + """status -> count of jobs that reached a terminal state at or after + `since` (by completed_at). Only succeeded/dead set completed_at, so + those are the only keys returned. Lets a one-shot batch report the + work it finished, independent of terminal rows from earlier runs.""" + async with self._lock: + async with self._conn.execute( + "SELECT status, COUNT(*) AS n FROM jobs " + "WHERE completed_at >= ? GROUP BY status", + (since.isoformat(),), + ) as cursor: + rows = await cursor.fetchall() + return {row["status"]: row["n"] for row in rows} + async def count_succeeded_since(self, seconds: int) -> int: """How many jobs reached `succeeded` in the last `seconds` seconds. Drives the dashboard's rolling-throughput chips.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index a14afe96..122866b8 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -114,7 +114,7 @@ class WorkerPool: async def drain_once(self, worker_id: str = "drain") -> int: """Drain every currently-claimable job to completion on the calling - coroutine. Used by run-once and by tests; not by `start()`.""" + coroutine. Used by tests; not by `start()`.""" processed = 0 while True: job = await self._jobs.claim_next(worker_id) diff --git a/tests/ingester/test_cli.py b/tests/ingester/test_cli.py new file mode 100644 index 00000000..b03fd445 --- /dev/null +++ b/tests/ingester/test_cli.py @@ -0,0 +1,46 @@ +"""haiku-ingester run-batch CLI: echoes the batch report and exits non-zero +when any job dead-letters. IngesterApp is patched so no real ingestion runs.""" + +from unittest.mock import AsyncMock + +from typer.testing import CliRunner + +from haiku.rag.ingester.app import BatchReport +from haiku.rag.ingester.cli import _cli as cli + +runner = CliRunner() + + +def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock: + fake = AsyncMock() + fake.run_batch.return_value = report + monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake) + return fake + + +def test_run_batch_reports_and_exits_zero(monkeypatch): + fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch) + + result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"]) + + assert result.exit_code == 0 + assert "3 succeeded, 0 dead" in result.output + fake.run_batch.assert_awaited_once() + + +def test_run_batch_exits_nonzero_when_dead(monkeypatch): + _fake_app(BatchReport(succeeded=1, dead=2), monkeypatch) + + result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"]) + + assert result.exit_code == 1 + assert "2 dead" in result.output + + +def test_run_batch_exits_nonzero_when_sweep_fails(monkeypatch): + _fake_app(BatchReport(succeeded=2, dead=0, failed_sweeps=["docs"]), monkeypatch) + + result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"]) + + assert result.exit_code == 1 + assert "failed to sweep: docs" in result.output diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py new file mode 100644 index 00000000..c646d5d6 --- /dev/null +++ b/tests/ingester/test_run_batch.py @@ -0,0 +1,307 @@ +"""IngesterApp lifecycle: run_batch (one sweep -> drain -> exit) and serve +(pollers + workers + API until shutdown). The document store (HaikuRAG) is +patched out — the behavior under test is the orchestration and orphan +pruning, not embedding.""" + +import asyncio +from contextlib import asynccontextmanager +from pathlib import Path +from unittest.mock import AsyncMock +from urllib.parse import unquote, urlparse + +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.client.exceptions import UnsupportedSourceError +from haiku.rag.config import ( + APIConfig, + AppConfig, + FSSourceConfig, + IngesterConfig, + QueueConfig, + RetryPolicyConfig, + WorkerConfig, +) +from haiku.rag.ingester.app import IngesterApp +from haiku.rag.ingester.workers.pool import WorkerPool +from haiku.rag.store.models.document import Document + + +def _config(tmp_path, **worker_kwargs) -> AppConfig: + workers = WorkerConfig( + worker_count=2, + poll_idle_interval_s=0.05, + retry=RetryPolicyConfig(max_attempts=1), + **worker_kwargs, + ) + return AppConfig( + ingester=IngesterConfig( + queue=QueueConfig(path=tmp_path / "queue.db"), + sources=[ + FSSourceConfig( + type="fs", + id="local", + root=tmp_path, + poll_interval_s=3600.0, + ) + ], + workers=workers, + api=APIConfig(enabled=False), + ), + ) + + +def _mock_client() -> AsyncMock: + """A HaikuRAG mock whose upserts return a fresh Document (with an md5 so + the worker records sync_state) and whose lookups return a Document with a + deterministic id so the DELETE path resolves.""" + client = AsyncMock(spec=HaikuRAG) + counter = {"n": 0} + + async def _create(uri, *_, **__): + counter["n"] += 1 + # Mirror the real client: persist the FS revision (mtime_ns) so the + # poller's change-detection skips unchanged files on the next sweep. + path = Path(unquote(urlparse(uri).path)) + return Document( + id=f"doc-{counter['n']}", + content="x", + uri=uri, + metadata={ + "content_type": "text/markdown", + "md5": f"md5-{counter['n']}", + "source_revision": str(path.stat().st_mtime_ns), + }, + ) + + async def _get_by_uri(uri): + return Document(id=f"doc-for-{uri}", content="x", uri=uri, metadata={}) + + client.create_document_from_source.side_effect = _create + client.get_document_by_uri.side_effect = _get_by_uri + return client + + +@pytest.fixture +def use_client(monkeypatch): + """Make IngesterApp's internally-created HaikuRAG resolve to the given + mock client. Call again to swap the client for a later run in the same + test.""" + + def _install(client): + @asynccontextmanager + async def _cm(*_, **__): + yield client + + monkeypatch.setattr("haiku.rag.client.HaikuRAG", lambda *a, **k: _cm()) + + return _install + + +@pytest.mark.asyncio +async def test_run_batch_drains_upserts(tmp_path, use_client): + (tmp_path / "a.md").write_text("hello") + (tmp_path / "b.md").write_text("world") + + client = _mock_client() + use_client(client) + + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch() + + assert report.succeeded == 2 + assert report.dead == 0 + assert client.create_document_from_source.await_count == 2 + ingested = { + call.args[0] for call in client.create_document_from_source.await_args_list + } + assert ingested == { + (tmp_path / "a.md").as_uri(), + (tmp_path / "b.md").as_uri(), + } + + +@pytest.mark.asyncio +async def test_run_batch_prunes_orphans(tmp_path, use_client): + (tmp_path / "a.md").write_text("hello") + (tmp_path / "b.md").write_text("world") + + client = _mock_client() + use_client(client) + config = _config(tmp_path) + db_path = tmp_path / "db.lancedb" + + # First batch ingests both files and records sync_state for each. + first = await IngesterApp(config=config, db_path=db_path).run_batch() + assert first.succeeded == 2 + client.delete_document.assert_not_awaited() + + # b.md disappears from the source tree. The next sweep sees it in + # sync_state but not on disk -> enqueues a DELETE for it. + (tmp_path / "b.md").unlink() + + second = await IngesterApp(config=config, db_path=db_path).run_batch() + + # a.md is unchanged (same mtime) so it's not re-ingested; only the orphan + # delete runs. + assert client.create_document_from_source.await_count == 2 + client.delete_document.assert_awaited_once() + assert second.succeeded == 1 + assert second.dead == 0 + + +@pytest.mark.asyncio +async def test_run_batch_reports_dead_on_permanent_failure(tmp_path, use_client): + (tmp_path / "a.md").write_text("hello") + + client = _mock_client() + client.create_document_from_source.side_effect = UnsupportedSourceError("nope") + use_client(client) + + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch() + + assert report.succeeded == 0 + assert report.dead == 1 + + +@pytest.mark.asyncio +async def test_run_batch_recovered_doc_is_not_counted_as_dead(tmp_path, use_client): + """A doc that dead-lettered in an earlier run and succeeds in this one + prunes its dead row. The report counts only this run's terminal jobs, so + the recovery shows as a success, never a negative or stale dead count.""" + (tmp_path / "a.md").write_text("hello") + config = _config(tmp_path) + db_path = tmp_path / "db.lancedb" + + failing = _mock_client() + failing.create_document_from_source.side_effect = UnsupportedSourceError("nope") + use_client(failing) + first = await IngesterApp(config=config, db_path=db_path).run_batch() + assert first.dead == 1 + + healthy = _mock_client() + use_client(healthy) + second = await IngesterApp(config=config, db_path=db_path).run_batch() + assert second.dead == 0 + assert second.succeeded == 1 + + +@pytest.mark.asyncio +async def test_run_batch_reports_failed_sweep( + tmp_path, use_client, monkeypatch, caplog +): + """A source whose discover() raises is reported in failed_sweeps so the + run can be treated as failed rather than a silent empty success.""" + (tmp_path / "a.md").write_text("hello") + use_client(_mock_client()) + + async def _failing_discover(self, **kwargs): + raise RuntimeError("discover blew up") + yield # unreachable; makes this an async generator + + monkeypatch.setattr( + "haiku.rag.ingester.sources.fs.FSSource.discover", _failing_discover + ) + + with caplog.at_level("ERROR", logger="haiku.rag.ingester.pollers.base"): + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch() + + assert report.failed_sweeps == ["local"] + assert report.succeeded == 0 + assert report.dead == 0 + assert "discover() failed" in caplog.text + + +@pytest.mark.asyncio +async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client): + client = _mock_client() + use_client(client) + + report = await asyncio.wait_for( + IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch(), + timeout=5.0, + ) + + assert report.succeeded == 0 + assert report.dead == 0 + client.create_document_from_source.assert_not_awaited() + + +async def _wait_until(predicate, *, timeout: float = 5.0): + deadline = asyncio.get_running_loop().time() + timeout + while not predicate(): + if asyncio.get_running_loop().time() > deadline: + raise AssertionError("condition not reached within timeout") + await asyncio.sleep(0.02) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api", [True, False]) +async def test_serve_starts_workers_pollers_and_shuts_down(tmp_path, use_client, api): + """serve brings up pollers, workers and (when enabled) the HTTP API, then + tears them all down. Shutdown is driven here by cancelling the serve task, + which runs the same drain-and-close path as a SIGINT/SIGTERM.""" + use_client(_mock_client()) + config = _config(tmp_path) + config.ingester.api = APIConfig(enabled=api, host="127.0.0.1", port=0) + app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") + + task = asyncio.create_task(app.serve(api=api)) + try: + await _wait_until( + lambda: ( + app._pool is not None + and app._pool.live_workers > 0 + and app._pollers is not None + and app._pollers.live_pollers > 0 + ) + ) + finally: + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5.0) + + # Workers and pollers are stopped after the shutdown path runs. + assert app._pool is not None and app._pollers is not None + assert app._pool.live_workers == 0 + assert app._pollers.live_pollers == 0 + + +class _SlowPool(WorkerPool): + """A WorkerPool whose stop never finishes within the grace, to exercise + _stop_pool's timeout path. The real wiring is bypassed since _stop_pool + only calls stop() and drain_pending_releases().""" + + def __init__(self) -> None: + self.released = 0 + + async def stop(self) -> None: + await asyncio.sleep(1.0) + + async def drain_pending_releases(self, timeout: float = 2.0) -> int: + self.released += 1 + return 2 + + +@pytest.mark.asyncio +async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog): + """When a worker doesn't stop within the shutdown grace, _stop_pool logs a + warning and still drains any pending cancel-cleanup releases.""" + config = _config(tmp_path, shutdown_grace_s=0.01) + app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") + pool = _SlowPool() + app._pool = pool + + with caplog.at_level("WARNING", logger="haiku.rag.ingester.app"): + await app._stop_pool() + + assert pool.released == 1 + assert "Shutdown grace" in caplog.text