diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index c88902c2..559fd74f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING from pydantic import BaseModel from haiku.rag.config import AppConfig -from haiku.rag.ingester.batch import BatchDryRunReport +from haiku.rag.ingester.batch import BatchDryRunReport, BatchManifest from haiku.rag.ingester.metadata import build_providers, load_metadata_providers from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.queue.migrations import open_queue @@ -179,6 +179,28 @@ class IngesterApp: if landed: logger.info("Drained %d cancel-cleanup release(s) before close", landed) + async def _drain_batch(self, started_at: datetime) -> BatchReport: + assert self._pool is not None and self._jobs is not None + while True: + counts = await self._jobs.counts_by_status() + if not counts.get("queued") and not counts.get("claimed"): + break + if self._pool.live_workers == 0: + outstanding = counts.get("queued", 0) + counts.get("claimed", 0) + logger.error( + "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) + completed = await self._jobs.counts_by_status_since(started_at) + return BatchReport( + succeeded=completed.get("succeeded", 0), + dead=completed.get("dead", 0), + ) + 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.""" @@ -250,26 +272,9 @@ class IngesterApp: 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 - if self._pool.live_workers == 0: - outstanding = counts.get("queued", 0) + counts.get("claimed", 0) - logger.error( - "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) - 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, - ) + report = await self._drain_batch(started_at) + report.failed_sweeps = failed_sweeps + return report finally: await self._stop_pool() await self._pollers.close_sources() @@ -282,6 +287,91 @@ class IngesterApp: manifest, failed_sweeps = await self._pollers.dry_run_manifest() return BatchDryRunReport(manifest=manifest, failed_sweeps=failed_sweeps) + async def run_batch_from_manifest(self, manifest: BatchManifest) -> BatchReport: + """Enqueue and drain a dry-run manifest without running a fresh + discovery sweep.""" + if manifest.version != 1: + raise ValueError(f"Unsupported manifest version: {manifest.version}") + async with self._resources(): + assert ( + self._pollers is not None + and self._pool is not None + and self._jobs is not None + ) + configured = {source.source_id for source in self._pollers.sources} + manifest_sources = {change.source_id for change in manifest.changes} + missing = sorted(manifest_sources - configured) + if missing: + await self._pollers.close_sources() + raise ValueError( + "Manifest references unconfigured source(s): " + ", ".join(missing) + ) + + pending = [ + source_id + for source_id in sorted(manifest_sources) + if await self._jobs.has_pending(source_id) + ] + if pending: + await self._pollers.close_sources() + raise ValueError( + "Cannot replay manifest while source(s) have pending work: " + + ", ".join(pending) + ) + + seen: set[tuple[str, str]] = set() + duplicates: set[tuple[str, str]] = set() + for change in manifest.changes: + key = (change.source_id, change.uri) + if key in seen: + duplicates.add(key) + seen.add(key) + if duplicates: + await self._pollers.close_sources() + rendered = ", ".join( + f"{source_id}:{uri}" for source_id, uri in duplicates + ) + raise ValueError(f"Manifest contains duplicate change(s): {rendered}") + + default_max_attempts = self._config.ingester.workers.retry.max_attempts + max_attempts_by_source = { + poller.source_id: ( + poller.config.retry.max_attempts + if poller.config.retry is not None + else default_max_attempts + ) + for poller in self._pollers.pollers + } + for change in manifest.changes: + job = await self._jobs.enqueue( + change.source_id, + change.uri, + op=change.op, + revision=change.revision, + max_attempts=max_attempts_by_source[change.source_id], + extra={ + "_manifest": { + "version": manifest.version, + "generated_at": manifest.generated_at.isoformat(), + "discovered_at": change.discovered_at.isoformat(), + } + }, + ) + if job is None: + await self._pollers.close_sources() + raise ValueError( + "Cannot replay manifest because a live job already exists " + f"for {change.source_id}:{change.uri}" + ) + + started_at = datetime.now(UTC) + await self._pool.start() + try: + return await self._drain_batch(started_at) + finally: + await self._stop_pool() + await self._pollers.close_sources() + 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 7ecb1941..0fb63bd0 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -150,6 +150,11 @@ def _write_manifest(manifest: BatchManifest, path: Path) -> None: path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") +def _load_manifest(path: Path) -> BatchManifest: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + return BatchManifest.model_validate(data) + + @_cli.command("serve") def serve( db: Path | None = typer.Option( @@ -212,12 +217,29 @@ def run_batch( "-o", help="Dry-run manifest path (defaults to manifest-.yaml).", ), + manifest: Path | None = typer.Option( + None, + "--manifest", + help="Replay a dry-run manifest instead of running discovery.", + ), ) -> None: """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, dry_run=dry_run, output=output)) + if manifest is not None and dry_run: + raise typer.BadParameter("--manifest cannot be combined with --dry-run") + if manifest is not None and output is not None: + raise typer.BadParameter("--output is only valid with --dry-run") + asyncio.run( + _run_batch( + get_config(), + db, + dry_run=dry_run, + output=output, + manifest_path=manifest, + ) + ) async def _run_batch( @@ -226,6 +248,7 @@ async def _run_batch( *, dry_run: bool = False, output: Path | None = None, + manifest_path: Path | None = None, ) -> None: db = _resolve_db_path(app_config, db_path) app = IngesterApp(config=app_config, db_path=db) @@ -248,6 +271,20 @@ async def _run_batch( ) return + if manifest_path is not None: + try: + manifest = _load_manifest(manifest_path) + report = await app.run_batch_from_manifest(manifest) + except ValueError as exc: + typer.echo(f"Error: {exc}") + raise typer.Exit(1) from exc + typer.echo( + f"Manifest batch complete: {report.succeeded} succeeded, {report.dead} dead" + ) + if report.dead: + raise typer.Exit(1) + return + report = await app.run_batch() typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead") if report.failed_sweeps: diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index f6c4bc90..977abbac 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -110,23 +110,38 @@ async def run_job( ), ): try: + manifest_context = extra.get("_manifest") if job.op is JobOp.DELETE: # An atomic-rename save can let a spurious DELETE win the # enqueue race while the file is mid-rewrite. If the resource # is already back, skip the delete (it would blackhole a live - # document) and let the next sweep re-ingest it. - try: - source = resolve_configured_source(job.uri, job.source_id, sources) - restored = await source.head(job.uri) is not None - except Exception: - restored = False - if restored: - return JobResult(deleted=False) + # document) and let the next sweep re-ingest it. Manifest + # replay intentionally follows the frozen dry-run changeset. + if manifest_context is None: + try: + source = resolve_configured_source( + job.uri, job.source_id, sources + ) + restored = await source.head(job.uri) is not None + except Exception: + restored = False + if restored: + return JobResult(deleted=False) doc = await client.get_document_by_uri(job.uri) if doc is not None and doc.id is not None: await client.delete_document(doc.id) return JobResult(deleted=True) + if manifest_context is not None and job.revision is not None: + source = resolve_configured_source(job.uri, job.source_id, sources) + current_revision = await source.head(job.uri) + if current_revision != job.revision: + raise PermanentError( + "manifest revision is stale for " + f"{job.uri}: expected {job.revision!r}, " + f"current {current_revision!r}" + ) + result = await client.create_document_from_source( job.uri, sources=sources, diff --git a/tests/ingester/test_cli.py b/tests/ingester/test_cli.py index 2b5a0903..07ef4e24 100644 --- a/tests/ingester/test_cli.py +++ b/tests/ingester/test_cli.py @@ -68,6 +68,17 @@ def _fake_dry_run_app(report: BatchDryRunReport, monkeypatch) -> AsyncMock: return fake +def _fake_manifest_app(report: BatchReport, monkeypatch) -> AsyncMock: + fake = AsyncMock() + fake.run_batch_from_manifest.return_value = report + monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake) + return fake + + +def _write_manifest(path) -> None: + path.write_text(yaml.safe_dump(_manifest().model_dump(mode="json"))) + + def test_run_batch_reports_and_exits_zero(monkeypatch): fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch) @@ -152,6 +163,80 @@ def test_run_batch_dry_run_exits_nonzero_when_sweep_fails(monkeypatch, tmp_path) assert not output.exists() +def test_run_batch_manifest_replays_manifest(monkeypatch, tmp_path): + manifest_path = tmp_path / "manifest.yaml" + _write_manifest(manifest_path) + fake = _fake_manifest_app(BatchReport(succeeded=2, dead=0), monkeypatch) + + result = runner.invoke( + cli, ["run-batch", "--manifest", str(manifest_path), "--db", "x.lancedb"] + ) + + assert result.exit_code == 0 + assert "Manifest batch complete: 2 succeeded, 0 dead" in result.output + fake.run_batch_from_manifest.assert_awaited_once() + loaded = fake.run_batch_from_manifest.await_args.args[0] + assert isinstance(loaded, BatchManifest) + assert loaded.changes[0].uri == "file:///a.md" + fake.run_batch.assert_not_awaited() + fake.run_batch_dry_run.assert_not_awaited() + + +def test_run_batch_manifest_exits_nonzero_when_dead(monkeypatch, tmp_path): + manifest_path = tmp_path / "manifest.yaml" + _write_manifest(manifest_path) + _fake_manifest_app(BatchReport(succeeded=1, dead=1), monkeypatch) + + result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)]) + + assert result.exit_code == 1 + assert "1 dead" in result.output + + +def test_run_batch_manifest_reports_validation_error(monkeypatch, tmp_path): + manifest_path = tmp_path / "manifest.yaml" + _write_manifest(manifest_path) + fake = AsyncMock() + fake.run_batch_from_manifest.side_effect = ValueError("bad manifest") + monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake) + + result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)]) + + assert result.exit_code == 1 + assert "Error: bad manifest" in result.output + + +def test_run_batch_manifest_conflicts_with_dry_run(tmp_path): + manifest_path = tmp_path / "manifest.yaml" + _write_manifest(manifest_path) + + result = runner.invoke( + cli, ["run-batch", "--manifest", str(manifest_path), "--dry-run"] + ) + + assert result.exit_code != 0 + assert "--manifest cannot be combined with --dry-run" in result.output + + +def test_run_batch_manifest_conflicts_with_output(tmp_path): + manifest_path = tmp_path / "manifest.yaml" + _write_manifest(manifest_path) + + result = runner.invoke( + cli, + [ + "run-batch", + "--manifest", + str(manifest_path), + "--output", + str(tmp_path / "out.yaml"), + ], + ) + + assert result.exit_code != 0 + assert "--output is only valid with --dry-run" in result.output + + # --- serve --- diff --git a/tests/ingester/test_pipeline.py b/tests/ingester/test_pipeline.py index 34468670..882c9ee2 100644 --- a/tests/ingester/test_pipeline.py +++ b/tests/ingester/test_pipeline.py @@ -16,6 +16,7 @@ def _job( *, op: JobOp = JobOp.UPSERT, uri: str = "https://example.com/a.pdf", + revision: str | None = None, extra: dict | None = None, attempts: int = 0, ) -> Job: @@ -25,6 +26,7 @@ def _job( source_id="src", uri=uri, op=op, + revision=revision, status=JobStatus.CLAIMED, attempts=attempts, max_attempts=5, @@ -255,6 +257,22 @@ async def test_delete_skipped_when_resource_restored_on_source(): client.delete_document.assert_not_awaited() +@pytest.mark.asyncio +async def test_manifest_delete_proceeds_when_resource_restored_on_source(): + client = _mock_client() + client.get_document_by_uri.return_value = Document(id="doc-9", content="", uri="u") + sources: list[Source] = [_StubSource("src", "12345")] + + result = await run_job( + client, + _job(op=JobOp.DELETE, extra={"_manifest": {"version": 1}}), + sources=sources, + ) + + assert result.deleted is True + client.delete_document.assert_awaited_once_with("doc-9") + + @pytest.mark.asyncio async def test_delete_proceeds_when_resource_absent_on_source(): client = _mock_client() @@ -267,6 +285,42 @@ async def test_delete_proceeds_when_resource_absent_on_source(): client.delete_document.assert_awaited_once_with("doc-9") +@pytest.mark.asyncio +async def test_manifest_upsert_rejects_stale_revision_before_fetch(): + client = _mock_client() + sources: list[Source] = [_StubSource("src", "r2")] + + with pytest.raises(PermanentError, match="manifest revision is stale"): + await run_job( + client, + _job(revision="r1", extra={"_manifest": {"version": 1}}), + sources=sources, + ) + + client.create_document_from_source.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_manifest_upsert_ingests_when_revision_matches(): + client = _mock_client() + client.create_document_from_source.return_value = Document( + id="doc-42", + content="x", + uri="https://example.com/a.pdf", + metadata={"md5": "abcd", "source_revision": "r1"}, + ) + sources: list[Source] = [_StubSource("src", "r1")] + + result = await run_job( + client, + _job(revision="r1", extra={"_manifest": {"version": 1}}), + sources=sources, + ) + + assert result.document_id == "doc-42" + client.create_document_from_source.assert_awaited_once() + + @pytest.mark.asyncio async def test_delete_proceeds_when_source_unresolvable(): """No configured source for the job: the probe can't run, so the delete diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index c9030965..636d0d9b 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -5,6 +5,7 @@ pruning, not embedding.""" import asyncio from contextlib import asynccontextmanager +from datetime import UTC, datetime from pathlib import Path from unittest.mock import AsyncMock from urllib.parse import unquote, urlparse @@ -23,6 +24,7 @@ from haiku.rag.config import ( WorkerConfig, ) from haiku.rag.ingester.app import IngesterApp +from haiku.rag.ingester.batch import BatchChange, BatchManifest from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.queue.migrations import open_queue from haiku.rag.ingester.queue.models import JobOp @@ -86,6 +88,10 @@ def _mock_client() -> AsyncMock: return client +def _manifest(*changes: BatchChange) -> BatchManifest: + return BatchManifest(generated_at=datetime.now(UTC), changes=list(changes)) + + @pytest.fixture def use_client(monkeypatch): """Make IngesterApp's internally-created HaikuRAG resolve to the given @@ -278,6 +284,138 @@ async def test_run_batch_dry_run_reports_manifest_without_mutating_queue(tmp_pat await engine.dispose() +@pytest.mark.asyncio +async def test_run_batch_from_manifest_drains_changes_without_sweeping( + tmp_path, use_client, monkeypatch +): + (tmp_path / "a.md").write_text("hello") + revision = str((tmp_path / "a.md").stat().st_mtime_ns) + client = _mock_client() + use_client(client) + sweep_all = AsyncMock(side_effect=AssertionError("manifest replay must not sweep")) + monkeypatch.setattr(PollerManager, "sweep_all", sweep_all) + + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch_from_manifest( + _manifest( + BatchChange( + op=JobOp.UPSERT, + source_id="local", + uri=(tmp_path / "a.md").as_uri(), + revision=revision, + discovered_at=datetime.now(UTC), + ) + ) + ) + + assert report.succeeded == 1 + assert report.dead == 0 + client.create_document_from_source.assert_awaited_once() + sweep_all.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_batch_from_manifest_rejects_stale_upsert_revision( + tmp_path, use_client +): + (tmp_path / "a.md").write_text("hello") + client = _mock_client() + use_client(client) + + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch_from_manifest( + _manifest( + BatchChange( + op=JobOp.UPSERT, + source_id="local", + uri=(tmp_path / "a.md").as_uri(), + revision="stale", + discovered_at=datetime.now(UTC), + ) + ) + ) + + assert report.succeeded == 0 + assert report.dead == 1 + client.create_document_from_source.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_batch_from_manifest_delete_uses_manifest_even_if_file_reappears( + tmp_path, use_client +): + path = tmp_path / "gone.md" + path.write_text("back") + client = _mock_client() + use_client(client) + + report = await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch_from_manifest( + _manifest( + BatchChange( + op=JobOp.DELETE, + source_id="local", + uri=path.as_uri(), + discovered_at=datetime.now(UTC), + ) + ) + ) + + assert report.succeeded == 1 + assert report.dead == 0 + client.delete_document.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_batch_from_manifest_rejects_pending_work(tmp_path, use_client): + (tmp_path / "a.md").write_text("hello") + config = _config(tmp_path) + client = _mock_client() + use_client(client) + engine = await open_queue(config.ingester.queue) + try: + jobs = JobRepo(engine) + await jobs.enqueue("local", (tmp_path / "a.md").as_uri(), op=JobOp.UPSERT) + finally: + await engine.dispose() + + with pytest.raises(ValueError, match="pending work"): + await IngesterApp( + config=config, db_path=tmp_path / "db.lancedb" + ).run_batch_from_manifest( + _manifest( + BatchChange( + op=JobOp.UPSERT, + source_id="local", + uri=(tmp_path / "a.md").as_uri(), + discovered_at=datetime.now(UTC), + ) + ) + ) + + +@pytest.mark.asyncio +async def test_run_batch_from_manifest_rejects_duplicate_changes(tmp_path, use_client): + path = tmp_path / "a.md" + path.write_text("hello") + client = _mock_client() + use_client(client) + change = BatchChange( + op=JobOp.UPSERT, + source_id="local", + uri=path.as_uri(), + discovered_at=datetime.now(UTC), + ) + + with pytest.raises(ValueError, match="duplicate change"): + await IngesterApp( + config=_config(tmp_path), db_path=tmp_path / "db.lancedb" + ).run_batch_from_manifest(_manifest(change, change)) + + @pytest.mark.asyncio async def test_run_batch_aborts_when_all_workers_die( tmp_path, use_client, monkeypatch, caplog