diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 32614a8e5f..4e8915fb07 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -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) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index ef69fcfed3..bbf657b851 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -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, diff --git a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts index ca355274b8..0fe7949e2b 100644 --- a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts +++ b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts @@ -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) => diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 338679d26b..9802bb9d72 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -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) => diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index 60cfff8727..67d1cfae79 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -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 { + 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?.(); } diff --git a/studio/frontend/src/features/rag/components/ingestion-progress.tsx b/studio/frontend/src/features/rag/components/ingestion-progress.tsx index 89f5efafab..0e85d416aa 100644 --- a/studio/frontend/src/features/rag/components/ingestion-progress.tsx +++ b/studio/frontend/src/features/rag/components/ingestion-progress.tsx @@ -40,6 +40,14 @@ export function IngestionProgress({ ); } + if (event.type === "cancelled") { + return ( +
+ Cancelled +
+ ); + } + if (event.type === "complete") { const chunks = event.num_chunks; return ( diff --git a/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx index 68ce3628c6..54f8e1972f 100644 --- a/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx +++ b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx @@ -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 | 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() { )} - + {allDone ? ( + + ) : ( + + )} diff --git a/studio/frontend/src/features/rag/stores/index-progress-store.ts b/studio/frontend/src/features/rag/stores/index-progress-store.ts index fe616ce05e..e4ea4d6eaf 100644 --- a/studio/frontend/src/features/rag/stores/index-progress-store.ts +++ b/studio/frontend/src/features/rag/stores/index-progress-store.ts @@ -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; } 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; + cancelAll: () => Promise; clear: () => void; } @@ -42,7 +48,7 @@ function patch( }); } -export const useIndexProgressStore = create((set) => ({ +export const useIndexProgressStore = create((set, get) => ({ entries: {}, add: (id, filename) => set((s) => ({ @@ -57,5 +63,15 @@ export const useIndexProgressStore = create((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 => Boolean(c)); + await Promise.allSettled(handles.map((c) => c())); + set({ entries: {} }); + }, clear: () => set({ entries: {} }), })); diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index f20ba22656..6409da4992 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -328,7 +328,11 @@ export const useRagStore = create((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?.(); } },