diff --git a/studio/backend/core/rag/db.py b/studio/backend/core/rag/db.py index f5fc7ccb98..1cc7dcb34e 100644 --- a/studio/backend/core/rag/db.py +++ b/studio/backend/core/rag/db.py @@ -73,9 +73,7 @@ def _backfill_fts(conn: sqlite3.Connection) -> None: """One-time: seed FTS from existing vectors for rag.db files created before FTS5 replaced the on-disk bm25s index. Runs only when FTS is empty but vectors exist; idempotent thereafter.""" - fts_seeded = conn.execute( - "SELECT 1 FROM rag_chunks_fts LIMIT 1" - ).fetchone() + fts_seeded = conn.execute("SELECT 1 FROM rag_chunks_fts LIMIT 1").fetchone() if fts_seeded is not None: return has_vectors = conn.execute("SELECT 1 FROM rag_vectors LIMIT 1").fetchone() diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py index 1a8c660b09..0ba1c52599 100644 --- a/studio/backend/core/rag/vector_store.py +++ b/studio/backend/core/rag/vector_store.py @@ -89,9 +89,7 @@ def upsert_chunks(scope: str, points: Iterable[dict]) -> None: rows, ) if fts_rows: - conn.executemany( - "DELETE FROM rag_chunks_fts WHERE chunk_id = ?", fts_delete - ) + conn.executemany("DELETE FROM rag_chunks_fts WHERE chunk_id = ?", fts_delete) conn.executemany( "INSERT INTO rag_chunks_fts (text, chunk_id, scope) VALUES (?, ?, ?)", fts_rows, diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 523cd3056a..7583dcc29f 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -255,7 +255,8 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async () => { + const refresh = useCallback(async (options?: { signal?: AbortSignal }) => { + const signal = options?.signal; setModelsError(null); try { const [listRes, statusRes, lorasRes] = await Promise.all([ @@ -264,6 +265,11 @@ export function useChatModelRuntime() { listLoras(), ]); + // Cancellation can land while the requests above are in flight (e.g. the + // user cancels a load during this refresh). Bail before writing any + // backend state back into the store -- cancelLoading already cleared it. + if (signal?.aborted) return; + setModels(listRes.models.map(toChatModelSummary)); setLoras(lorasRes.loras.map(toLoraSummary)); @@ -405,6 +411,7 @@ export function useChatModelRuntime() { }); } } catch (error) { + if (signal?.aborted) return; const message = error instanceof Error ? error.message : "Failed to load models"; setModelsError(message); @@ -766,7 +773,7 @@ export function useChatModelRuntime() { store.setParams({ ...store.params, ...p }); } } - await refresh(); + await refresh({ signal: abortCtrl.signal }); } catch (error) { // Skip rollback if user cancelled -- model is already being unloaded. if (abortCtrl.signal.aborted) throw error; @@ -1061,6 +1068,8 @@ export function useChatModelRuntime() { try { await performLoad(); + // User cancelled mid-refresh; cancelLoading handles teardown. + if (abortCtrl.signal.aborted) return; if (loadToastDismissedRef.current) { toast.success(`${toastDisplayName} loaded`, { classNames: MODEL_LOADED_TOAST_CLASSNAMES, diff --git a/tests/python/test_rag_bm25.py b/tests/python/test_rag_bm25.py index 303b47b9a0..f73ec58a6d 100644 --- a/tests/python/test_rag_bm25.py +++ b/tests/python/test_rag_bm25.py @@ -46,7 +46,9 @@ def test_lexical_search_roundtrip(isolated_rag_db): scope, [ _chunk("c1", "the quick brown fox jumps over the lazy dog", index = 0), - _chunk("c2", "machine learning models predict outputs from inputs", index = 1), + _chunk( + "c2", "machine learning models predict outputs from inputs", index = 1 + ), _chunk("c3", "fox terriers are small dogs", index = 2), ], )