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

CLAUDE.md

@@ -248,12 +248,36 @@ Per-user file quota enforcement tracked via the `file_count` and `file_quota` co
 
 **How it works:**
 - `files.created_by` records which user uploaded each file (nullable for pre-migration files)
-- `users.file_count` is incremented/decremented in application code (in `PartitionFileManager`) — no SQL triggers
-- Decrements use `func.greatest(file_count - N, 0)` to prevent negative values from race conditions
+- `users.file_count` is incremented/decremented in application code — no SQL triggers
+- Decrements use `GREATEST(file_count - N, 0)` to prevent negative values from race conditions
 - `delete_partition` queries per-uploader counts before cascade delete, then bulk decrements
-- Quota check (`check_user_file_quota` in `openrag/api/dependencies/auth.py`) runs on upload, considering both indexed files and pending tasks
 
-**Quota logic (`file_quota` column):**
+**Atomic reserve/release (issue #664).** `file_count` is a **reserved + completed**
+counter, not a "completed files" counter. Admission (`check_user_file_quota` in
+`openrag/api/dependencies/auth.py`) charges a slot with one conditional UPDATE
+(`UserRepository.try_reserve_file_slot`) so concurrent uploads cannot all read the
+same pre-increment count and overshoot the quota. There is **no** completion-time
+increment — `add_file_to_partition` only consumes the existing reservation.
+
+Consequences to respect when touching this code:
+- The in-memory `TaskStateManager` pending count is **not** an admission input. It is
+  volatile (a restart zeroes it) and reserved uploads are already inside `file_count`.
+  Never add it back into a quota decision, and never add it to `file_count` when
+  reporting usage — that double-counts in-flight uploads.
+- A reservation has an **owner**. Before dispatch it is the request's: the
+  `check_user_file_quota` yield-teardown releases it unless the router calls
+  `commit_quota_reservation(...)`. After dispatch it is the worker's:
+  `IndexerWorker.process_file` releases it in a `finally` unless the catalog write
+  reports a new row. Both quota-gated routes (`add_file`, `copy_file`) reserve, and
+  so do the MCP tools that create rows (`MCPService.index_url`, `MCPService.copy_file`),
+  which have no dependency chain and therefore reserve inline;
+  `put_file` (replace re-index) does not, since it reuses an existing row.
+- Any new early return between admission and dispatch is automatically covered by the
+  teardown — but any new code path that *creates a file row without reserving*, or
+  *reserves without either committing or releasing*, leaks the counter. A leak is
+  silent and permanently narrows the user's quota.
+
+**Quota logic (`file_quota` column)** — the reserve SQL predicate reproduces exactly this:
 - `None` → use global default (`DEFAULT_FILE_QUOTA` env var, default `-1`)
 - `< 0` → unlimited
 - `>= 0` → specific limit

openrag/api/dependencies/auth.py

@@ -254,40 +254,90 @@ def require_admin_or_self(
     )
 
 
+class QuotaReservation:
+    """One reserved file slot, owned by the request that admitted it.
+
+    Issue #664: admission increments ``users.file_count`` *before* the
+    upload is dispatched, so the counter is now "reserved + completed"
+    rather than "completed". That makes the reservation a resource with an
+    owner — it must either be **consumed** (a ``files`` row is created for
+    it) or **released**.
+
+    The handover point is dispatch: until the job is queued the request
+    owns the slot and :func:`check_user_file_quota`'s teardown releases it
+    on any error; once ``commit()`` is called the indexing worker owns it
+    and is responsible for releasing on failure/cancellation. So the
+    router must call ``commit()`` — via :func:`commit_quota_reservation` —
+    at exactly the moment responsibility transfers, and never earlier.
+    """
+
+    __slots__ = ("user_id", "committed")
+
+    def __init__(self, user_id: int) -> None:
+        self.user_id = user_id
+        self.committed = False
+
+    def commit(self) -> None:
+        """Hand the slot off to the worker; teardown will not release it."""
+        self.committed = True
+
+
+def commit_quota_reservation(reservation: object) -> None:
+    """Commit ``reservation`` when it is a real one; no-op otherwise.
+
+    Routers receive whatever ``check_user_file_quota`` yields, and tests
+    routinely override that dependency with a plain stub. Guarding on the
+    type here keeps the routers free of ``if ... is not None`` noise and
+    keeps an overridden dependency from turning into an AttributeError.
+    """
+    if isinstance(reservation, QuotaReservation):
+        reservation.commit()
+
+
 async def check_user_file_quota(
     user=Depends(current_user),
     auth_service=Depends(get_auth_service),
-    job_service=Depends(get_job_service),
     config=Depends(get_config),
 ):
-    """Check if user has reached their file quota."""
-    default_file_quota = config.rdb.default_file_quota
-    if user.get("is_admin", False):
-        return user
-    user_quota = user.get("file_quota")
-    effective_quota = default_file_quota if user_quota is None else user_quota
-    if effective_quota < 0:
-        return user
-
+    """Atomically reserve one file slot for this upload, or reject it.
+
+    This is the admission gate for the two quota-bearing routes (``add_file``
+    and ``copy_file``). It replaces the pre-#664 read-then-check, which
+    compared the request's stale ``file_count`` snapshot plus an in-memory
+    pending-task count against the quota and admitted every racer.
+
+    The in-memory ``TaskStateManager`` count is deliberately **not** an
+    input any more: it is not durable (a restart zeroes it, reopening the
+    gate) and it is no longer needed — a reserved slot is already counted
+    in the durable ``file_count``.
+
+    Yields a :class:`QuotaReservation`; on the way out, an uncommitted
+    reservation is released. That covers every path between admission and
+    dispatch — a 409 duplicate, a rejected/oversize upload, workspace
+    validation, a dispatch error, or a client disconnect — because FastAPI
+    propagates the endpoint's exception into this generator.
+    """
     user_id = user.get("id")
-    pending_count = await job_service.get_user_pending_task_count(user_id)
-
-    logger.debug(
-        "User file quota check",
-        user_id=user_id,
-        pending_count=pending_count,
-    )
+    if user_id is None:
+        # No durable identity to charge (e.g. auth disabled). Nothing to
+        # reserve, so nothing can be enforced or leaked.
+        yield None
+        return
 
     try:
-        auth_service.validate_file_quota(
-            user,
-            pending_task_count=pending_count,
-            default_quota=default_file_quota,
+        new_count = await auth_service.reserve_file_slot(
+            user_id,
+            default_quota=config.rdb.default_file_quota,
         )
     except OpenRAGError as exc:
-        raise HTTPException(
-            status_code=exc.status_code,
-            detail=exc.message,
-        ) from exc
+        logger.bind(user_id=user_id).info("Upload rejected: file quota exceeded.")
+        raise HTTPException(status_code=exc.status_code, detail=exc.message) from exc
 
-    return user
+    logger.bind(user_id=user_id, file_count=new_count).debug("Reserved a file slot.")
+    reservation = QuotaReservation(user_id)
+    try:
+        yield reservation
+    finally:
+        if not reservation.committed:
+            logger.bind(user_id=user_id).debug("Releasing an unconsumed file-slot reservation.")
+            await auth_service.release_file_slot(user_id)

openrag/api/routers/admin/indexing.py

@@ -17,6 +17,7 @@
 
 from api.dependencies.auth import (
     check_user_file_quota,
+    commit_quota_reservation,
     current_user,
     current_user_partitions,
     ensure_partition_role,
@@ -182,7 +183,13 @@ async def add_file(
         original_filename=original_filename,
         user=user,
         workspace_ids=parsed_workspace_ids,
+        quota_reserved=True,
     )
+    # The job is queued: the worker now owns the reserved file slot and
+    # releases it if indexing fails or is cancelled. Committing before this
+    # point would leak the slot on a dispatch error; committing is skipped
+    # on every early return above, which releases it instead (#664).
+    commit_quota_reservation(_quota_check)
 
     return JSONResponse(
         status_code=status.HTTP_201_CREATED,
@@ -400,14 +407,20 @@ async def copy_file_between_partitions(
         partition_service=partition_service,
     )
 
-    await service.copy_file(
+    created = await service.copy_file(
         source_file_id=source_file_id,
         source_partition=source_partition,
         target_file_id=file_id,
         target_partition=partition,
         metadata=metadata,
         user=user,
     )
+    # Copy is synchronous, so the reservation is consumed here or not at all.
+    # ``created`` is False when no catalog row was written — the source had no
+    # chunks, or the target already existed (duplicate-at-catalog race) — in
+    # which case the slot goes back rather than leaking (#664).
+    if created:
+        commit_quota_reservation(_quota_check)
     return JSONResponse(
         status_code=status.HTTP_201_CREATED,
         content={"message": "File copied successfully."},

openrag/core/indexing/dispatcher.py

@@ -39,8 +39,14 @@ async def dispatch_indexing(
         embedder_name: str | None = None,
         require_existing_partition: bool = False,
         allow_legacy_require_existing_partition_retry: bool = False,
+        quota_reserved: bool = False,
     ) -> str:
-        """Queue an (re)indexing job, register its task state, return its id."""
+        """Queue an (re)indexing job, register its task state, return its id.
+
+        ``quota_reserved`` marks that one ``users.file_count`` slot was
+        charged at admission (#664) and now belongs to the job: the worker
+        releases it if the file never reaches the catalog.
+        """
         ...
 
     @abstractmethod
@@ -66,8 +72,13 @@ async def copy_file(
         metadata: dict,
         partition: str,
         user: dict | None,
-    ) -> None:
-        """Copy a file's chunks into another partition / file id."""
+    ) -> bool:
+        """Copy a file's chunks into another partition / file id.
+
+        Returns whether a catalog row was actually created (False for an
+        empty source or an already-existing target) — the caller needs it to
+        settle the reserved quota slot (#664).
+        """
         ...
 
     @abstractmethod

openrag/core/models/catalog.py

@@ -53,12 +53,32 @@ class DocumentRecord(BaseModel):
 
 
 class IndexationJob(BaseModel):
-    """An indexation job tracking batch document processing."""
+    """A durable indexation job — one row per dispatched indexing task.
+
+    The record mirrors what the in-memory ``TaskStateManager`` Ray actor holds,
+    but survives restarts and is operator-visible (issue #660). ``id`` is the
+    dispatcher's ``task_id``, so the durable row and the hot-cache entry share
+    one identity.
+
+    ``status`` reuses :class:`DocumentStatus` rather than :class:`JobStatus`
+    because a job here tracks exactly one file and must reproduce the actor's
+    state taxonomy verbatim (``QUEUED`` -> ``SERIALIZING`` -> ... ->
+    ``COMPLETED`` / ``FAILED`` / ``CANCELLED``); :class:`JobStatus` is the
+    coarser roll-up reserved for a future batch-level job entity.
+
+    ``error`` is bounded at write time (see
+    :func:`core.utils.text.truncate_error_text`) — a full traceback is
+    unbounded input and this row is retained.
+    """
 
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    status: JobStatus = JobStatus.QUEUED
-    total_documents: int = 0
+    status: DocumentStatus = DocumentStatus.QUEUED
     partition: str = "default"
+    file_id: str | None = None
+    user_id: int | None = None
+    job_metadata: dict[str, Any] = Field(default_factory=dict)
+    error: str | None = None
     created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
     started_at: datetime | None = None
     completed_at: datetime | None = None
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

openrag/core/ports/job_repo.py

@@ -3,10 +3,17 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from datetime import datetime
 from typing import Any
 
 from openrag.core.models.catalog import IndexationJob
 
+# The non-terminal states an indexation job passes through. Shared by the
+# repository implementations and the read models so "active"/"pending" means
+# exactly one thing across the stack.
+ACTIVE_JOB_STATES: tuple[str, ...] = ("QUEUED", "SERIALIZING", "CHUNKING", "INSERTING")
+TERMINAL_JOB_STATES: tuple[str, ...] = ("COMPLETED", "FAILED", "CANCELLED")
+
 
 class JobRepository(ABC):
     """CRUD operations for indexation jobs."""
@@ -18,7 +25,51 @@ async def create_job(self, job: IndexationJob) -> IndexationJob: ...
     async def get_job(self, job_id: str) -> IndexationJob | None: ...
 
     @abstractmethod
-    async def list_jobs(self, status: str | None = None, offset: int = 0, limit: int = 50) -> list[IndexationJob]: ...
+    async def list_jobs(
+        self,
+        status: str | None = None,
+        offset: int = 0,
+        limit: int = 50,
+        user_id: int | None = None,
+    ) -> list[IndexationJob]:
+        """Return jobs newest-first.
+
+        ``status`` accepts an exact state (case-insensitive) or the pseudo-status
+        ``"active"``, which expands to :data:`ACTIVE_JOB_STATES`. ``user_id``
+        scopes the result to one uploader; ``None`` means every job.
+        """
 
     @abstractmethod
     async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: ...
+
+    @abstractmethod
+    async def mark_failed_if_not_cancelled(self, job_id: str, *, error: str, completed_at: datetime) -> bool:
+        """Record a FAILED outcome unless the row is already CANCELLED.
+
+        Arbitration lives here rather than in the in-memory task actor because
+        the durable row is the only participant guaranteed to still exist. An
+        actor that restarted or evicted the entry cannot say whether the user
+        cancelled, and treating "I don't know" as "cancelled" leaves the job in a
+        non-terminal state that retention never sweeps — a permanent phantom in
+        the queue views, which is the failure #660 exists to remove.
+
+        Returns ``True`` if the row moved to FAILED.
+        """
+
+    @abstractmethod
+    async def count_by_status(self) -> dict[str, int]:
+        """Return ``{status: count}`` over every retained job.
+
+        Cheaper than paging the whole table for the queue-info roll-up.
+        """
+
+    @abstractmethod
+    async def purge_terminal_jobs(self, *, older_than_seconds: int, keep_last: int) -> int:
+        """Evict terminal jobs, returning how many rows were removed.
+
+        Retention is what keeps the durable store from repeating the unbounded
+        growth of the in-memory actor it replaced. Both bounds apply: a terminal
+        job is removed once it is older than ``older_than_seconds`` **or** once
+        it falls outside the ``keep_last`` most recently completed rows.
+        In-flight jobs are never touched.
+        """

openrag/core/ports/user_repo.py

@@ -51,6 +51,30 @@ async def delete_user(self, user_id: int) -> bool: ...
     @abstractmethod
     async def count_users(self) -> int: ...
 
+    # ── File-quota reserve / release ──────────────────────────────────
+    #
+    # ``users.file_count`` is a *reserved + completed* counter (issue #664).
+    # Admission reserves a slot with a single conditional UPDATE so N
+    # concurrent uploads can never overshoot the quota; whoever holds an
+    # unconsumed reservation must release it.
+
+    @abstractmethod
+    async def try_reserve_file_slot(self, user_id: int, *, default_quota: int) -> int | None:
+        """Atomically claim one file slot against the user's quota.
+
+        Returns the post-increment ``file_count`` when the slot was
+        granted, or ``None`` when the user is at (or over) quota — or does
+        not exist. Quota semantics: admins bypass; a ``NULL`` per-user
+        ``file_quota`` falls back to ``default_quota``; a *resolved* quota
+        ``< 0`` means unlimited.
+        """
+        ...
+
+    @abstractmethod
+    async def release_file_slot(self, user_id: int) -> None:
+        """Give back one reserved slot (clamped at zero). Idempotent-ish."""
+        ...
+
     # ── API keys ──────────────────────────────────────────────────────
 
     @abstractmethod