Add bulk cancellation for active jobs

openrag/api/routers/admin/indexing.py

@@ -30,6 +30,7 @@
     validate_metadata,
 )
 from api.routers.admin.task_logs import collect_task_logs
+from core.models.catalog import TERMINAL_TASK_STATES
 from core.utils.exceptions import OpenRAGError
 from core.utils.filename import sanitize_filename
 from core.utils.log_tail import app_log_file
@@ -557,5 +558,11 @@ async def cancel_task(
 ):
     cancelled = await service.cancel_task(task_id)
     if not cancelled:
+        task_state = await service.get_task_state(task_id)
+        if task_state in TERMINAL_TASK_STATES:
+            raise HTTPException(
+                status.HTTP_409_CONFLICT,
+                f"Task {task_id} is already {task_state.lower()} and cannot be cancelled",
+            )
         raise HTTPException(404, f"No ObjectRef stored for task {task_id}")
     return {"message": f"Cancellation signal sent for task {task_id}"}

openrag/core/indexing/dispatcher.py

@@ -82,8 +82,7 @@ async def get_task_error(self, task_id: str) -> str | None:
     async def cancel_task(self, task_id: str) -> bool:
         """Cancel a running/queued task.
 
-        Returns ``False`` when no object ref is stored for ``task_id``
-        (the caller maps that to a 404), ``True`` once the cancel signal
-        has been sent.
+        Returns ``False`` when the task cannot be cancelled, ``True`` once
+        the cancel signal has been sent.
         """
         ...

openrag/core/models/__init__.py

@@ -1,6 +1,6 @@
 """Domain models — pure Pydantic, no infrastructure imports."""
 
-from .catalog import DocumentRecord, DocumentStatus, IndexationJob, JobStatus
+from .catalog import TERMINAL_TASK_STATES, DocumentRecord, DocumentStatus, IndexationJob, JobStatus
 from .chunk import Chunk, ChunkType
 from .contextualization import ContextualizedQuery
 from .conversation import Conversation, Message
@@ -13,6 +13,7 @@
 from .workspace import Workspace
 
 __all__ = [
+    "TERMINAL_TASK_STATES",
     "ApiKey",
     "Chunk",
     "ChunkType",

openrag/core/models/catalog.py

@@ -20,6 +20,13 @@ class DocumentStatus(str, Enum):
     CANCELLED = "CANCELLED"
 
 
+# States a task/document cannot leave once reached. Single source of truth for
+# the cancellation guard shared by the indexing route, dispatcher, and
+# TaskStateManager — keep those in sync via this constant rather than
+# re-declaring the set.
+TERMINAL_TASK_STATES = frozenset({DocumentStatus.COMPLETED, DocumentStatus.FAILED, DocumentStatus.CANCELLED})
+
+
 class JobStatus(str, Enum):
     QUEUED = "QUEUED"
     RUNNING = "RUNNING"

openrag/services/workers/dispatcher.py

@@ -237,11 +237,19 @@ async def cancel_task(self, task_id: str) -> bool:
         if obj_ref is None:
             return False
 
-        ray.cancel(obj_ref["ref"], recursive=True)
-        await self._call(
-            self._tsm.set_state.remote(task_id, "CANCELLED"),
-            task_description=f"set_state({task_id})",
+        # Atomically claim the cancellation first: if ray.cancel() ran before
+        # this and the RPC then failed, a killed worker could never report
+        # back and the task would be stuck active forever (a zombie).
+        # Claiming first means a failed ray.cancel() just leaves a task
+        # marked CANCELLED that later self-corrects to COMPLETED/FAILED.
+        cancelled = await self._call(
+            self._tsm.set_cancelled_if_active.remote(task_id),
+            task_description=f"set_cancelled_if_active({task_id})",
         )
+        if not cancelled:
+            return False
+
+        ray.cancel(obj_ref["ref"], recursive=True)
         return True
 
 

openrag/services/workers/task_state.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 import ray
+from core.models.catalog import TERMINAL_TASK_STATES
 
 try:
     from core.config import load_config as _load_config
@@ -65,6 +66,15 @@ async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool:
             info.error = tb_str
             return True
 
+    @ray.method(concurrency_group="set")
+    async def set_cancelled_if_active(self, task_id: str) -> bool:
+        async with self.lock:
+            info = self.tasks.get(task_id)
+            if info is None or info.state in TERMINAL_TASK_STATES:
+                return False
+            info.state = "CANCELLED"
+            return True
+
     @ray.method(concurrency_group="set")
     async def set_details(
         self,

tests/unit/services/workers/test_dispatcher.py

@@ -61,6 +61,7 @@ def _workspace_repo() -> MagicMock:
 def _task_state_manager() -> MagicMock:
     tsm = MagicMock()
     tsm.set_state = _remote_mock()
+    tsm.set_cancelled_if_active = _remote_mock(True)
     tsm.set_details = _remote_mock()
     tsm.set_object_ref = _remote_mock()
     tsm.get_state = _remote_mock("SERIALIZING")
@@ -195,7 +196,6 @@ async def test_cancel_task_uses_stored_pool_object_ref() -> None:
     ref = object()
     tsm = _task_state_manager()
     tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref})
-    tsm.get_state.remote = AsyncMock(return_value="SERIALIZING")
     dispatcher = WorkerDispatcher(
         pool=_pool_with_ref(object()),
         task_state_manager=tsm,
@@ -209,18 +209,18 @@ async def test_cancel_task_uses_stored_pool_object_ref() -> None:
         result = await dispatcher.cancel_task("task-1")
 
     assert result is True
+    tsm.set_cancelled_if_active.remote.assert_called_once_with("task-1")
     cancel.assert_called_once_with(ref, recursive=True)
-    tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED")
 
 
 @pytest.mark.asyncio
-async def test_cancel_task_marks_cancelled_even_if_worker_finished_first() -> None:
+async def test_cancel_task_does_not_cancel_terminal_task() -> None:
     from services.workers.dispatcher import WorkerDispatcher
 
     ref = object()
     tsm = _task_state_manager()
     tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref})
-    tsm.get_state.remote = AsyncMock(return_value="COMPLETED")
+    tsm.set_cancelled_if_active.remote = AsyncMock(return_value=False)
     dispatcher = WorkerDispatcher(
         pool=_pool_with_ref(object()),
         task_state_manager=tsm,
@@ -230,11 +230,12 @@ async def test_cancel_task_marks_cancelled_even_if_worker_finished_first() -> No
         collection="default",
     )
 
-    with patch("ray.cancel"):
+    with patch("ray.cancel") as cancel:
         result = await dispatcher.cancel_task("task-1")
 
-    assert result is True
-    tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED")
+    assert result is False
+    tsm.set_cancelled_if_active.remote.assert_called_once_with("task-1")
+    cancel.assert_not_called()
 
 
 @pytest.mark.asyncio

ui/src/components/shared/data-table.test.tsx

@@ -119,4 +119,18 @@ describe("DataTable selection", () => {
     fireEvent.click(screen.getAllByRole("checkbox")[0]);
     expect(screen.queryByText(/selected/)).toBeNull();
   });
+
+  it("disables select-all when the visible rows cannot be selected", () => {
+    render(
+      <DataTable
+        columns={columns}
+        data={makeRows(2)}
+        enableSelection
+        canSelectRow={() => false}
+        getRowId={(r) => String(r.id)}
+      />,
+    );
+
+    expect(screen.getByRole("checkbox", { name: /select visible rows/i }).hasAttribute("disabled")).toBe(true);
+  });
 });

ui/src/components/shared/data-table.tsx

@@ -31,6 +31,8 @@ interface BaseDataTableProps<TData, TValue> {
   initialSorting?: SortingState;
   /** Render a leading checkbox column. */
   enableSelection?: boolean;
+  /** Optional row-level selection guard for pages with state-dependent bulk actions. */
+  canSelectRow?: (row: TData) => boolean;
   /** Stable row id (e.g. file_id) so selection survives a data refetch. */
   getRowId?: (row: TData) => string;
 }
@@ -54,6 +56,7 @@ export function DataTable<TData, TValue>({
   pageSize = 10,
   initialSorting = [],
   enableSelection = false,
+  canSelectRow,
   getRowId,
   rowSelection: controlledRowSelection,
   onRowSelectionChange,
@@ -71,20 +74,25 @@ export function DataTable<TData, TValue>({
     const selectColumn: ColumnDef<TData, TValue> = {
       id: "__select",
       enableSorting: false,
-      header: ({ table }) => (
-        <Checkbox
-          checked={
-            table.getIsAllPageRowsSelected() ||
-            (table.getIsSomePageRowsSelected() && "indeterminate")
-          }
-          onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
-          aria-label="Select visible rows"
-        />
-      ),
+      header: ({ table }) => {
+        const hasSelectablePageRows = table.getRowModel().rows.some((row) => row.getCanSelect());
+        return (
+          <Checkbox
+            checked={
+              table.getIsAllPageRowsSelected() ||
+              (table.getIsSomePageRowsSelected() && "indeterminate")
+            }
+            onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
+            disabled={!hasSelectablePageRows}
+            aria-label="Select visible rows"
+          />
+        );
+      },
       cell: ({ row }) => (
         <Checkbox
           checked={row.getIsSelected()}
           onCheckedChange={(value) => row.toggleSelected(!!value)}
+          disabled={!row.getCanSelect()}
           aria-label="Select row"
         />
       ),
@@ -103,7 +111,9 @@ export function DataTable<TData, TValue>({
     onColumnFiltersChange: setColumnFilters,
     onPaginationChange: setPagination,
     onRowSelectionChange: setRowSelection,
-    enableRowSelection: enableSelection,
+    enableRowSelection: enableSelection
+      ? (row) => (canSelectRow ? canSelectRow(row.original) : true)
+      : false,
     getRowId,
     // Keep the user on their current page when the rows change underneath them
     // (e.g. deleting a row refetches the list). TanStack resets pageIndex to 0

ui/src/pages/admin/jobs/list.test.tsx

@@ -2,10 +2,18 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { MemoryRouter } from "react-router-dom";
+import { toast } from "sonner";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { getQueueInfo, listTasks } from "@/lib/api/jobs";
+import { cancelTask, getQueueInfo, listTasks, type TaskState } from "@/lib/api/jobs";
 import JobListPage from "./list";
 
+vi.mock("sonner", () => ({
+  toast: {
+    success: vi.fn(),
+    error: vi.fn(),
+  },
+}));
+
 const permissions = vi.hoisted(() => ({ isAdmin: true }));
 
 vi.mock("@/lib/permissions", () => ({
@@ -14,15 +22,22 @@ vi.mock("@/lib/permissions", () => ({
   }),
 }));
 
-vi.mock("@/lib/api/jobs", () => ({
-  getQueueInfo: vi.fn(),
-  listTasks: vi.fn(),
-}));
+vi.mock("@/lib/api/jobs", async () => {
+  const actual = await vi.importActual<typeof import("@/lib/api/jobs")>("@/lib/api/jobs");
+  return {
+    ...actual,
+    cancelTask: vi.fn(),
+    getQueueInfo: vi.fn(),
+    listTasks: vi.fn(),
+  };
+});
 
+const cancelTaskMock = vi.mocked(cancelTask);
 const getQueueInfoMock = vi.mocked(getQueueInfo);
 const listTasksMock = vi.mocked(listTasks);
+const toastErrorMock = vi.mocked(toast.error);
 
-const task = (task_id: string, state: "COMPLETED" | "FAILED", filename: string) => ({
+const task = (task_id: string, state: TaskState, filename: string) => ({
   task_id,
   state,
   details: {
@@ -54,6 +69,7 @@ describe("JobListPage filters", () => {
   beforeEach(() => {
     permissions.isAdmin = true;
     vi.clearAllMocks();
+    cancelTaskMock.mockResolvedValue({ message: "Cancellation signal sent" });
     getQueueInfoMock.mockResolvedValue({
       workers: { total_slots: 4, pool_size: 2, max_per_actor: 2 },
       tasks: {
@@ -90,6 +106,70 @@ describe("JobListPage filters", () => {
     expect(await screen.findByText("failed.pdf")).not.toBeNull();
   });
 
+  it("bulk-cancels only selected active jobs", async () => {
+    listTasksMock.mockResolvedValue({
+      tasks: [
+        task("queued-task", "QUEUED", "queued.pdf"),
+        task("completed-task", "COMPLETED", "completed.pdf"),
+      ],
+    });
+
+    renderJobs();
+
+    expect(await screen.findByText("queued.pdf")).not.toBeNull();
+    expect(screen.getByText("completed.pdf")).not.toBeNull();
+
+    const rowCheckboxes = screen.getAllByRole("checkbox", { name: /select row/i });
+    expect(rowCheckboxes[0].hasAttribute("disabled")).toBe(false);
+    expect(rowCheckboxes[1].hasAttribute("disabled")).toBe(true);
+
+    await userEvent.click(screen.getByRole("checkbox", { name: /select visible rows/i }));
+    expect(screen.getByText("1 selected")).not.toBeNull();
+
+    await userEvent.click(screen.getByRole("button", { name: /cancel selected/i }));
+    await userEvent.click(screen.getByRole("button", { name: /confirm/i }));
+
+    await waitFor(() => expect(cancelTaskMock).toHaveBeenCalledWith("queued-task"));
+    expect(cancelTaskMock).not.toHaveBeenCalledWith("completed-task");
+  });
+
+  it("rechecks selected jobs before bulk cancelling", async () => {
+    let tasks = [task("queued-task", "QUEUED", "queued.pdf")];
+    listTasksMock.mockImplementation(async () => ({ tasks }));
+
+    renderJobs();
+
+    expect(await screen.findByText("queued.pdf")).not.toBeNull();
+
+    await userEvent.click(screen.getByRole("checkbox", { name: /select row/i }));
+    expect(screen.getByText("1 selected")).not.toBeNull();
+
+    tasks = [task("queued-task", "COMPLETED", "queued.pdf")];
+
+    await userEvent.click(screen.getByRole("button", { name: /cancel selected/i }));
+    await userEvent.click(screen.getByRole("button", { name: /confirm/i }));
+
+    await waitFor(() => expect(toastErrorMock).toHaveBeenCalledWith("1 task(s) were no longer active"));
+    expect(cancelTaskMock).not.toHaveBeenCalled();
+  });
+
+  it("clears selected jobs when the search changes", async () => {
+    listTasksMock.mockResolvedValue({
+      tasks: [task("queued-task", "QUEUED", "queued.pdf"), task("other-task", "QUEUED", "other.pdf")],
+    });
+
+    renderJobs();
+
+    expect(await screen.findByText("queued.pdf")).not.toBeNull();
+
+    await userEvent.click(screen.getAllByRole("checkbox", { name: /select row/i })[0]);
+    expect(screen.getByText("1 selected")).not.toBeNull();
+
+    await userEvent.type(screen.getByPlaceholderText("Search jobs..."), "other");
+
+    await waitFor(() => expect(screen.queryByText("1 selected")).toBeNull());
+  });
+
   it("shows queue and worker pressure from the backend", async () => {
     getQueueInfoMock.mockResolvedValue({
       workers: { total_slots: 2, pool_size: 1, max_per_actor: 2 },