Handle shutdown more gracefully, by stopping pollers and cancelling jobs after timeout. Skip periodic poll if a source has pending jobs

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

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:
 

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):

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

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 = {

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(

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

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):
+async def test_repeated_sweep_skipped_when_queue_has_pending(fs_config, jobs, sync):
-    """Live-uniqueness: enqueueing the same (source_id, uri, op) when a job
+    """Backpressure: once a job is queued/claimed, the next sweep skips
-    is already queued/claimed is a no-op."""
+    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")
+    claimed = await jobs.claim_next("worker")
-    assert len(queued) == 1
+    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

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)

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
+    """An Exception subclass the pipeline classifier didn't recognise still
-    # slips past the pool catches with a final defensive net.
+    gets marked dead by the pool's defensive `except Exception` net."""
-    client.create_document_from_source.side_effect = KeyboardInterrupt("nope")
+
+    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)