diff --git a/CHANGELOG.md b/CHANGELOG.md index 8861dcac..55717a09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ - `haiku-rag serve` renamed to `haiku-rag mcp` (only MCP is left). `--mcp-port` renamed to `--port`. Update any `claude_desktop_config.json` from `["serve", "--mcp", "--stdio"]` to `["mcp", "--stdio"]`. - `document.metadata` now uses source-agnostic keys: `source_revision` (was `etag` — S3-only and never populated for FS, so periodic sweeps re-ingested every file) and `content_type` (was `contentType`, snake_case for consistency). The v0.50.0 startup migration rewrites existing documents. All four source adapters (FS, HTTP, S3, future WebDAV) now write their native revision (mtime_ns, ETag, etc.) under the same key, fixing the regression where FS sources never short-circuited on unchanged files. +- Ingester pollers skip their periodic sweep when the source already has queued or claimed jobs in the queue — saves the listing round-trip (`PROPFIND` / `S3 LIST` / FS walk) when work is backed up. FS push events from `watchfiles` keep flowing during skipped sweeps. Visible in Logfire as `ingester.poller.sweep` spans with `skipped=true reason=pending_work`. +- Ingester now drains in-flight jobs on `SIGINT` / `SIGTERM` up to `workers.shutdown_grace_s` (default 60s) before cancelling. Cancelled jobs stay `claimed` and the reaper resets them on next start. Bonus: the pipeline no longer wraps `KeyboardInterrupt` / `SystemExit` / `CancelledError` as `TransientError` — those now propagate as intended. - Drop `list_documents` and `get_document` from the default RAG skill's tool set; the skill now exposes only `search` and `cite`. Both tools dumped unbounded content into the agent's context (full document lists, full document bodies) and `get_document` returned no chunk_ids so its output was structurally uncitable. The analysis skill already covers these uses programmatically — `await list_documents()` and `Path('/documents/{id}/content.txt').read_text()` inside `execute_code`. The tool branches remain in `create_skill_tools` and the `skill_generator` `AVAILABLE_TOOLS` set so users can still opt in when building custom skills. ## [0.48.1] - 2026-05-21 diff --git a/docs/ingester.md b/docs/ingester.md index dd439e9f..d267333e 100644 --- a/docs/ingester.md +++ b/docs/ingester.md @@ -152,6 +152,7 @@ ingester: poll_idle_interval_s: 1.0 claim_timeout_s: 1800 reaper_interval_s: 60 + shutdown_grace_s: 60 # SIGTERM drains in-flight up to this long retry: max_attempts: 5 base_delay_s: 2.0 @@ -168,6 +169,18 @@ extension, 4xx HTTP except 408/429, etc.) skips retry entirely. A reaper task resets jobs whose `claimed_at` is older than `claim_timeout_s` so a crashed worker doesn't strand its job. +**Backpressure.** Each poller skips its periodic sweep when its source +already has queued or claimed jobs in the queue. The unique-index dedup +would coalesce a re-sweep anyway; the skip saves the listing round-trip +(`PROPFIND` / `S3 LIST` / FS walk). FS push events from `watchfiles` +still flow during a skipped sweep, so new files aren't lost. + +**Graceful shutdown.** On `SIGINT` / `SIGTERM`, pollers stop immediately +and workers are given `shutdown_grace_s` to finish in-flight jobs. Jobs +still running after the grace window are cancelled — they stay +`claimed` in the queue and are reset by the reaper on the next start +once `claim_timeout_s` elapses. + **Per-source override.** A source can opt out of the global retry policy: diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index ac5768db..7d4fcc47 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -271,6 +271,12 @@ class WorkerConfig(BaseModel): claim_timeout_s: int = 1800 reaper_interval_s: int = 60 retry: RetryPolicyConfig = Field(default_factory=RetryPolicyConfig) + shutdown_grace_s: float = Field( + default=60.0, + description="On SIGINT/SIGTERM, how long to wait for in-flight jobs to " + "finish before forcing cancellation. Cancelled jobs stay 'claimed' in " + "the queue; the reaper resets them after claim_timeout_s.", + ) class APIConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 16447b17..cb11de50 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -103,7 +103,18 @@ class IngesterApp: if api_task is not None: await asyncio.gather(api_task, return_exceptions=True) await self._pollers.stop() - await self._pool.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, + ) finally: # Close the queue connection unconditionally. aiosqlite runs the # underlying sqlite3 in a background thread; leaving it open holds diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index 41ad66e0..b6cf1e20 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -83,13 +83,25 @@ class BasePoller: async def _sweep_once(self) -> bool: """One discover() sweep. Returns True on success, False if the - breaker is open or the sweep failed (and was recorded).""" + breaker is open, the source has pending work already queued, or the + sweep failed (and was recorded).""" if self._breaker.is_open: logger.debug( "Skipping discover() — circuit breaker open for %s", self.source_id ) return False with logfire.span("ingester.poller.sweep", source_id=self.source_id) as span: + if await self._jobs.has_pending(self.source_id): + # The unique index would dedupe a re-sweep into a saturated + # queue anyway; skipping saves the listing round-trip + # (PROPFIND / S3 LIST / FS walk) and keeps Logfire readable. + span.set_attribute("skipped", True) + span.set_attribute("skip_reason", "pending_work") + logger.debug( + "Skipping discover() — %s has pending work in the queue", + self.source_id, + ) + return False try: snapshot = await self._sync.get_snapshot(self.source_id) counts = { diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 61eaa8da..ba277e8f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -243,6 +243,21 @@ class JobRepo: rows = await cursor.fetchall() return [_row_to_job(r) for r in rows] + async def has_pending(self, source_id: str) -> bool: + """True iff at least one queued/claimed job exists for the source. + + Cheap probe used by pollers to skip sweeps when there's already + outstanding work — the queue's unique index would dedupe new enqueues + anyway, so a sweep into a saturated queue is pure wasted listing work. + """ + async with self._lock: + async with self._conn.execute( + "SELECT 1 FROM jobs WHERE source_id=? AND status IN ('queued','claimed') LIMIT 1", + (source_id,), + ) as cursor: + row = await cursor.fetchone() + return row is not None + async def counts_by_status(self) -> dict[str, int]: async with self._lock: async with self._conn.execute( diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index 66c50b0f..d2c52a0c 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -117,5 +117,8 @@ async def run_job(client: "HaikuRAG", job: Job) -> JobResult: revision=metadata.get("source_revision"), content_hash=metadata.get("md5"), ) - except BaseException as exc: + except Exception as exc: + # CancelledError, KeyboardInterrupt, SystemExit are BaseException + # subclasses; they signal the runtime is shutting us down, not a + # job-level failure, so we let them propagate untouched. raise _classify(exc) from exc diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index 8950aeb8..5f6f8210 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -167,16 +167,33 @@ async def test_delete_event_skipped_when_delete_orphans_false(fs_config, jobs, s @pytest.mark.asyncio -async def test_repeated_sweep_does_not_duplicate_jobs(fs_config, jobs, sync): - """Live-uniqueness: enqueueing the same (source_id, uri, op) when a job - is already queued/claimed is a no-op.""" +async def test_repeated_sweep_skipped_when_queue_has_pending(fs_config, jobs, sync): + """Backpressure: once a job is queued/claimed, the next sweep skips + discover() entirely instead of churning the listing operation.""" event = _event("file:///a.md", revision="r1") source = _StubSource("src", [[event], [event]]) poller = _periodic(source, fs_config, jobs, sync) + assert await poller._sweep_once() is True + assert source.discover_calls == 1 + # Second sweep: queue still has the live job → skip without calling discover. + assert await poller._sweep_once() is False + assert source.discover_calls == 1 + + +@pytest.mark.asyncio +async def test_sweep_resumes_after_queue_drains(fs_config, jobs, sync): + """Once the queue clears (success, dead, or cancel), sweeps resume.""" + event = _event("file:///a.md", revision="r1") + source = _StubSource("src", [[event], []]) + poller = _periodic(source, fs_config, jobs, sync) await poller._sweep_once() - await poller._sweep_once() - queued = await jobs.list_jobs(source_id="src") - assert len(queued) == 1 + + claimed = await jobs.claim_next("worker") + assert claimed is not None + await jobs.mark_succeeded(claimed.id) + + assert await poller._sweep_once() is True + assert source.discover_calls == 2 @pytest.mark.asyncio diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index a00157b9..5d1be9d4 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -124,6 +124,47 @@ async def test_enqueue_after_succeeded_succeeds(jobs): assert second is not None +@pytest.mark.asyncio +async def test_has_pending_returns_false_on_empty_queue(jobs): + assert await jobs.has_pending("src") is False + + +@pytest.mark.asyncio +async def test_has_pending_true_for_queued_job(jobs): + await jobs.enqueue("src", "u", JobOp.UPSERT) + assert await jobs.has_pending("src") is True + + +@pytest.mark.asyncio +async def test_has_pending_true_for_claimed_job(jobs): + await jobs.enqueue("src", "u", JobOp.UPSERT) + await jobs.claim_next("w") + assert await jobs.has_pending("src") is True + + +@pytest.mark.asyncio +async def test_has_pending_false_for_terminal_states(jobs): + await jobs.enqueue("src", "u-ok", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_succeeded(claimed.id) + + await jobs.enqueue("src", "u-bad", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(claimed.id, "permanent") + + assert await jobs.has_pending("src") is False + + +@pytest.mark.asyncio +async def test_has_pending_is_per_source(jobs): + """A backed-up source must not block an idle one.""" + await jobs.enqueue("busy", "u", JobOp.UPSERT) + assert await jobs.has_pending("busy") is True + assert await jobs.has_pending("idle") is False + + @pytest.mark.asyncio async def test_enqueue_different_ops_coexist(jobs): upsert = await jobs.enqueue("s", "u", JobOp.UPSERT) diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 1eed49e1..e6370fba 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -159,23 +159,42 @@ async def test_transient_error_at_max_attempts_marks_dead(client, jobs, sync, co @pytest.mark.asyncio async def test_unknown_exception_caught_and_marked_dead(client, jobs, sync): - # Pipeline classifies BaseException → TransientError, but if something - # slips past the pool catches with a final defensive net. - client.create_document_from_source.side_effect = KeyboardInterrupt("nope") + """An Exception subclass the pipeline classifier didn't recognise still + gets marked dead by the pool's defensive `except Exception` net.""" + + class _Weird(Exception): + pass + + client.create_document_from_source.side_effect = _Weird("surprise") job = await jobs.enqueue("src", "u", JobOp.UPSERT, max_attempts=1) assert job is not None pool = _pool( client, jobs, sync, retry_policy=RetryPolicy(base_delay_s=0.0, jitter=0.0) ) - # KeyboardInterrupt is a BaseException — pipeline wraps it to TransientError. - # With max_attempts=1, the worker marks dead. await pool.drain_once() refreshed = await jobs.get_job(job.id) assert refreshed is not None assert refreshed.status is JobStatus.DEAD +@pytest.mark.asyncio +async def test_keyboard_interrupt_propagates_not_classified(client, jobs, sync): + """KeyboardInterrupt / SystemExit / CancelledError signal runtime shutdown. + The pipeline must not wrap them — the job stays 'claimed' for the reaper.""" + client.create_document_from_source.side_effect = KeyboardInterrupt("ctrl-c") + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool(client, jobs, sync) + with pytest.raises(KeyboardInterrupt): + await pool.drain_once() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + + # --- start / stop lifecycle --- @@ -203,6 +222,67 @@ async def test_workers_drain_queue_after_start(client, jobs, sync): assert counts.get("succeeded", 0) == 5 +@pytest.mark.asyncio +async def test_shutdown_grace_lets_inflight_job_complete(client, jobs, sync): + """A short-running job in flight when stop() is called must finish before + the pool returns. Cancellation is the timeout path, not the default.""" + finished = asyncio.Event() + + async def _slow_then_finish(*args, **kwargs): + await asyncio.sleep(0.1) + finished.set() + return Document( + id="doc", + content="x", + uri="u", + metadata={"md5": "m", "source_revision": "e"}, + ) + + client.create_document_from_source.side_effect = _slow_then_finish + await jobs.enqueue("src", "u", JobOp.UPSERT) + + pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) + await pool.start() + # Yield long enough for the worker to claim and enter _process. + await asyncio.sleep(0.02) + await asyncio.wait_for(pool.stop(), timeout=5.0) + + assert finished.is_set() + counts = await jobs.counts_by_status() + assert counts.get("succeeded", 0) == 1 + + +@pytest.mark.asyncio +async def test_shutdown_grace_timeout_cancels_long_job(client, jobs, sync): + """When grace elapses, wait_for raises TimeoutError and the job stays + 'claimed' for the reaper to reset later.""" + cancelled = asyncio.Event() + + async def _hangs_forever(*args, **kwargs): + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cancelled.set() + raise + return Document(id="doc", content="x", uri="u") + + client.create_document_from_source.side_effect = _hangs_forever + await jobs.enqueue("src", "u", JobOp.UPSERT) + + pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) + await pool.start() + await asyncio.sleep(0.05) + + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.stop(), timeout=0.2) + + assert cancelled.is_set() + counts = await jobs.counts_by_status() + # Job was cancelled mid-_process before mark_succeeded/dead could fire, + # so it stays in 'claimed' for the reaper to pick up later. + assert counts.get("claimed", 0) == 1 + + @pytest.mark.asyncio async def test_double_start_raises(client, jobs, sync): pool = _pool(client, jobs, sync, worker_count=1)