fix: atomic file-quota reserve + durable indexation job state

[Ahmath-Gadji]

Two related fixes that share a root cause — indexing state living only in the in-memory TaskStateManager Ray actor — plus the follow-ups from review. The first two commits are the substance, quota first, then durability; the rest are fixes found reviewing them.

Closes #664
Closes #660

Rebased onto develop @ d450f14, which includes #685 (bulk cancellation) and #671 (milvus cleanup before catalog delete). Both rewrote paths this PR also rewrites, so each merge needed its own review: the #685 interaction is described in §13–15, and the #671 interaction — 14 overlapping files, including one collision that would have broken every replace re-index — in §16. A finding from the #685 review is split out into #700, stacked on this branch.

---

1. fix(quota): reserve file slots atomically at admission (#664)

The per-user file quota was a check/admit TOCTOU race: check_user_file_quota compared a stale users.file_count snapshot plus an in-memory pending count against the quota, and the durable count was only incremented at the end of indexing. Concurrent uploads all observed the same pre-increment state and all passed; a restart zeroed the pending count and reopened the gate.

Admission is now a single conditional UPDATE (try_reserve_file_slot), so the read and write happen under one row lock:

UPDATE users SET file_count = file_count + 1
 WHERE id = $1
   AND (is_admin
        OR COALESCE(file_quota, $2::int) < 0
        OR file_count < COALESCE(file_quota, $2::int))
RETURNING file_count

The predicate reproduces the documented resolved-quota semantics exactly. Note the COALESCE resolves the per-user quota against the global default before the < 0 test — the naive file_quota < 0 from the issue body would let a negative global default silently override an explicit per-user limit. No row returned means no slot.

file_count therefore becomes a reserved + completed counter, so the completion-time increment in add_file_to_partition is removed (it would double-count), and every reservation now needs an owner:

• Before dispatch the request owns it. check_user_file_quota is a yield dependency whose teardown releases anything uncommitted — covering the 409 duplicate, upload rejection, workspace validation, dispatch errors, and client disconnects.
• After dispatch the worker owns it. IndexerWorker.process_file releases in a finally — deliberately not an except, which cancellation would skip, since ray.cancel raises CancelledError (a BaseException). This covers indexing failure, cancellation, and the duplicate-at-catalog race.
• The MCP tools (index_url, copy_file) have no dependency chain, so they reserve inline. put_file (replace re-index) never reserves; it reuses an existing row.

The in-memory pending count is no longer a correctness input anywhere, and get_current_user_info stops adding it to file_count — reserved uploads are already counted, so that double-counted in-flight files.

Status code is 403, not the 429 the issue floated: validate_file_quota already raised FILE_QUOTA_EXCEEDED as 403 and the API tests assert it. Easy to flip if reviewers prefer 429.

Admins and unlimited users are incremented too. Removing the completion-time increment while reserving only for capped users would have frozen every admin's file_count at zero, and /users/info displays it. This keeps release symmetric for all callers.

2. feat(jobs): persist indexation job state in Postgres (#660)

Job state lived only in the in-memory actor: unbounded (insert-only, no TTL/cap), volatile (wiped on restart), and unpersisted (PgJobRepository was a stub whose every method raised, with no jobs table and zero writers).

• Migration d4e5f6a7b8c9 adds a jobs table (one row per dispatched task, id = task_id, status CHECK over the 7 states, three indexes). Guarded with table_exists/index_exists in both directions, because metadata.create_all() runs at startup and an unguarded CREATE would raise DuplicateTable. Constraints stay unnamed so the migration and create_all converge on the same Postgres-default names. No FK to partitions — a job row is a historical record and must outlive the partition it targeted.
• Real PgJobRepository replaces the stub. IndexationJob becomes a per-task record reusing DocumentStatus (matches the actor's taxonomy verbatim; JobStatus lacked SERIALIZING/CHUNKING/INSERTING/CANCELLED).
• Lifecycle writes: QUEUED from the dispatcher before submit (so a crash between submit and the worker's first write still leaves the job visible rather than silently in-flight), SERIALIZING/COMPLETED/FAILED from the worker, CANCELLED from cancel_task. All durable writes are best-effort by design — this is bookkeeping about the work, not the work; a Postgres blip must not fail a file that indexed correctly, nor mask the real exception on the failure path.
• Reads hit Postgres and fall back to the actor — on a failure or on an empty result (see §6). This is required, not optional: evicting terminal entries from the actor without it would make completed work vanish from the queue views.
• Bounding, both sides: terminal jobs swept by age (7d) + row cap (10k), throttled to once/5min and run inline from the dispatch path (no new background task to own across replicas); the actor evicts terminal tasks by cap (2000) + TTL (1h), never in-flight ones — their object_ref is not serializable. Stored tracebacks are truncated (4000 chars, tail kept).

Incidental fix: the worker was discarding set_failed_if_not_cancelled's return value. It now honours the verdict, so a durable FAILED cannot overwrite a cancellation the user already asked for and already saw.

On the in-memory pending count

get_user_pending_task_count was intentionally not routed through Postgres. A job row only leaves the active states when a worker writes a terminal transition, so a job orphaned by a crash would occupy the quota of that user permanently (retention sweeps terminal rows only) — the in-memory count is wrong in the opposite, self-healing direction. §1 makes this moot by removing the pending count from the quota decision entirely, so the combined design has no reader of it as a correctness input.

3–12. Review follow-ups

Found reviewing the two commits above; each is standalone and independently reviewable.

• fix(quota): release the reserved slot when the SERIALIZING write fails. The pre-flight set_state ran outside the try whose finally releases, so an unreachable state actor leaked the slot permanently — the request had already handed ownership over at dispatch.
• fix(quota): reserve for the MCP index_url / copy_file tools. They create files rows without passing through the dependency, so removing the completion-time increment made MCP files invisible on the way in but still counted on the way out — driving file_count below reality and handing back free slots on a repeatable loop. Measured at quota 5: upload 5 via HTTP, index 3 via MCP, delete those 3 → the user holds 5 files with 3 free slots.
• fix(jobs): fail closed when listing tasks for an anonymous non-admin. list_jobs(user_id=None) means "every job". Not reachable today, but the MCP _USER_ID ContextVar defaults to None.
• fix(jobs): treat an empty durable read as a miss, not an empty queue. The fallback only triggered when the repo raised — []/{} are not None, so zero rows read as authoritative. The actor is detached and survives the API restart that first deploys this, so every pre-cutover task is live in the cache and absent from jobs: /queue/tasks answered [] and /queue/info reported 0 active while workers were indexing. Both aggregate readers now fall through on empty; get_task_details keeps is not None, since for a single row None already is the miss.
• fix(workers): drop a late write instead of resurrecting an evicted task. set_error/set_details/set_object_ref went through _ensure_task, which recreates on miss; the recreated entry has state=None, never re-enters terminal_at, and is never evictable again — the unbounded growth Persist indexing job state to Postgres; retire the in-memory-only TaskStateManager (unbounded, wiped on restart) #660 exists to fix. Not reachable today (the dispatcher writes set_state first), but silent and permanent if it ever is.
• fix(jobs): upper-case a job status before it reaches the CHECK. _expand_status normalized case, _status_value did not; a lower-case status would violate ck_jobs_status, and best-effort writes would swallow it, freezing the job at its previous status.
• refactor(workers): use the project logger in indexer_actor.
• test(workers): pin the retention bounds and prove the throttle recovers. The shipped 7d/10k were asserted nowhere, and the throttle test also passed for an implementation that purges once per process and never again.
• test(repos): revive the delete-cascade file_count assertion. Dead on develop (create_legacy_user signature drift), and it guards the very counter this branch redefines. Fixing the signature alone would have made it pass vacuously (0 → 0) now that the insert no longer increments, so it reserves the slots at admission the way a real upload does and asserts 2 → 0.
• test(quota): pin that dispatch tells the worker the slot is reserved. The release paths were covered, but nothing asserted the signal that arms them: commit_quota_reservation only stops the request's teardown from releasing, while quota_reserved=True is what makes the worker take ownership. Dropping that keyword from the router leaves the rest of the unit suite green — teardown declines to release (committed) and the worker declines too (never told it owns one), so every failed upload would leak a slot in silence. It now fails exactly one test.
• docs: stop two comments overstating their guarantees — the setup-failure release in indexer_pool releases through the store whose init just failed; truncate_error_text bounds the retained text, not the returned string.
• fix(jobs): stop list_tasks truncating the queue silently. _LIST_LIMIT caps the page at 1000 against a retention cap of 10k, so the cap is reachable and the route returned a short list with nothing to say it was short. It now reads one row past the limit to tell an exactly-full page from a truncated one, drops the probe, and logs a warning. Real pagination is the proper fix and is left as follow-up.

13–15. Found reviewing the rebase onto #685

#685 landed a bulk-cancel UI and rewrote cancel_task to claim the cancellation before signalling ray.cancel. That is the same code path this branch makes durable, so the merge needed its own review.

• fix(workers): mark a claimed cancellation terminal so it can be evicted. set_cancelled_if_active (from Add bulk cancellation for active jobs #685) is the only state write that does not call _mark_terminal (this branch's eviction hook). Eviction is driven entirely off terminal_at, and nothing writes that task's state again — ray.cancel raises CancelledError, a BaseException that process_file's except Exception sails past. So every user-initiated cancel retained a TaskInfo, its user_index entry and its pinned object_ref for the lifetime of the detached actor: verbatim the unbounded growth Persist indexing job state to Postgres; retire the in-memory-only TaskStateManager (unbounded, wiped on restart) #660 exists to fix, and reachable in bulk from Add bulk cancellation for active jobs #685's cancel-all button. Neither change is wrong alone; the gap is the merge. test_cancelled_tasks_are_evictable missed it because it drives set_state(..., "CANCELLED"), a path no cancellation actually takes.
• fix(jobs): keep a cancellation sticky in the durable row too. mark_failed_if_not_cancelled arbitrates FAILED-vs-CANCELLED in SQL, but the worker's SERIALIZING/COMPLETED writes were blind UPDATEs racing the cancel's blind UPDATE — so only the direction where the cancel loses was covered. When the cancel wins: COMPLETED-over-CANCELLED leaves the actor and the table disagreeing permanently (and the two read paths answer differently for the same task); SERIALIZING-over-CANCELLED is worse, stranding the row non-terminal after ray.cancel has killed the only writer that could finish it — and retention sweeps terminal rows only, so it never ages out. update_job now carries the same guard, mirroring the stickiness TaskStateManager.set_state already enforces in memory.
• docs(quota): name the residual ambiguity at the submit boundary. A submit timeout leaves ownership of the reserved slot genuinely undecidable — the worker may or may not have started. Teardown releases, which is right for the common case and off by one the other way; both directions undercount rather than leak.

A fourth finding from the same review — the cancel-before-start slot leak, which #685's cancel-all makes reachable at scale — is fixed in #700, stacked on this branch. It is a change to the ownership model rather than a bug fix, so it is kept separate and reviewable on its own.

Verification

Re-run on the rebased branch:

• ruff check / ruff format --check clean; scripts/check_layer_imports.py → layer import guard: OK.
• Unit: 1914 passed, 0 failed.
• Integration against real Postgres: 120 passed, 1 failed. The 1 (test_create_is_idempotent_per_name — create_partition raises PARTITION_EXISTS where the test expects the conflict swallowed) is pre-existing, confirmed identical on clean origin/develop with the same DSN, and fixed separately in test(repos): assert create_partition rejects duplicates, as it does #691.
• On the "6 failed" reported in the previous revision: those were an artefact of the checkout, not the code. load_dotenv() (core/config/loader.py:271) is called with no argument, so it walks up from the CWD and takes the closest .env; a checkout without one of its own inherits a parent directory's, and an empty RERANKER_PROVIDER= there breaks a pydantic discriminator. A checkout with its own .env shadows it and is fully green. Unrelated to this PR either way, but it means test runs are not hermetic depending on checkout location — worth a separate look.
• The quota concurrency claim is measured, not asserted: against a real Postgres, 20 parallel admits at quota 1 and quota 5 grant exactly 1 and exactly 5, with contiguous returned counts. The regression test was checked against a read-then-check reimplementation of the reserve — it overshoots the quota by a nondeterministic margin, observed between 14 and all 20 of the 20 concurrent admits across runs and machines — i.e. the test reliably fails on the bug it guards, though not at a fixed count.
• Migration verified end-to-end against real Postgres: create_all → upgrade head does not raise, re-upgrade is a no-op, downgrade drops and re-upgrade re-creates, and the two schemas are byte-identical across columns, PK, FK, indexes and check constraints. Names converge on jobs_pkey / jobs_user_id_fkey. Single head (d4e5f6a7b8c9).
• Two guards were mutation-checked rather than assumed: dropping quota_reserved=True from the router, and dropping the status upper-casing in _status_value, each fail exactly one test.

Known gaps — tracked in #676

• No startup reconciliation. A hard process death between reserve and dispatch leaks a quota slot, and an orphaned job row stays active forever. There is now a durable record to reconcile against; the sweep itself is follow-up work.
  • Retention only sweeps terminal rows, so an orphaned row is never purged; and cancel_task returns early when the actor no longer holds an object_ref, so an operator cannot force one terminal by hand either. The active roll-up in /queue/info therefore drifts upward across such events until reconciliation lands.
• A cancel of a not-yet-started task leaks its reserved quota slot — fixed in fix(quota): arbitrate the reserved file-slot release instead of leaking it #700, stacked on this branch, and left out of this one so the ownership-model change stays separately reviewable. After dispatch the worker owns the slot and only releases it in process_file's finally, so if ray.cancel retires the task before that body runs the release never happens — and cancel_task cannot release itself without double-releasing a task cancelled mid-flight (whose finally does run). Because the cancel writes a terminal CANCELLED row, indistinguishable from a clean cancel, recovering the slot on this branch alone needs the Reconcile orphaned in-flight jobs on startup + live-Ray e2e coverage for the durable job/quota paths #676 reconciliation to recount file_count (completed files + active job rows), not merely sweep orphaned active rows.
  • Worth weighing when reviewing: Add bulk cancellation for active jobs #685's cancel-all fires one cancel per selected task, so on this branch the leak scales with a single click.
• A dispatch that partially succeeds (pool accepts the job, set_object_ref then raises) double-releases the slot: the worker consumes it and teardown releases it. Chosen deliberately here — under-counting is recoverable, whereas a leak permanently locks the user out. Also closed by fix(quota): arbitrate the reserved file-slot release instead of leaking it #700.
• No live-Ray e2e. The dispatch → worker → Postgres path is proven at unit and repo-integration level only.

Review notes

• JobService.list_tasks caps at 1000 rows and now says so when it truncates (§12). Filtering happens in SQL before the cap, so the cap cannot leak another user's jobs.
• JobStatus is now unused; left in place rather than change a public model export. IndexationJob.total_documents was dropped — it had no readers and the model is internal (it appears in no API response).
• The actor's TTL is swept lazily, only when some task settles — an idle deployment keeps its last terminal entries cached past the hour. Harmless (a terminal state is immutable, so a stale read is not a wrong read) and the cap is the real memory bound; noted because the comment used to claim otherwise.

---

16. Found reviewing the rebase onto #671

Rebased onto develop @ d450f149, which includes #671 (fix/658-milvus-cleanup-before-catalog-delete). #671 rewrote the indexing dispatch, cancellation and delete paths that this branch also rewrites — 14 files overlap — so, as with #685 in §13–15, the merge needed its own review rather than a textual resolution. Seven of the 33 commits conflicted; the decisions that were not mechanical are below.

The single most important one is §16.1: a naive merge is not merely untidy, it breaks every replace re-index.

16.1 _write_catalog_record — two incompatible meanings for one bool

The two branches gave the same return value different meanings:

• This branch: "was a new row created?" — deliberately False on a replace, because that is the verdict that decides whether the reserved quota slot was consumed (File-quota check/admit is a TOCTOU race (quota bypass under concurrent uploads) #664).
• fix(indexing): delete vectors before catalog rows #671: "did the catalog write land?" — True on either branch, because it added a fail-closed if not wrote_catalog: raise RuntimeError("Catalog row was not written after vector indexing") on that value, so a require_existing_partition rejection surfaces instead of silently no-opping.

Merged textually, replace keeps returning False and #671's new guard turns every re-index into a hard failure. Nothing in either test suite catches it: this branch's tests assert the return value directly, and #671's exercise the insert path.

The two questions are now separate. _write_catalog_record reports whether the write landed (#671's meaning, so the fail-closed guard is correct), and the quota verdict is derived at the call site from replace:

if not replace:
    release_slot = False

The duplicate-at-catalog race — the other case that used to return False — now reaches the same outcome by a different route: add_file_to_partition returns False for the loser, the fail-closed guard raises, and process_file's finally releases the slot on the way out. test_duplicate_at_catalog_releases_the_slot was updated to assert both halves (it raises and the slot comes back) so neither can regress silently, and test_a_replace_reindex_does_not_consume_the_reservation was rewritten against the new contract as test_a_replace_reindex_updates_the_row_without_creating_one.

Verified live, not just in tests: a PUT re-index over an existing file returns 202 and settles COMPLETED, with users.file_count unchanged.

16.2 4ec0a634 is obsoleted by #671, and its test contradicted the new behaviour

fix(quota): don't report a dispatch failure once the worker has started argued that once submit returns, the worker owns the reserved slot and will run to completion — so a failing set_object_ref must not fail the dispatch, or the router skips commit_quota_reservation and teardown releases a slot the worker also owns.

#671 invalidated the premise. dispatch_indexing now rolls the dispatch back: _cancel_submitted_task cancels the worker and _cleanup_submitted_vectors sweeps what it wrote. The worker no longer runs to completion, so the error is truthful and must propagate — and #671's test_dispatch_indexing_cancels_worker_when_ref_registration_fails asserts exactly that.

Kept #671's behaviour and removed this branch's test_dispatch_does_not_fail_once_the_worker_has_started, which asserted the opposite of a test already on develop. The residual cost is a double release (the cancelled worker's finally and the request teardown both give the slot back), which under-counts rather than leaks — the direction this branch already chose deliberately (see Known gaps), and which #700 closes.

Worth noting the two signals from that one call are genuinely different and only one may fail the dispatch: a raised exception means the state actor is unreachable, while False is #671's deliberate delete-fence refusal. They happen to want the same handling here only because #671 rolls back in both cases.

16.3 The durable QUEUED write moves behind #671's admission gate

set_state(QUEUED) + set_details(...) were collapsed into a single atomic _set_queued_details(...) -> bool, which can now refuse: it returns False when a delete fence covers the file, and forces the in-memory state to CANCELLED.

The durable create_job(QUEUED) therefore moved after that gate. Writing it before would leave a row for a job that was never admitted — and since no worker ever runs, nothing settles it: non-terminal forever (retention sweeps terminal rows only) and counted active in /queue/info for good. Exactly the orphan class §16.5 exists to close, introduced by the rebase rather than fixed by it.

16.4 quota_reserved threads through #671's submit_kwargs, unconditionally

#671 hoisted the nine inline submit.remote(...) keywords into a submit_kwargs dict, and adds require_existing_partition to it only when true so a rolling-deploy actor without the kwarg still accepts the common case.

quota_reserved is passed unconditionally, like the nine base kwargs. The conditional pattern exists for rolling-deploy tolerance, and _is_legacy_require_existing_partition_rejection is hard-coded to that one kwarg name; generalising it for a flag that ships in the same release would be speculative. Flagging the asymmetry since it is deliberate: an actor old enough to reject quota_reserved is old enough to predate the whole quota model.

16.5 The dispatch-orphan fix lands in #671's failure path, not around submit

The fix from §16's sibling commit (settle a task that never reaches a worker) originally wrapped the submit call. On develop that call lives in _submit_indexing_task_once, which _submit_indexing_task retries for legacy actors — settling there would write a spurious FAILED for an attempt that is about to be retried successfully.

It moved into _mark_submit_failed, which #671 already calls on the failure path and which is gated on _cancel_submitted_task confirming the task really was cancelled. So the durable FAILED inherits that gate and a task that actually completed is never recorded failed in either store.

One improvement kept from this branch: #671's _mark_submit_failed drove its actor write through an unguarded self._call, so an unreachable state actor there replaced the exception the caller needs ("the pool is down") with one about the bookkeeping ("the actor is gone"). It is now guarded, matching how every other durable write on this branch behaves.

16.6 _live_task composes with #671's _record_details

b7e0b966 stops set_error / set_details / set_object_ref recreating an evicted task, while #671 refactored set_details onto a shared _record_details helper. Merged so set_details drops the write on a miss and still uses the shared helper. _ensure_task now has exactly two callers — set_state and #671's set_queued_details — and both are legitimate creation points for a task admission.

16.7 copy_file keeps both changes

#671 added strip_internal_metadata so _openrag* keys cannot leak into copied chunks; this branch made the return value the quota-commit verdict. Both kept. The empty-source path is an explicit return False rather than the bare return #671 left — under a -> bool signature that returned None, which is accidentally falsy but not something to rely on.

Not addressed — a gap this rebase opens

Cancellations from the delete path write no durable row. #671 added task_cancellation.cancel_active_indexing_tasks, which delete_file and delete_partition call to cancel every active task for a file or partition; it writes set_state(task_id, "CANCELLED") directly to the actor. That path did not exist at the merge base, so this branch does not mirror it to Postgres — a delete now silently cancels tasks with no durable record, and the two read paths disagree for those task ids until the actor evicts them.

Left out deliberately: it is a new writer to cover rather than a merge conflict to resolve, and _cancel_refs writes its state after ray.cancel and a settle-wait, which is the opposite ordering from cancel_task (§14). Making that path durable should be a deliberate decision about which ordering is right, not a rebase artefact. Tracked in #676 alongside the other reconciliation work.

Verification after the rebase

• ruff check / ruff format --check clean; layer import guard: OK.
• Unit: 1914 passed, 0 failed.
• Integration against real Postgres: 120 passed, 1 failed — test_create_is_idempotent_per_name, pre-existing and unrelated, fixed separately in test(repos): assert create_partition rejects duplicates, as it does #691.
• Re-ran the live stack on the rebased code: clean boot, durable job history survived an API restart, a new index settles COMPLETED, a replace re-index settles COMPLETED without consuming a slot (§16.1), and a cancel writes a terminal CANCELLED and returns the slot.

Summary by CodeRabbit

• New Features

  • Added durable indexation job history with lifecycle statuses, error details, filtering, queue counts, cancellation tracking, and retention.
  • Job status and details remain available after in-memory task data is evicted or services restart.
  • Added atomic file-quota reservations for uploads, URL indexing, and file copies.

• Bug Fixes

  • Prevented concurrent uploads from exceeding quotas or double-counting in-progress files.
  • Ensured reservations are released after failed, cancelled, duplicate, or incomplete operations.
  • Limited stored error details to prevent excessive diagnostic data growth.

• Documentation

  • Clarified quota usage, reservation ownership, and release behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request replaces pending-task quota checks with atomic durable reservations, adds reservation cleanup across upload and copy flows, implements durable Postgres indexation jobs with lifecycle and retention handling, and bounds the Ray task-state cache and stored error text.

Changes

File quota admission

Layer / File(s)	Summary
Quota contracts and atomic reservation CLAUDE.md, openrag/api/dependencies/auth.py, openrag/services/persistence/user_repo.py, openrag/services/orchestrators/auth_service.py, openrag/services/orchestrators/user_service.py	File counts now include reserved and completed files; admission uses conditional reservation, release is clamped, and pending tasks are informational.
Reservation propagation and cleanup openrag/api/routers/admin/indexing.py, openrag/services/orchestrators/mcp_service.py, openrag/services/workers/*, openrag/services/persistence/document_repo.py	Reservation ownership is propagated to indexing workers, committed after successful dispatch or catalog creation, and released on failure or duplicate paths.

Durable indexation state

Layer / File(s)	Summary
Durable job model and repository openrag/core/models/catalog.py, openrag/core/ports/job_repo.py, openrag/services/persistence/schema.py, openrag/services/persistence/job_repo.py, openrag/services/persistence/migrations/...	Single-task job records, lifecycle operations, status filtering, cancellation guards, retention purging, and error truncation are persisted in Postgres.
Durable dispatch and read paths openrag/services/workers/dispatcher.py, openrag/services/workers/indexer_actor.py, openrag/services/orchestrators/job_service.py, openrag/di/container.py	Dispatch, worker transitions, cancellation, failure settlement, queue reads, and task details use durable jobs with Ray fallback behavior.
Bounded task-state cache openrag/services/workers/task_state.py	Terminal task entries are evicted by cap or TTL, and late writes do not recreate evicted entries.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant JobService
  participant WorkerDispatcher
  participant PgJobRepository
  participant TaskStateManager
  Client->>WorkerDispatcher: dispatch_indexing
  WorkerDispatcher->>PgJobRepository: create QUEUED job
  WorkerDispatcher->>TaskStateManager: submit task state
  WorkerDispatcher-->>Client: task id
  Client->>JobService: list_tasks or get_task_details
  JobService->>PgJobRepository: read durable jobs
  PgJobRepository-->>JobService: job state and details
  JobService-->>Client: task information
  WorkerDispatcher->>PgJobRepository: update lifecycle or failure

Loading

Possibly related issues

• linagora/openrag#660 — Covers the durable Postgres job state, lifecycle persistence, retention, and bounded task-state work.
• linagora/openrag#676 — Covers reconciliation and end-to-end validation of quota reservation and job-state behavior.

Possibly related PRs

• linagora/openrag#671 — Both changes modify the copy_file flow in openrag/services/workers/dispatcher.py.
• linagora/openrag#700 — Both changes address reserved file-slot release arbitration and cleanup.

Suggested labels: fix

Suggested reviewers: hedhoud, enjoybacon7

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 35.36% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly captures the two main changes: atomic file-quota reservation and durable indexation job state.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/660-durable-job-state

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[EnjoyBacon7]

Review: two bugs fixed on the branch, two design gaps left for you

Reviewed both commits, ran the suites, and verified the claims in the description rather than taking them on trust. The core of this is solid — see the end for what I checked and confirmed. Four commits pushed to the branch (704d00f6..dcffc1fd); two findings below are not fixed because they're your design calls, not mine.

---

Fixed: MCP paths bypass the quota entirely (00da466e)

This is the one I'd have blocked a merge on.

Removing the completion-time increment assumes every new files row is preceded by a reservation. MCPService.index_url (mcp_service.py:623) and MCPService.copy_file (mcp_service.py:555) create rows without one — they aren't HTTP routes, so they never pass through the check_user_file_quota dependency.

It's worse than an unenforced quota, because the delete path still decrements unconditionally on created_by (document_repo.py:201, :225, :374). MCP files are invisible going in but counted going out, driving the counter below reality. Measured against a real Postgres, at file_quota = 5:

upload 5 via HTTP   -> file_count=5   (correctly at quota, 6th rejected)
index 3 via MCP     -> file_count=5   (3 real files invisible to the quota)
delete those 3      -> file_count=2   (count driven below reality)
=> user holds 5 files, has 3 free slots again. Repeatable => quota unbounded.

GREATEST(..., 0) doesn't save this — the count is driven down, not negative.

Both tools now reserve at admission and release anything unconsumed. index_url reserves before the download so an over-quota call costs no bandwidth; copy_file is synchronous, so its slot settles on the created verdict. I made auth_service/default_file_quota required constructor args rather than optional — a missed wiring would silently reopen exactly this hole.

Note this is the failure mode your own CLAUDE.md warns about ("any new code path that creates a file row without reserving … leaks the counter"), so I extended that section to name the MCP tools.

Fixed: quota slot leaks when the SERIALIZING write fails (7d0dd791)

process_file takes ownership of the slot at dispatch, but set_state.remote(task_id, "SERIALIZING") ran before the try whose finally releases it. TaskStateManager is a detached actor that can be unreachable (restart, node loss), so that await can raise — and nothing else covers the window:

• the request already handed ownership over (commit_quota_reservation runs once dispatch returns a task id), so its teardown won't release;
• IndexerWorkerActor.process_file's except BaseException only wraps the catalog/registry setup, not the worker call.

Result: the slot leaks permanently and narrows that user's quota until an admin fixes file_count by hand. Moved the pre-flight writes inside the try; regression test added. Your pool comment already names this exact window for the file-purge concern — the release just hadn't been extended to it.

---

Not fixed — your call

1. The durable FAILED write is gated on the volatile actor's verdict (indexer_actor.py, the if failed: gate)

set_failed_if_not_cancelled returns False for two different reasons (task_state.py:126-133): the task was cancelled (correct to skip), or info is None — the actor has no record of it. So if the actor is restarted or has evicted the entry between SERIALIZING and the exception, no durable FAILED is written and the row is stranded in SERIALIZING forever.

This inverts the stated design: the source of truth defers to the cache for a decision the cache is allowed to forget.

The clean fix is to move the arbitration into SQL rather than distinguish the two in the actor's return value — make the durable transition conditional:

UPDATE jobs SET status='FAILED', error=$2, completed_at=$3, updated_at=now()
 WHERE id=$1 AND status <> 'CANCELLED'

Then Postgres arbitrates the durable row and the actor arbitrates only its own cache, which is what "Postgres is the source of truth" should mean. I left this alone because test_failing_after_eviction_does_not_resurrect_the_task shows the no-resurrect behaviour is deliberate, and reworking your failure arbitration felt like your decision rather than a review fix.

2. Jobs stuck in active states are never purged — partly #676, but with a consequence worth calling out

purge_terminal_jobs only matches terminal states, and every durable write is best-effort (_record_job, _update_job both swallow). So a worker OOM, a lost node, or a Postgres blip at the terminal write leaves a row no retention pass can ever reclaim — contradicting the module docstring's "bounded rows … otherwise the durable store would just reproduce the in-memory leak on disk".

The extra consequence beyond #676: it permanently inflates get_queue_info's active count, which reads count_by_status from Postgres. This is precisely the hazard you documented as the reason to keep get_user_pending_task_count on the actor (job_service.py:171-179) — the same reasoning applies to get_queue_info but wasn't applied there. Worth either an age-based sweep for active rows or a note in #676.

Minor: _LIST_LIMIT = 1000 silently truncates an unpaginated route whose retention cap is 10k — no indicator to the caller that rows were dropped.

---

Also pushed

• dcffc1fd — list_tasks fails closed for an anonymous non-admin. list_jobs(user_id=None) means every job, and there's no guard for is_admin=False, user_id=None. Unreachable today (the HTTP path always resolves an id, MCP 403s first), but the MCP _USER_ID ContextVar's unset default is exactly None — the escalating value. One middleware-ordering change from a data leak.
• 704d00f6 — logging/naming consistency in indexer_actor: the two new helpers each invented an idiom (function-local get_logger import; stdlib logging.getLogger(__name__) with %s-formatting) where every other module under services/workers/ binds logger = get_logger() at import. Also dropped the underscore from release_quota_slot, since indexer_pool imports it — it's module API, and its neighbour delete_uploaded_file is already public.

Claims I verified

• The 6 unit / 2 integration failures are genuinely pre-existing and environmental. Confirmed both directions: the same 6 fail on a clean origin/develop checked out inside a directory whose parent holds a .env, and all pass on the same commit checked out outside it. Your load_dotenv() diagnosis is exactly right.
• The atomicity claim holds under harsher conditions than the test asserts. Re-derived independently with 50 concurrent racers on separate pool connections: quota 7 → exactly 7 granted, quota 1 → 1, quota 0 → 0, admin → all 50, unlimited → all 50, returned counts contiguous every time. The COALESCE-before-< 0 ordering is right, and the point about a negative global default not overriding an explicit per-user limit is a good catch.
• Migration chain: b7c1d2e3f4a5 really is the current head, single head after this revision, up/down idempotent, create_all → upgrade head converges.
• update_job's enumerate(..., start=2) parameter numbering is correct — dict ordering keeps assignments and updates.values() aligned.
• TaskStateManager eviction is sound under stress (50k tasks/100 users → caps hold at 2000, no dangling user_index, in-flight never evicted).

Final state: 1735 unit passed (+10 from the new tests), integration 111 passed, ruff/format/layer-guard clean, merges cleanly into current develop.

The acknowledged double-release on a partially-succeeded dispatch I left alone — agreed that under-counting beats a permanent lockout.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

openrag/services/workers/task_state.py (1)

300-311: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Return False for missing task registrations

set_object_ref can return None for an unknown or evicted task, but dispatch_indexing only treats False as a failed registration, so this falls through as success. Return False here to match the delete-fence branch and the -> bool contract.

🩹 Proposed fix

             info = self._live_task(task_id)
             if info is None:
-                return
+                return False

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/workers/task_state.py` around lines 300 - 311, Update
set_object_ref so the missing-task branch when _live_task(task_id) returns None
explicitly returns False, preserving the bool contract and ensuring
dispatch_indexing treats unknown or evicted registrations as failures.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@openrag/services/storage/postgres_store.py`:
- Around line 87-88: Move the job_repo property from the stub repositories
section into the Real repos section, alongside the initialized _job_repo in the
constructor. Keep the property implementation unchanged and update its
surrounding placement so the stub-repository header only covers actual stubs.

In `@tests/unit/services/workers/test_indexer_pool.py`:
- Around line 1256-1301: Add the `@pytest.mark.asyncio` decorator to both
test_setup_failure_releases_the_reserved_slot and
test_setup_failure_releases_nothing_when_no_slot_was_reserved so pytest-asyncio
awaits them in strict mode and executes their assertions.

---

Outside diff comments:
In `@openrag/services/workers/task_state.py`:
- Around line 300-311: Update set_object_ref so the missing-task branch when
_live_task(task_id) returns None explicitly returns False, preserving the bool
contract and ensuring dispatch_indexing treats unknown or evicted registrations
as failures.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 4233d284-ea88-47fd-b52c-af436a7b5629

📥 Commits

Reviewing files that changed from the base of the PR and between d450f14 and 4402d7e.

📒 Files selected for processing (42)

• CLAUDE.md
• openrag/api/dependencies/auth.py
• openrag/api/routers/admin/indexing.py
• openrag/core/indexing/dispatcher.py
• openrag/core/models/catalog.py
• openrag/core/ports/job_repo.py
• openrag/core/ports/user_repo.py
• openrag/core/utils/text.py
• openrag/di/container.py
• openrag/services/orchestrators/auth_service.py
• openrag/services/orchestrators/indexing_service.py
• openrag/services/orchestrators/job_service.py
• openrag/services/orchestrators/mcp_service.py
• openrag/services/orchestrators/user_service.py
• openrag/services/persistence/__init__.py
• openrag/services/persistence/document_repo.py
• openrag/services/persistence/job_repo.py
• openrag/services/persistence/migrations/alembic/versions/d4e5f6a7b8c9_add_jobs_table.py
• openrag/services/persistence/schema.py
• openrag/services/persistence/user_repo.py
• openrag/services/storage/postgres_store.py
• openrag/services/workers/dispatcher.py
• openrag/services/workers/indexer_actor.py
• openrag/services/workers/indexer_pool.py
• openrag/services/workers/task_state.py
• tests/integration/repos/conftest.py
• tests/integration/repos/test_job_repo.py
• tests/integration/repos/test_partition_repo.py
• tests/integration/repos/test_user_repo_quota_reserve.py
• tests/unit/api/dependencies/test_auth.py
• tests/unit/api/routers/admin/test_indexing_quota_release.py
• tests/unit/core/utils/test_text.py
• tests/unit/services/orchestrators/test_indexing_service.py
• tests/unit/services/orchestrators/test_job_service.py
• tests/unit/services/orchestrators/test_mcp_service.py
• tests/unit/services/orchestrators/test_user_service.py
• tests/unit/services/persistence/test_job_repo.py
• tests/unit/services/workers/test_dispatcher.py
• tests/unit/services/workers/test_indexer_pool.py
• tests/unit/services/workers/test_indexer_worker.py
• tests/unit/services/workers/test_indexer_worker_quota_release.py
• tests/unit/services/workers/test_task_state.py

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Mirror this move in the property section. The __init__ now correctly places _job_repo above the stubs block, but the job_repo property (Lines 178-180) still sits under the "Stub repos — methods raise StubRepositoryError…" header (Lines 173-176). That comment is now inaccurate for a durable repo and can mislead readers into thinking job_repo is unimplemented. Move the property up into the "Real repos" section.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/storage/postgres_store.py` around lines 87 - 88, Move the
job_repo property from the stub repositories section into the Real repos
section, alongside the initialized _job_repo in the constructor. Keep the
property implementation unchanged and update its surrounding placement so the
stub-repository header only covers actual stubs.

[coderabbitai]

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check pytest-asyncio mode; strict/legacy means unmarked async tests won't run.
rg -n 'asyncio_mode' -g 'pyproject.toml' -g 'pytest.ini' -g 'setup.cfg' -g 'tox.ini' || \
  echo "asyncio_mode not set → pytest-asyncio defaults to strict (markers required)"

Repository: linagora/openrag

Length of output: 230

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="tests/unit/services/workers/test_indexer_pool.py"

# Show the relevant range with line numbers.
sed -n '1225,1315p' "$file" | cat -n

# Look for module-level pytest markers in this file.
rg -n 'pytestmark|mark\.asyncio|pytest\.asyncio' "$file"

Repository: linagora/openrag

Length of output: 4561

---

Add @pytest.mark.asyncio to both setup-failure tests. With no asyncio_mode configured, pytest-asyncio stays in strict mode here, so these coroutines are collected but never awaited and the quota-release assertions never run.

🧰 Tools

🪛 ast-grep (0.44.1)

[info] 1275-1275: Do not hardcode temporary file or directory names
Context: "/tmp/doc.txt"
Note: [CWE-377] Insecure Temporary File.

(hardcoded-tmp-file)

---

[info] 1293-1293: Do not hardcode temporary file or directory names
Context: "/tmp/doc.txt"
Note: [CWE-377] Insecure Temporary File.

(hardcoded-tmp-file)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/services/workers/test_indexer_pool.py` around lines 1256 - 1301,
Add the `@pytest.mark.asyncio` decorator to both
test_setup_failure_releases_the_reserved_slot and
test_setup_failure_releases_nothing_when_no_slot_was_reserved so pytest-asyncio
awaits them in strict mode and executes their assertions.

[hedhoud]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4402d7e290

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Release the reservation when cancellation wins before work starts

When a user cancels a task while it is still queued (a common case when the pool is busy), ray.cancel can retire the Ray call before IndexerWorker.process_file enters its finally. The request already committed its reservation after dispatch, and this path deliberately does not release it, so each such cancellation permanently consumes one file_count slot and eventually blocks the user from uploading despite having no corresponding file.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Use the API catalog database for worker reservation releases

With an explicit settings.rdb.database, the API store preserves that configured database (di/repositories.py only derives a name when it is None), while IndexerWorkerActor still constructs _catalog_store using the collection-derived database (indexer_pool.py:89). Passing this worker-local user_repo here means failed/cancelled uploads release against a different database than the API admission increment, typically doing nothing and permanently leaking quota; the same split also leaves worker job-state updates unable to settle the API-created job rows.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Merge hot-cache tasks when durable history is only partial

If a durable job write is transiently lost (the dispatcher explicitly swallows such failures), while the jobs table already contains any historical row, this truthy result returns only Postgres jobs and omits the live actor-only task. Thus a successfully accepted upload can disappear from the task list during precisely the database degradation this fallback is intended to handle; get_queue_info has the same non-empty-table behavior for its active counts.

Useful? React with 👍 / 👎.