diff --git a/CHANGELOG.md b/CHANGELOG.md index fa46810d..f0b89008 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,7 @@ - `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, 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. +- Ingester now drains in-flight jobs on `SIGINT` / `SIGTERM` up to `workers.shutdown_grace_s` (default 60s) before cancelling. Cancelled jobs release their claim back to `queued` (and decrement `attempts` since a cancel isn't a failure) so the next process picks them up immediately instead of waiting on the reaper's `claim_timeout_s`. Bonus: the pipeline no longer wraps `KeyboardInterrupt` / `SystemExit` / `CancelledError` as `TransientError` — those now propagate as intended. - `providers.docling_serve.base_url` now accepts a list. Jobs round-robin across the entries with each job's submit/poll/result pinned to one instance (task IDs are instance-local). The counter is per-process; for cross-process load balancing or failover, put an LB in front and pass a single URL here. - 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. diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index ba277e8f..f275848f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -266,6 +266,31 @@ class JobRepo: rows = await cursor.fetchall() return {row["status"]: row["n"] for row in rows} + async def release_if_claimed(self, job_id: str) -> bool: + """Reset a still-claimed job back to queued, immediately reclaimable. + Idempotent — a no-op if the job already transitioned to + succeeded/dead/rescheduled. Decrements attempts to undo the increment + from `claim_next`, since a cancellation isn't a failed attempt. + Returns True if the row was released.""" + now = _utcnow_iso() + async with self._lock: + async with self._conn.execute( + """ + UPDATE jobs + SET status='queued', + claimed_at=NULL, + claimed_by=NULL, + scheduled_at=?, + attempts=MAX(0, attempts - 1) + WHERE id=? AND status='claimed' + RETURNING id + """, + (now, job_id), + ) as cursor: + row = await cursor.fetchone() + await self._conn.commit() + return row is not None + async def reap_stale(self, claim_timeout_seconds: int) -> int: """Return claimed jobs whose claimed_at is older than the timeout to the queue. Used by the reaper to recover from crashed workers.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 3ff6875c..edbd1aef 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -118,6 +118,13 @@ class WorkerPool: logger.info("Processing %s %s (job %s)", job.op.value, job.uri, job.id) try: result = await run_job(self._client, job) + except asyncio.CancelledError: + # Graceful shutdown cancelled us mid-flight. Release the claim so + # the next process can pick the job up immediately instead of + # waiting on the reaper's claim_timeout_s. + await self._jobs.release_if_claimed(job.id) + logger.info("Job %s released back to queue on cancel", job.id) + raise except PermanentError as e: await self._jobs.mark_dead(job.id, str(e)) logger.info("Job %s dead (permanent): %s", job.id, e) diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index 5d1be9d4..a4daec74 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -372,6 +372,56 @@ async def test_reap_stale_leaves_fresh_claims_alone(jobs): assert refreshed.status is JobStatus.CLAIMED +# --- release_if_claimed --- + + +@pytest.mark.asyncio +async def test_release_if_claimed_resets_claimed_job_and_decrements_attempts(jobs): + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert job is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + assert claimed.attempts == 1 + + released = await jobs.release_if_claimed(job.id) + assert released is True + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + assert refreshed.claimed_at is None + assert refreshed.claimed_by is None + assert refreshed.attempts == 0 + + +@pytest.mark.asyncio +async def test_release_if_claimed_noop_on_already_queued(jobs): + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert job is not None + + released = await jobs.release_if_claimed(job.id) + assert released is False + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + assert refreshed.attempts == 0 + + +@pytest.mark.asyncio +async def test_release_if_claimed_noop_on_succeeded(jobs): + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert job is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_succeeded(job.id) + + released = await jobs.release_if_claimed(job.id) + assert released is False + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.SUCCEEDED + + # --- list / counts --- diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index e6370fba..37cf8045 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -253,9 +253,10 @@ async def test_shutdown_grace_lets_inflight_job_complete(client, jobs, sync): @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.""" +async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync): + """When grace elapses and the worker is cancelled mid-job, the claim is + released back to 'queued' so the next process can pick it up immediately + — no waiting on the reaper's claim_timeout_s.""" cancelled = asyncio.Event() async def _hangs_forever(*args, **kwargs): @@ -267,7 +268,8 @@ async def test_shutdown_grace_timeout_cancels_long_job(client, jobs, sync): return Document(id="doc", content="x", uri="u") client.create_document_from_source.side_effect = _hangs_forever - await jobs.enqueue("src", "u", JobOp.UPSERT) + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) await pool.start() @@ -277,10 +279,13 @@ async def test_shutdown_grace_timeout_cancels_long_job(client, jobs, sync): 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 + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + assert refreshed.claimed_by is None + # claim_next incremented attempts to 1; release_if_claimed decremented it + # back to 0 because a cancellation isn't a failed attempt. + assert refreshed.attempts == 0 @pytest.mark.asyncio