ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: cancel RAG indexing from the toast and reset the batch

Browse source
This commit is contained in:
Roland Tannous 2026-05-29 11:02:54 +04:00
commit 68646a7abf
9 changed files with 281 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

46
studio/backend/core/rag/ingestion.py
View file

@ -428,6 +428,9 @@ class _JobState:
self.stage: str | None = None
self.progress: float = 0.0
self.error: str | None = None
self.cancelled = False
self.proc: Any = None
self.out_queue: Any = None
self.subscribers: list[queue_module.Queue[dict]] = []
self.lock = threading.Lock()
@ -664,6 +667,8 @@ def _pump(
try:
while True:
if state.cancelled:
break
try:
msg = out_queue.get(timeout = _QUEUE_TIMEOUT_SECONDS)
except queue_module.Empty:
@ -672,6 +677,9 @@ def _pump(
break
continue
mtype = msg.get("type")
if mtype == "__cancel__":
state.cancelled = True
break
if mtype == "progress":
state.stage = msg.get("stage")
state.progress = float(msg.get("progress", 0.0))
@ -733,6 +741,20 @@ def _pump(
proc.join(timeout = 5)
finished_at = int(time.time())
if state.cancelled:
# User cancelled mid-flight. The route-side deleteDocument removes the
# row, file, and chunk artifacts; here we just mark terminal and notify
# subscribers so the SSE stream closes cleanly.
_update_document_row(state.document_id, status = "cancelled")
_update_job_row(
state.job_id,
status = "cancelled",
stage = "cancelled",
finished_at = finished_at,
)
state.status = "cancelled"
state.push_event({"type": "cancelled"})
return
if final_status == "completed":
full_scope_chunks = _all_scope_chunks(state.scope)
bm25.rebuild_index(state.scope, full_scope_chunks)
@ -861,6 +883,7 @@ def enqueue_ingestion(
_jobs[job_id] = state
out_queue = _CTX.Queue()
state.out_queue = out_queue
proc = _CTX.Process(
target = _subprocess_worker,
args = (
@ -879,6 +902,7 @@ def enqueue_ingestion(
daemon = True,
)
proc.start()
state.proc = proc
pump_thread = threading.Thread(
target = _pump,
args = (state, proc, out_queue),
@ -889,6 +913,28 @@ def enqueue_ingestion(
return job_id
def cancel_ingestion(job_id: str) -> bool:
"""Stop an in-flight ingestion: wake the pump via a sentinel and kill the
worker subprocess so it stops consuming GPU/CPU. Returns False if the job
is unknown or already terminal. Artifact/row cleanup is the caller's job
(the route deletes the document)."""
state = get_job_state(job_id)
if state is None:
return False
if state.status in ("completed", "failed", "cancelled"):
return False
state.cancelled = True
if state.out_queue is not None:
try:
state.out_queue.put_nowait({"type": "__cancel__"})
except Exception:
pass
proc = state.proc
if proc is not None and proc.is_alive():
proc.terminate()
return True
def delete_document_artifacts(document_id: str, scope: str) -> None:
"""Drop the doc's vectors, rebuild BM25. Caller deletes the rag_documents row."""
vector_store.delete_document(scope, document_id)

11
studio/backend/routes/rag.py
View file

@ -1029,6 +1029,17 @@ def clear_thread_documents(
# --- Ingestion job SSE ---
@router.post("/jobs/{job_id}/cancel")
def cancel_job(
job_id: str,
current_subject: str = Depends(get_current_subject),
) -> dict:
"""Stop an in-flight ingestion job. The caller deletes the document
afterwards to reset the index; this only halts the worker."""
cancelled = ingestion.cancel_ingestion(job_id)
return {"ok": True, "cancelled": cancelled}
@router.get("/jobs/{job_id}/events")
async def job_events(
job_id: str,

93
studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts
View file

@ -5,7 +5,7 @@ import { useAui } from "@assistant-ui/react";
import { useCallback, useState } from "react";
import { toast } from "sonner";
import { subscribeToJobEvents } from "@/features/rag/api/rag-api";
import { cancelJob, subscribeToJobEvents } from "@/features/rag/api/rag-api";
import { useIndexProgressStore } from "@/features/rag/stores/index-progress-store";
import { useRagStore } from "@/features/rag/stores/rag-store";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
@ -86,6 +86,44 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
const addDoc = useCallback(
(file: File) => {
const localChipId = crypto.randomUUID();
// Lifecycle state shared between the upload flow and the cancel thunk.
// The cancel thunk closes over these `let`s by reference, so it always
// sees the latest job/document ids no matter when the user cancels.
const abort = new AbortController();
let jobId: string | undefined;
let documentId: string | undefined;
let scopeKey: string | null = null;
let unsubscribe: (() => void) | undefined;
let slotAcquired = false;
let slotReleased = false;
let cleaned = false;
const releaseSlot = () => {
if (slotAcquired && !slotReleased) {
slotReleased = true;
releaseIndexSlot();
}
};
const removeChip = () => {
setPendingDocs((prev) => prev.filter((d) => d.id !== localChipId));
setChipScopeKeys((m) => {
const { [localChipId]: _gone, ...rest } = m;
return rest;
});
};
// Stop the backend job (if started) and delete its document so the
// index resets. Idempotent: both a late in-flight abort and the toast
// cancel can reach here.
const cleanupBackend = async () => {
if (cleaned) return;
cleaned = true;
if (jobId) await cancelJob(jobId);
if (documentId && scopeKey) {
try {
await useRagStore.getState().deleteDocument(documentId, scopeKey);
} catch {}
}
};
setPendingDocs((prev) => [
...prev,
{ id: localChipId, file, status: "uploading" },
@ -94,32 +132,40 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
// whole batch) so the single toast counts queued files too.
const indexProgress = useIndexProgressStore.getState();
indexProgress.add(localChipId, file.name);
indexProgress.setCancel(localChipId, async () => {
abort.abort();
unsubscribe?.();
releaseSlot();
await cleanupBackend();
removeChip();
});
void (async () => {
// Hold an indexing slot for this document's whole lifecycle so bulk /
// folder uploads drain at the configured concurrency instead of
// spawning every ingestion at once. Released on every terminal path.
await acquireIndexSlot();
slotAcquired = true;
if (abort.signal.aborted) {
releaseSlot();
return;
}
indexProgress.setIndexing(localChipId);
let slotReleased = false;
const releaseSlot = () => {
if (!slotReleased) {
slotReleased = true;
releaseIndexSlot();
}
};
const ragSource = useChatRuntimeStore.getState().ragSource;
let scope:
| { kind: "kb"; kbId: string }
| { kind: "thread"; threadId: string }
| null = null;
let scopeKey: string | null = null;
if (ragSource.kind === "kb") {
scope = { kind: "kb", kbId: ragSource.kbId };
scopeKey = `kb:${ragSource.kbId}`;
} else {
// ragSource is "thread" or "off" — fall back to thread.
const threadId = await ensureThreadId();
if (abort.signal.aborted) {
releaseSlot();
return;
}
if (!threadId) {
setPendingDocs((prev) =>
prev.map((d) =>
@ -143,10 +189,21 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
setChipScopeKeys((m) => ({ ...m, [localChipId]: scopeKey }));
const uploadDocument = useRagStore.getState().uploadDocument;
try {
const { documentId, jobId, alreadyIndexed } = await uploadDocument(
scope,
file,
);
const {
documentId: did,
jobId: jid,
alreadyIndexed,
} = await uploadDocument(scope, file);
documentId = did;
jobId = jid;
if (abort.signal.aborted) {
// Cancelled while the upload was in flight: the document now
// exists on the backend, so tear it down here.
releaseSlot();
await cleanupBackend();
removeChip();
return;
}
if (alreadyIndexed) {
// Identical file already in this scope — no re-index. If a
// chip for this document already exists, drop the one we just
@ -154,14 +211,14 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
// otherwise mark this chip ready.
setPendingDocs((prev) => {
const dupExists = prev.some(
(d) => d.id !== localChipId && d.documentId === documentId,
(d) => d.id !== localChipId && d.documentId === did,
);
if (dupExists) {
return prev.filter((d) => d.id !== localChipId);
}
return prev.map((d) =>
d.id === localChipId
? { ...d, status: "ready", documentId }
? { ...d, status: "ready", documentId: did }
: d,
);
});
@ -179,11 +236,11 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
setPendingDocs((prev) =>
prev.map((d) =>
d.id === localChipId
? { ...d, status: "ingesting", jobId, documentId }
? { ...d, status: "ingesting", jobId: jid, documentId: did }
: d,
),
);
subscribeToJobEvents(jobId, {
unsubscribe = subscribeToJobEvents(jid, {
onEvent: (event) => {
if (event.type === "progress") {
indexProgress.setProgress(localChipId, event.progress);
@ -206,6 +263,8 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
}
indexProgress.setReady(localChipId, event.num_chunks);
releaseSlot();
} else if (event.type === "cancelled") {
releaseSlot();
} else if (event.type === "error") {
setPendingDocs((prev) =>
prev.map((d) =>

87
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -37,7 +37,7 @@ import {
} from "lucide-react";
import { Image03Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { subscribeToJobEvents } from "@/features/rag/api/rag-api";
import { cancelJob, subscribeToJobEvents } from "@/features/rag/api/rag-api";
import { useIndexProgressStore } from "@/features/rag/stores/index-progress-store";
import { useRagStore } from "@/features/rag/stores/rag-store";
import { acquireIndexSlot, releaseIndexSlot } from "./utils/rag-index-queue";
@ -587,6 +587,37 @@ export function SharedComposer({
const addDoc = useCallback(
(file: File) => {
const localChipId = crypto.randomUUID();
// Lifecycle state shared between the upload flow and the cancel thunk;
// the thunk closes over these `let`s so it sees the latest ids whenever
// the user cancels.
const abort = new AbortController();
let jobId: string | undefined;
let documentId: string | undefined;
let scopeKey: string | null = null;
let unsubscribe: (() => void) | undefined;
let slotAcquired = false;
let slotReleased = false;
let cleaned = false;
const releaseSlot = () => {
if (slotAcquired && !slotReleased) {
slotReleased = true;
releaseIndexSlot();
}
};
const removeChip = () => {
setPendingDocs((prev) => prev.filter((d) => d.id !== localChipId));
};
const cleanupBackend = async () => {
if (cleaned) return;
cleaned = true;
if (jobId) await cancelJob(jobId);
if (documentId && scopeKey) {
try {
await useRagStore.getState().deleteDocument(documentId, scopeKey);
} catch {}
}
};
setPendingDocs((prev) => [
...prev,
{ id: localChipId, file, status: "uploading" },
@ -595,19 +626,24 @@ export function SharedComposer({
// single toast counts queued files too.
const indexProgress = useIndexProgressStore.getState();
indexProgress.add(localChipId, file.name);
indexProgress.setCancel(localChipId, async () => {
abort.abort();
unsubscribe?.();
releaseSlot();
await cleanupBackend();
removeChip();
});
void (async () => {
// Hold an indexing slot for the document's whole lifecycle so bulk /
// folder uploads drain at the configured concurrency. Released on
// every terminal path below.
await acquireIndexSlot();
slotAcquired = true;
if (abort.signal.aborted) {
releaseSlot();
return;
}
indexProgress.setIndexing(localChipId);
let slotReleased = false;
const releaseSlot = () => {
if (!slotReleased) {
slotReleased = true;
releaseIndexSlot();
}
};
const ragSource = useChatRuntimeStore.getState().ragSource;
let scope:
| { kind: "kb"; kbId: string }
@ -615,8 +651,13 @@ export function SharedComposer({
| null = null;
if (ragSource.kind === "kb") {
scope = { kind: "kb", kbId: ragSource.kbId };
scopeKey = `kb:${ragSource.kbId}`;
} else {
const threadId = await ensureThreadId();
if (abort.signal.aborted) {
releaseSlot();
return;
}
if (!threadId) {
setPendingDocs((prev) =>
prev.map((d) =>
@ -635,26 +676,38 @@ export function SharedComposer({
return;
}
scope = { kind: "thread", threadId };
scopeKey = `thread:${threadId}`;
}
const uploadDocument = useRagStore.getState().uploadDocument;
try {
const { documentId, jobId, alreadyIndexed } = await uploadDocument(
scope,
file,
);
const {
documentId: did,
jobId: jid,
alreadyIndexed,
} = await uploadDocument(scope, file);
documentId = did;
jobId = jid;
if (abort.signal.aborted) {
// Cancelled while uploading: the document now exists on the
// backend, so tear it down here.
releaseSlot();
await cleanupBackend();
removeChip();
return;
}
if (alreadyIndexed) {
// Drop the just-added chip if this doc is already represented
// so the composer never shows the same document twice.
setPendingDocs((prev) => {
const dupExists = prev.some(
(d) => d.id !== localChipId && d.documentId === documentId,
(d) => d.id !== localChipId && d.documentId === did,
);
if (dupExists) {
return prev.filter((d) => d.id !== localChipId);
}
return prev.map((d) =>
d.id === localChipId
? { ...d, status: "ready", documentId }
? { ...d, status: "ready", documentId: did }
: d,
);
});
@ -672,11 +725,11 @@ export function SharedComposer({
setPendingDocs((prev) =>
prev.map((d) =>
d.id === localChipId
? { ...d, status: "ingesting", jobId, documentId }
? { ...d, status: "ingesting", jobId: jid, documentId: did }
: d,
),
);
subscribeToJobEvents(jobId, {
unsubscribe = subscribeToJobEvents(jid, {
onEvent: (event) => {
if (event.type === "progress") {
indexProgress.setProgress(localChipId, event.progress);
@ -696,6 +749,8 @@ export function SharedComposer({
}
indexProgress.setReady(localChipId, event.num_chunks);
releaseSlot();
} else if (event.type === "cancelled") {
releaseSlot();
} else if (event.type === "error") {
setPendingDocs((prev) =>
prev.map((d) =>

15
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -144,6 +144,7 @@ export type JobEvent =
}
| { type: "progress"; stage: string; progress: number }
| { type: "complete"; num_chunks: number }
| { type: "cancelled" }
| { type: "error"; error: string };
function parseErrorText(status: number, body: unknown): string {
@ -457,6 +458,14 @@ export async function prefetchRag(
// --- Ingestion SSE ---
/** Cancel an in-flight ingestion job. Best-effort: a 404/already-terminal job
* resolves without error so batch cancellation never throws on stale ids. */
export async function cancelJob(jobId: string): Promise<void> {
await authFetch(`/api/rag/jobs/${encodeURIComponent(jobId)}/cancel`, {
method: "POST",
}).catch(() => {});
}
/** Subscribe to a job's SSE stream; returns an unsubscribe fn.
* Use the EventSource polyfill so the bearer token rides in an
* Authorization header instead of leaking through URL query params. */
@ -485,7 +494,11 @@ export function subscribeToJobEvents(
try {
const parsed = JSON.parse(e.data) as JobEvent;
handlers.onEvent?.(parsed);
if (parsed.type === "complete" || parsed.type === "error") {
if (
parsed.type === "complete" ||
parsed.type === "error" ||
parsed.type === "cancelled"
) {
source.close();
handlers.onClose?.();
}

8
studio/frontend/src/features/rag/components/ingestion-progress.tsx
View file

@ -40,6 +40,14 @@ export function IngestionProgress({
);
}
if (event.type === "cancelled") {
return (
<div className={cn("text-xs text-muted-foreground", className)}>
Cancelled
</div>
);
}
if (event.type === "complete") {
const chunks = event.num_chunks;
return (

43
studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx
View file

@ -6,7 +6,7 @@ import { Progress } from "@/components/ui/progress";
import { Cancel01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { AnimatePresence, motion, useReducedMotion } from "motion/react";
import { useEffect, useRef } from "react";
import { useEffect, useRef, useState } from "react";
import { useIndexProgressStore } from "../stores/index-progress-store";
/** Single aggregate indexing toast (top-right). One entry per upload batch:
@ -18,8 +18,19 @@ const DISMISS_DELAY_MS = 4000;
export function IngestionToastStack() {
const entries = useIndexProgressStore((s) => s.entries);
const clear = useIndexProgressStore((s) => s.clear);
const cancelAll = useIndexProgressStore((s) => s.cancelAll);
const reduced = useReducedMotion();
const dismissTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
const [cancelling, setCancelling] = useState(false);
const onCancel = async () => {
setCancelling(true);
try {
await cancelAll();
} finally {
setCancelling(false);
}
};
const items = Object.values(entries);
const total = items.length;
@ -111,15 +122,27 @@ export function IngestionToastStack() {
</div>
)}
</div>
<Button
variant="ghost"
size="icon"
aria-label="Dismiss"
className="h-5 w-5 shrink-0 text-muted-foreground hover:text-foreground"
onClick={() => clear()}
>
<HugeiconsIcon icon={Cancel01Icon} size={12} />
</Button>
{allDone ? (
<Button
variant="ghost"
size="icon"
aria-label="Dismiss"
className="h-5 w-5 shrink-0 text-muted-foreground hover:text-foreground"
onClick={() => clear()}
>
<HugeiconsIcon icon={Cancel01Icon} size={12} />
</Button>
) : (
<Button
variant="ghost"
size="sm"
disabled={cancelling}
className="h-6 shrink-0 px-2 text-xs text-muted-foreground hover:text-destructive"
onClick={onCancel}
>
{cancelling ? "Cancelling…" : "Cancel"}
</Button>
)}
</div>
</motion.div>
</AnimatePresence>

18
studio/frontend/src/features/rag/stores/index-progress-store.ts
View file

@ -18,6 +18,10 @@ export interface IndexEntry {
progress: number;
/** Chunks this file produced (from the job's complete event); 0 until done. */
chunks: number;
/** Tear down this upload and remove its document from the index. Registered
* by the upload surface so the aggregate toast can cancel the whole batch
* without owning the per-file job/SSE/semaphore handles. */
cancel?: () => Promise<void> | void;
}
interface IndexProgressState {
@ -27,6 +31,8 @@ interface IndexProgressState {
setProgress: (id: string, progress: number) => void;
setReady: (id: string, chunks?: number) => void;
setError: (id: string) => void;
setCancel: (id: string, cancel: () => Promise<void> | void) => void;
cancelAll: () => Promise<void>;
clear: () => void;
}
@ -42,7 +48,7 @@ function patch(
});
}
export const useIndexProgressStore = create<IndexProgressState>((set) => ({
export const useIndexProgressStore = create<IndexProgressState>((set, get) => ({
entries: {},
add: (id, filename) =>
set((s) => ({
@ -57,5 +63,15 @@ export const useIndexProgressStore = create<IndexProgressState>((set) => ({
setReady: (id, chunks = 0) =>
patch(set, id, { status: "ready", progress: 1, chunks }),
setError: (id) => patch(set, id, { status: "error" }),
setCancel: (id, cancel) => patch(set, id, { cancel }),
// Cancel every file in the batch (running, queued, and already-finished) so
// the index returns to its pre-batch state, then drop all toast entries.
cancelAll: async () => {
const handles = Object.values(get().entries)
.map((e) => e.cancel)
.filter((c): c is NonNullable<typeof c> => Boolean(c));
await Promise.allSettled(handles.map((c) => c()));
set({ entries: {} });
},
clear: () => set({ entries: {} }),
}));

6
studio/frontend/src/features/rag/stores/rag-store.ts
View file

@ -328,7 +328,11 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
const unsubscribe = subscribeToJobEvents(jobId, {
onEvent: (event) => {
set((state) => ({ jobs: { ...state.jobs, [jobId]: event } }));
if (event.type === "complete" || event.type === "error") {
if (
event.type === "complete" ||
event.type === "error" ||
event.type === "cancelled"
) {
onComplete?.();
}
},