From cf6caf14fe977ce60f6589c47470a9f612828ccb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 26 May 2026 12:38:55 +0300 Subject: [PATCH] Tighten retry/reap_stale guards; fix SourceSummary.type leak --- .../haiku/rag/ingester/api/routes/sources.py | 2 +- .../haiku/rag/ingester/api/schemas.py | 3 +- .../haiku/rag/ingester/queue/repository.py | 21 ++++++---- tests/ingester/test_api.py | 2 +- tests/ingester/test_queue.py | 41 +++++++++++++++++++ 5 files changed, 59 insertions(+), 10 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py index b26f2e3e..077d3a07 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py @@ -17,7 +17,7 @@ async def list_sources( summaries.append( SourceSummary( source_id=poller.source_id, - type=type(poller.config).__name__, + type=poller.config.type, last_polled_at=poller.last_polled_at, circuit_breaker_open=poller.is_circuit_open, last_skip_reason=poller.last_skip_reason, diff --git a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py index 4e1f5caa..55206c2f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import Literal from pydantic import BaseModel @@ -17,7 +18,7 @@ class HealthResponse(BaseModel): class SourceSummary(BaseModel): source_id: str - type: str + type: Literal["fs", "http", "s3", "webdav"] last_polled_at: datetime | None circuit_breaker_open: bool # Reason the most recent sweep attempt was skipped (e.g. "pending_work"), diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 2c7e57df..7fef15ee 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -177,8 +177,11 @@ class JobRepo: await self._conn.commit() async def retry(self, job_id: str) -> Job: - """Rescue a dead job: status='queued', attempts=0, error cleared. - Raises KeyError if the job doesn't exist.""" + """Reset a `dead` or `queued` job: status='queued', attempts=0, + error cleared, scheduled for immediate re-claim. Refuses `claimed` + rows (would race with the worker still processing) and `succeeded` + rows (re-ingest via UPSERT instead). Raises KeyError when the row + is missing or in a non-retryable state.""" now = _utcnow_iso() async with self._lock: async with self._conn.execute( @@ -191,7 +194,7 @@ class JobRepo: claimed_by=NULL, completed_at=NULL, scheduled_at=? - WHERE id=? + WHERE id=? AND status IN ('dead', 'queued') RETURNING * """, (now, job_id), @@ -199,7 +202,7 @@ class JobRepo: row = await cursor.fetchone() await self._conn.commit() if not row: - raise KeyError(f"Job {job_id!r} not found") + raise KeyError(f"Job {job_id!r} not found or not retryable") return _row_to_job(row) async def cancel(self, job_id: str) -> bool: @@ -335,8 +338,9 @@ class JobRepo: 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.""" + """Reset claimed jobs whose claimed_at is older than the timeout + back to `queued`. Decrements `attempts` to undo the increment from + `claim_next` — a crashed worker isn't a consumed attempt.""" threshold = ( datetime.now(UTC) - timedelta(seconds=claim_timeout_seconds) ).isoformat() @@ -344,7 +348,10 @@ class JobRepo: cursor = await self._conn.execute( """ UPDATE jobs - SET status='queued', claimed_at=NULL, claimed_by=NULL + SET status='queued', + claimed_at=NULL, + claimed_by=NULL, + attempts=MAX(0, attempts - 1) WHERE status='claimed' AND claimed_at < ? """, (threshold,), diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py index 2a39f531..67d08972 100644 --- a/tests/ingester/test_api.py +++ b/tests/ingester/test_api.py @@ -344,7 +344,7 @@ async def test_sources_lists_configured(tmp_path, jobs, sync): payload = resp.json() assert len(payload) == 1 assert payload[0]["source_id"] == "local" - assert payload[0]["type"] == "FSSourceConfig" + assert payload[0]["type"] == "fs" assert payload[0]["circuit_breaker_open"] is False diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index 152a0b46..470d69db 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -315,6 +315,43 @@ async def test_retry_unknown_raises(jobs): await jobs.retry("not-a-real-id") +@pytest.mark.asyncio +async def test_retry_refuses_claimed_job(jobs): + """Resetting a `claimed` row would race with the worker still + processing it: claim_next would re-claim and a second worker would + process the same URI.""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + assert claimed.status is JobStatus.CLAIMED + + with pytest.raises(KeyError): + await jobs.retry(job.id) + + # Row is untouched. + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + assert refreshed.attempts == 1 + + +@pytest.mark.asyncio +async def test_retry_refuses_succeeded_job(jobs): + """Succeeded rows should be re-ingested through the UPSERT path, not + re-run from the queue.""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_succeeded(claimed.id) + + with pytest.raises(KeyError): + await jobs.retry(job.id) + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.SUCCEEDED + + # --- cancel --- @@ -342,6 +379,7 @@ async def test_reap_stale_resets_old_claims(conn, jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None + assert claimed.attempts == 1 # Backdate claimed_at to simulate a crashed worker. long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() @@ -357,6 +395,9 @@ async def test_reap_stale_resets_old_claims(conn, jobs): assert refreshed.status is JobStatus.QUEUED assert refreshed.claimed_at is None assert refreshed.claimed_by is None + # claim_next incremented to 1; reap undoes that since a crashed worker + # isn't a consumed attempt — matches release_if_claimed semantics. + assert refreshed.attempts == 0 @pytest.mark.asyncio