fix(quota): arbitrate the reserved file-slot release instead of leaking it

CLAUDE.md

@@ -272,6 +272,16 @@ Consequences to respect when touching this code:
   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.
+- **After dispatch, ownership is arbitrated, not assumed.** A task that
+  `ray.cancel` retires before `process_file`'s body runs never executes that
+  `finally`, so `WorkerDispatcher.cancel_task` releases instead. Both call
+  `TaskStateManager.claim_quota_release`, a one-shot compare-and-set in the
+  (single-threaded) state actor: the first caller wins and the other stands
+  down. Any new release path must go through the same claim — releasing
+  directly reintroduces the double-release, and skipping the release when the
+  claim is *lost* is correct, but skipping it when the claim cannot be *made*
+  is not. When the arbiter is unreachable, release: an undercount is
+  recoverable, a leak is permanent.
 - 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

openrag/di/container.py

@@ -584,6 +584,7 @@ def indexing_service(self) -> IndexingService:
                     workspace_repo=self.workspace_repo,
                     collection=settings.vectordb.collection_name,
                     job_repo=self.job_repo,
+                    user_repo=self.user_repo,
                 ),
                 config=settings,
                 partition_service=self.partition_service,

openrag/services/workers/dispatcher.py

@@ -63,6 +63,7 @@ def __init__(
         collection: str,
         timeout: float = DEFAULT_TIMEOUT,
         job_repo: Any = None,
+        user_repo: Any = None,
     ) -> None:
         self._pool = pool
         self._tsm = task_state_manager
@@ -76,6 +77,10 @@ def __init__(
         # is absent, job state degrades to the in-memory actor — the pre-#660
         # behaviour.
         self._job_repo = job_repo
+        # Optional for the same reason as ``job_repo``. Without it a cancel that
+        # beats the worker to the task cannot return the reserved slot, which is
+        # the pre-#664 behaviour rather than a new failure.
+        self._user_repo = user_repo
         self._last_job_purge_at: float | None = None
 
     async def _call(self, future: Any, task_description: str) -> Any:
@@ -95,6 +100,7 @@ async def _set_queued_details(
         partition: str,
         metadata: dict[str, Any],
         user_id: int | None,
+        quota_reserved: bool = False,
     ) -> bool:
         remote = _remote_actor_method(self._tsm, "set_queued_details")
         if remote is not None:
@@ -105,6 +111,7 @@ async def _set_queued_details(
                     partition=partition,
                     metadata=metadata,
                     user_id=user_id,
+                    quota_reserved=quota_reserved,
                 ),
                 task_description=f"set_queued_details({task_id})",
             )
@@ -121,6 +128,7 @@ async def _set_queued_details(
                 partition=partition,
                 metadata=metadata,
                 user_id=user_id,
+                quota_reserved=quota_reserved,
             ),
             task_description=f"set_details({task_id})",
         )
@@ -168,6 +176,9 @@ async def dispatch_indexing(
             partition=partition,
             metadata=user_metadata,
             user_id=user.get("id") if user else None,
+            # Tells the actor a slot is outstanding for this task, so
+            # ``claim_quota_release`` can arbitrate who gives it back.
+            quota_reserved=quota_reserved,
         )
         if not accepted:
             raise RuntimeError(
@@ -683,17 +694,53 @@ async def cancel_task(self, task_id: str) -> bool:
             # raises: the actor has already claimed the cancellation, and
             # leaving the worker running would contradict both records.
             ray.cancel(obj_ref["ref"], recursive=True)
-        # The reserved quota slot (#664) is deliberately not released here.
-        # A task cancelled mid-flight gives its slot back in
-        # ``IndexerWorker.process_file``'s ``finally``; releasing here too
-        # would double-release. But a task that ``ray.cancel`` retires
-        # *before* that body runs never executes the ``finally``, so its slot
-        # leaks -- and the CANCELLED row written just above is terminal, hence
-        # indistinguishable from a clean cancel. Recovering it needs the #676
-        # reconciliation to *recount* ``file_count`` (completed files + active
-        # job rows), not merely sweep orphaned active rows.
+        # Return the reserved quota slot (#664) if -- and only if -- this
+        # cancel is the reason nobody else will. A task cancelled mid-flight
+        # releases in ``IndexerWorker.process_file``'s ``finally``; a task that
+        # ``ray.cancel`` retires before that body runs never executes it, and
+        # used to leak the slot outright. ``claim_quota_release`` is a one-shot
+        # token in the state actor, so exactly one of the two wins and the
+        # outcome no longer depends on how far the worker got.
+        #
+        # After the release, not inside the ``finally`` above: the claim is
+        # only meaningful once ``ray.cancel`` has actually retired the task,
+        # and a failure here must not stop the worker from being killed.
+        await self._release_cancelled_slot(task_id)
         return True
 
+    async def _release_cancelled_slot(self, task_id: str) -> None:
+        """Give a cancelled task's reserved file slot back, if we own it.
+
+        Best-effort throughout: a cancellation the user already asked for and
+        already saw must not fail because quota bookkeeping did. A lost release
+        leaves ``file_count`` one too high, which the #676 reconciliation
+        recount is there to repair.
+        """
+        if self._user_repo is None:
+            return
+        try:
+            details = await self._call(
+                self._tsm.get_details.remote(task_id),
+                task_description=f"get_details({task_id})",
+            )
+            user_id = (details or {}).get("user_id")
+            if user_id is None:
+                return
+            claimed = await self._call(
+                self._tsm.claim_quota_release.remote(task_id),
+                task_description=f"claim_quota_release({task_id})",
+            )
+            if not claimed:
+                return
+            await self._user_repo.release_file_slot(user_id)
+        except Exception as exc:  # noqa: BLE001 - the cancellation itself already succeeded
+            logger.warning(
+                "Could not release the reserved file slot for a cancelled task; "
+                "the user file_count may be one too high until reconciliation.",
+                task_id=task_id,
+                error=str(exc),
+            )
+
 
 def from_ray_namespace(
     namespace: str = "openrag",
@@ -704,6 +751,7 @@ def from_ray_namespace(
     workspace_repo: Any,
     collection: str,
     job_repo: Any = None,
+    user_repo: Any = None,
 ) -> WorkerDispatcher:
     import ray
     from services.workers.indexer_pool import build_indexer_pool
@@ -717,6 +765,7 @@ def from_ray_namespace(
         collection=collection,
         timeout=timeout,
         job_repo=job_repo,
+        user_repo=user_repo,
     )
 
 

openrag/services/workers/indexer_actor.py

@@ -223,7 +223,7 @@ async def process_file(
             )
             raise
         finally:
-            if release_slot:
+            if release_slot and await claim_quota_release(self._tsm, task_id):
                 await release_quota_slot(self._user_repo, user_id)
         # The raw upload is purged (when configured) by the enclosing actor, not
         # here: cleanup must also cover failures that happen *before* this method
@@ -481,6 +481,28 @@ def _display_filename(path: str, metadata: dict[str, Any]) -> str:
     return Path(path).name
 
 
+async def claim_quota_release(tsm: Any, task_id: str) -> bool:
+    """Ask the state actor whether *we* are the ones who owe the slot back.
+
+    ``WorkerDispatcher.cancel_task`` can reach the same conclusion for the same
+    task, so the actor arbitrates with a one-shot token (#664). Only the winner
+    releases.
+
+    An unreachable actor answers ``True``: a slot was reserved and this is the
+    last code that knows it, so releasing risks an undercount while staying
+    silent risks a permanent leak. This branch prefers the recoverable failure,
+    which is also exactly the pre-arbitration behaviour.
+    """
+    try:
+        return bool(await tsm.claim_quota_release.remote(task_id))
+    except Exception:  # noqa: BLE001 - cleanup must never mask the indexing error
+        logger.warning(
+            "Could not arbitrate the reserved file slot release; releasing it anyway.",
+            task_id=task_id,
+        )
+        return True
+
+
 async def release_quota_slot(user_repo: Any, user_id: int | None) -> None:
     """Give the uploader's reserved file slot back, swallowing any error.
 

openrag/services/workers/indexer_pool.py

@@ -9,6 +9,7 @@
 import ray
 from services.workers.indexer_actor import (
     IndexerWorker,
+    claim_quota_release,
     delete_uploaded_file,
     mark_dispatch_orphan_failed,
     release_quota_slot,
@@ -69,8 +70,8 @@ def __init__(self) -> None:
         task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag")
         # Kept on the actor as well as handed to the worker: setup can fail
         # before the worker is ever entered, and that path has to settle the
-        # task itself (see ``process_file``).
-        self._task_state_manager = task_state_manager
+        # task itself and arbitrate the quota release (see ``process_file``).
+        self._tsm = task_state_manager
         pipeline = build_indexing_pipeline(
             parser=parser,
             chunker=chunker,
@@ -244,7 +245,11 @@ async def process_file(
                 # release too, and ``release_quota_slot`` swallows it, leaking the
                 # slot. Nothing local can fix that (the DB *is* the counter);
                 # recovering it needs the reconciliation sweep tracked in #676.
-                if quota_reserved:
+                #
+                # Routed through the same one-shot claim as every other release
+                # path: a cancellation *during* setup lands here and in
+                # ``cancel_task``, and without arbitration both would release.
+                if quota_reserved and await claim_quota_release(self._tsm, task_id):
                     await release_quota_slot(self._catalog_store.user_repo, (user or {}).get("id"))
                 # The task is QUEUED in both stores and ``IndexerWorker`` --
                 # the only writer of SERIALIZING/COMPLETED/FAILED -- is never
@@ -262,7 +267,7 @@ async def process_file(
                 # unreachable Postgres still leaves an evictable, terminal
                 # ``TaskInfo`` rather than a pinned one.
                 await mark_dispatch_orphan_failed(
-                    self._task_state_manager,
+                    self._tsm,
                     self._catalog_store.job_repo,
                     task_id,
                 )

openrag/services/workers/task_state.py

@@ -62,6 +62,12 @@ class TaskInfo:
     error: str | None = None
     details: dict[str, Any] = field(default_factory=dict)
     object_ref: ray.ObjectRef | None = None
+    # #664. ``quota_reserved`` records that admission charged one
+    # ``users.file_count`` slot for this task; ``quota_release_claimed`` is the
+    # one-shot token that decides *who* gives it back. Kept on the record rather
+    # than in ``details`` because ``details`` is surfaced in API responses.
+    quota_reserved: bool = False
+    quota_release_claimed: bool = False
 
 
 @ray.remote(concurrency_groups={"set": 1000, "get": 1000, "queue_info": 1000})
@@ -253,6 +259,7 @@ async def set_details(
         partition: str,
         metadata: dict[str, Any],
         user_id: int | None,
+        quota_reserved: bool = False,
     ) -> None:
         async with self.lock:
             info = self._live_task(task_id)
@@ -268,6 +275,7 @@ async def set_details(
                 metadata=metadata,
                 user_id=user_id,
             )
+            info.quota_reserved = quota_reserved
 
     @ray.method(concurrency_group="set")
     async def set_queued_details(
@@ -278,6 +286,7 @@ async def set_queued_details(
         partition: str,
         metadata: dict[str, Any],
         user_id: int | None,
+        quota_reserved: bool = False,
     ) -> bool:
         async with self.lock:
             info = await self._ensure_task(task_id)
@@ -291,12 +300,41 @@ async def set_queued_details(
                 metadata=metadata,
                 user_id=user_id,
             )
+            info.quota_reserved = quota_reserved
             if self._file_delete_fenced(partition=partition, file_id=file_id):
                 info.state = "CANCELLED"
                 return False
             info.state = "QUEUED"
             return True
 
+    @ray.method(concurrency_group="set")
+    async def claim_quota_release(self, task_id: str) -> bool:
+        """Hand the task's reserved file slot to exactly one releaser (#664).
+
+        Two parties can end up believing they owe the uploader a slot back:
+        ``IndexerWorker.process_file``'s ``finally``, and
+        ``WorkerDispatcher.cancel_task`` when ``ray.cancel`` retires the task
+        before that body ever runs. Letting both release double-counts; letting
+        neither leaks. This actor is single-threaded, so a compare-and-set here
+        settles it: the first caller wins, every later one is told to stand
+        down.
+
+        Returns ``True`` for a task the actor no longer knows. A slot was
+        reserved for *some* task and nothing else will give it back, and this
+        branch prefers an undercount (recoverable, and self-heals on the next
+        reconciliation) over a leak (permanent, and silently narrows the user's
+        quota). In practice it is unreachable: eviction is FIFO by settle time,
+        so a task that just settled is the last candidate to be dropped.
+        """
+        async with self.lock:
+            info = self.tasks.get(task_id)
+            if info is None:
+                return True
+            if not info.quota_reserved or info.quota_release_claimed:
+                return False
+            info.quota_release_claimed = True
+            return True
+
     @ray.method(concurrency_group="set")
     async def set_object_ref(self, task_id: str, object_ref: ray.ObjectRef) -> bool:
         async with self.lock: