diff --git a/studio/backend/core/rag/authorization.py b/studio/backend/core/rag/authorization.py
index 78df96b5ed..0dfc267333 100644
--- a/studio/backend/core/rag/authorization.py
+++ b/studio/backend/core/rag/authorization.py
@@ -16,7 +16,7 @@ import sqlite3
 
 from fastapi import HTTPException
 
-from storage.studio_db import get_connection
+from storage.studio_db import closing_connection
 
 _NOT_FOUND_DETAIL = "Document not found"
 
@@ -58,7 +58,7 @@ def document_for_subject_or_404(
     if not document_id or not current_subject:
         raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
 
-    with get_connection() as conn:
+    with closing_connection() as conn:
         row = conn.execute(
             "SELECT * FROM rag_documents WHERE id = ?",
             (document_id,),
diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py
index 27e5d54840..bd151855b2 100644
--- a/studio/backend/core/rag/ingestion.py
+++ b/studio/backend/core/rag/ingestion.py
@@ -20,7 +20,7 @@ from typing import Any
 from uuid import uuid4
 
 from loggers import get_logger
-from storage.studio_db import get_connection
+from storage.studio_db import closing_connection
 from utils.rag.config import (
     RAG_CHUNK_OVERLAP,
     RAG_CHUNK_SIZE,
@@ -475,7 +475,7 @@ def _update_job_row(job_id: str, **fields: Any) -> None:
     keys = list(fields.keys())
     set_clause = ", ".join(f"{k} = ?" for k in keys)
     values = list(fields.values()) + [job_id]
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute(f"UPDATE rag_ingestion_jobs SET {set_clause} WHERE id = ?", values)
         conn.commit()
 
@@ -486,7 +486,7 @@ def _update_document_row(document_id: str, **fields: Any) -> None:
     keys = list(fields.keys())
     set_clause = ", ".join(f"{k} = ?" for k in keys)
     values = list(fields.values()) + [document_id]
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute(f"UPDATE rag_documents SET {set_clause} WHERE id = ?", values)
         conn.commit()
 
@@ -554,7 +554,7 @@ def _insert_chunks_and_collect_for_bm25(
         )
         if kind in ("text", "caption") and meta["text"]:
             bm25_rows.append({"id": chunk_id, "text": meta["text"]})
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.executemany(
             """
             INSERT INTO rag_chunks
@@ -597,7 +597,7 @@ def _replace_document_pages(document_id: str, pages: list[dict]) -> None:
         )
         for page in pages
     ]
-    with get_connection() as conn:
+    with closing_connection() as conn:
         doc_row = conn.execute(
             "SELECT 1 FROM rag_documents WHERE id = ?",
             (document_id,),
@@ -639,7 +639,7 @@ def _all_scope_chunks(scope: str) -> list[dict]:
         bind = (thread_id,)
     else:
         return []
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(sql, bind).fetchall()
     return [{"id": r["id"], "text": r["text"]} for r in rows]
 
@@ -866,7 +866,7 @@ def enqueue_ingestion(
     else:
         logger.info("RAG ingest: figure captioning disabled for this upload")
     job_id = str(uuid4())
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute(
             """
             INSERT INTO rag_ingestion_jobs
@@ -961,7 +961,7 @@ def purge_thread_documents(thread_ids: list[str]) -> None:
 
     placeholders = ",".join("?" for _ in thread_ids)
     uploads_root = Path(os.path.realpath(rag_uploads_root()))
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             f"SELECT stored_path FROM rag_documents WHERE thread_id IN ({placeholders})",
             thread_ids,
@@ -984,7 +984,7 @@ def purge_thread_documents(thread_ids: list[str]) -> None:
 
 def purge_all_thread_documents() -> None:
     """Drop every per-thread RAG artifact."""
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             "SELECT DISTINCT thread_id FROM rag_documents WHERE thread_id IS NOT NULL"
         ).fetchall()
diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py
index db0dae319a..0ca02699da 100644
--- a/studio/backend/core/rag/locators.py
+++ b/studio/backend/core/rag/locators.py
@@ -12,7 +12,7 @@ from pathlib import Path
 from typing import Any
 
 from loggers import get_logger
-from storage.studio_db import get_connection
+from storage.studio_db import closing_connection
 
 from . import vector_store
 from .parsers import ParsedPage, parse
@@ -171,7 +171,7 @@ def _replace_document_pages(document_id: str, pages: list[ParsedPage]) -> None:
         )
         for index, page in enumerate(pages)
     ]
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute(
             "DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,)
         )
@@ -363,7 +363,7 @@ def backfill_document_locators(document_id: str, stored_path: Path) -> BackfillR
     pages = parsed.pages
     _replace_document_pages(document_id, pages)
 
-    with get_connection() as conn:
+    with closing_connection() as conn:
         doc_row = conn.execute(
             "SELECT kb_id, thread_id FROM rag_documents WHERE id = ?",
             (document_id,),
@@ -469,7 +469,7 @@ def backfill_document_locators(document_id: str, stored_path: Path) -> BackfillR
             }
 
     if sql_updates:
-        with get_connection() as conn:
+        with closing_connection() as conn:
             conn.executemany(
                 """
                 UPDATE rag_chunks
diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py
index 34089d6e2b..242446aac9 100644
--- a/studio/backend/core/rag/retrieval.py
+++ b/studio/backend/core/rag/retrieval.py
@@ -231,7 +231,13 @@ def retrieve_hybrid(
 
 
 def filter_by_min_score(hits: list[Hit], min_score: float) -> list[Hit]:
-    """Drop hits whose dense_score < min_score; BM25-only hits dropped too."""
+    """Apply the dense-similarity floor.
+
+    min_score is a cosine threshold, so it only gates hits that carry a
+    dense_score. BM25-only and figure-ref hits (dense_score is None) are matched
+    by a different signal that the cosine floor does not apply to, so they pass
+    through rather than being silently dropped when min_score is raised.
+    """
     if min_score <= 0.0:
         return hits
-    return [h for h in hits if h.dense_score is not None and h.dense_score >= min_score]
+    return [h for h in hits if h.dense_score is None or h.dense_score >= min_score]
diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py
index 054915ce24..da5c224cda 100644
--- a/studio/backend/core/rag/scope.py
+++ b/studio/backend/core/rag/scope.py
@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from storage.studio_db import get_connection, list_chat_settings
+from storage.studio_db import closing_connection, list_chat_settings
 from utils.rag.config import resolve_embedder
 
 RAG_DEFAULTS_KEY = "rag.defaults"
@@ -19,7 +19,7 @@ def resolve_scope_embedder(scope: str) -> str | None:
     """KB → kb.embedding_model; thread → per-thread/defaults/matrix. None = use default."""
     if scope.startswith("kb_"):
         kb_id = scope[len("kb_") :]
-        with get_connection() as conn:
+        with closing_connection() as conn:
             row = conn.execute(
                 "SELECT embedding_model FROM rag_knowledge_bases WHERE id = ?",
                 (kb_id,),
diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py
index 02cb443483..c2363b8156 100644
--- a/studio/backend/core/rag/tool.py
+++ b/studio/backend/core/rag/tool.py
@@ -159,7 +159,7 @@ def search_knowledge_base(
 
     from core.rag import retrieval
     from core.rag.vector_store import kb_scope, thread_scope
-    from storage.studio_db import get_connection
+    from storage.studio_db import closing_connection
 
     scope = kb_scope(scope_kb_id) if scope_kb_id else thread_scope(scope_thread_id)
     k = top_k if top_k is not None else default_top_k
@@ -225,7 +225,7 @@ def search_knowledge_base(
     lookup: dict[str, dict] = {}
     if chunk_ids:
         placeholders = ",".join("?" for _ in chunk_ids)
-        with get_connection() as conn:
+        with closing_connection() as conn:
             rows = conn.execute(
                 f"""
                 SELECT c.id AS chunk_id, c.text, c.page_number,
diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py
index a21f6109be..e6d8eaed99 100644
--- a/studio/backend/routes/rag.py
+++ b/studio/backend/routes/rag.py
@@ -47,7 +47,7 @@ from core.rag.locators import backfill_document_locators
 from core.rag.vector_store import kb_scope, thread_scope
 from loggers import get_logger
 from storage.studio_db import (
-    get_connection,
+    closing_connection,
     list_chat_settings,
     upsert_chat_settings_merge,
 )
@@ -220,7 +220,7 @@ def _row_to_document(row: Any) -> DocumentResponse:
 
 
 def _kb_or_404(kb_id: str) -> Any:
-    with get_connection() as conn:
+    with closing_connection() as conn:
         row = conn.execute(
             "SELECT * FROM rag_knowledge_bases WHERE id = ?",
             (kb_id,),
@@ -231,7 +231,7 @@ def _kb_or_404(kb_id: str) -> Any:
 
 
 def _thread_or_404(thread_id: str) -> None:
-    with get_connection() as conn:
+    with closing_connection() as conn:
         row = conn.execute(
             "SELECT id FROM chat_threads WHERE id = ?",
             (thread_id,),
@@ -241,7 +241,7 @@ def _thread_or_404(thread_id: str) -> None:
 
 
 def _document_or_404(document_id: str) -> Any:
-    with get_connection() as conn:
+    with closing_connection() as conn:
         row = conn.execute(
             "SELECT * FROM rag_documents WHERE id = ?",
             (document_id,),
@@ -309,7 +309,7 @@ def _start_ingestion(
     content_hash: str | None = None,
 ) -> UploadResponse:
     document_id = str(uuid4())
-    with get_connection() as conn:
+    with closing_connection() as conn:
         # Dedup: skip re-ingestion if the same content hash is already indexed
         # in this scope. Only 'completed' counts — failed/in-flight may retry.
         # Scope is the target kb_id or thread_id (a file in two KBs indexes in each).
@@ -398,7 +398,7 @@ def create_knowledge_base(
         payload.mode, payload.chunking_strategy
     )
     created_at = _now_ms()
-    with get_connection() as conn:
+    with closing_connection() as conn:
         try:
             conn.execute(
                 """
@@ -439,7 +439,7 @@ def create_knowledge_base(
 def list_knowledge_bases(
     current_subject: str = Depends(get_current_subject),
 ) -> KBListResponse:
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             "SELECT * FROM rag_knowledge_bases ORDER BY created_at DESC"
         ).fetchall()
@@ -677,7 +677,7 @@ def _reingest_scope(
 ) -> ReingestResponse:
     """Wipe scope artifacts and re-enqueue every document; metadata untouched."""
     scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id)  # type: ignore[arg-type]
-    with get_connection() as conn:
+    with closing_connection() as conn:
         if kb_id:
             rows = conn.execute(
                 "SELECT id, stored_path FROM rag_documents WHERE kb_id = ?",
@@ -757,7 +757,7 @@ def reingest_knowledge_base(
         else resolve_embedder(new_mode, new_strategy)
     )
 
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute(
             """
             UPDATE rag_knowledge_bases
@@ -825,7 +825,7 @@ def delete_knowledge_base(
     current_subject: str = Depends(get_current_subject),
 ) -> dict:
     _kb_or_404(kb_id)
-    with get_connection() as conn:
+    with closing_connection() as conn:
         doc_rows = conn.execute(
             "SELECT stored_path FROM rag_documents WHERE kb_id = ?",
             (kb_id,),
@@ -911,7 +911,7 @@ def list_kb_documents(
     current_subject: str = Depends(get_current_subject),
 ) -> DocumentListResponse:
     _kb_or_404(kb_id)
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             "SELECT * FROM rag_documents WHERE kb_id = ? ORDER BY created_at DESC",
             (kb_id,),
@@ -924,7 +924,7 @@ def list_thread_documents(
     thread_id: str,
     current_subject: str = Depends(get_current_subject),
 ) -> DocumentListResponse:
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             "SELECT * FROM rag_documents WHERE thread_id = ? ORDER BY created_at DESC",
             (thread_id,),
@@ -961,7 +961,7 @@ def delete_document(
 ) -> dict:
     row = _document_or_404(document_id)
     scope = kb_scope(row["kb_id"]) if row["kb_id"] else thread_scope(row["thread_id"])
-    with get_connection() as conn:
+    with closing_connection() as conn:
         conn.execute("DELETE FROM rag_documents WHERE id = ?", (document_id,))
         conn.commit()
     _unlink_if_under_uploads(Path(row["stored_path"]))
@@ -974,7 +974,7 @@ def list_thread_indexes(
     current_subject: str = Depends(get_current_subject),
 ) -> ThreadIndexListResponse:
     """List threads with >=1 RAG doc. LEFT JOIN keeps unpersisted threads (null title)."""
-    with get_connection() as conn:
+    with closing_connection() as conn:
         rows = conn.execute(
             """
             SELECT
@@ -1034,7 +1034,7 @@ async def job_events(
 ) -> StreamingResponse:
     state = ingestion.get_job_state(job_id)
     if state is None:
-        with get_connection() as conn:
+        with closing_connection() as conn:
             row = conn.execute(
                 "SELECT * FROM rag_ingestion_jobs WHERE id = ?",
                 (job_id,),
@@ -1462,7 +1462,7 @@ def get_document_preview_target(
     # TOCTOU window — if the chunk is deleted between the two calls, the fetch
     # returns None and the route 500s (D1.1). Cross-document collapses to the
     # same 404 — never 400 (would leak doc existence).
-    with get_connection() as conn:
+    with closing_connection() as conn:
         chunk_row = conn.execute(
             """
             SELECT id, chunk_index, page_number, text, kind, image_path,
@@ -1686,7 +1686,7 @@ def search(
     chunk_lookup: dict[str, dict] = {}
     if chunk_ids:
         placeholders = ",".join("?" for _ in chunk_ids)
-        with get_connection() as conn:
+        with closing_connection() as conn:
             rows = conn.execute(
                 f"""
                 SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text,
diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py
index 8f56d6fe06..9addf23932 100644
--- a/studio/backend/storage/studio_db.py
+++ b/studio/backend/storage/studio_db.py
@@ -16,6 +16,7 @@ import os
 import platform
 import sqlite3
 import threading
+from contextlib import contextmanager
 from datetime import datetime, timezone
 
 logger = logging.getLogger(__name__)
@@ -383,6 +384,24 @@ def get_connection() -> sqlite3.Connection:
     return conn
 
 
+@contextmanager
+def closing_connection():
+    """`get_connection()` that also closes the connection on exit.
+
+    Commits on success and rolls back on error like sqlite3's own context
+    manager, then always closes the connection. Use this instead of the bare
+    `with get_connection() as conn:` (which commits but never closes, leaning on
+    GC to release the handle) so connections are freed deterministically,
+    matching the explicit conn.close() convention used elsewhere in this module.
+    """
+    conn = get_connection()
+    try:
+        with conn:
+            yield conn
+    finally:
+        conn.close()
+
+
 def create_run(
     id: str,
     model_name: 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 2809f3cb04..c569d6df7d 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
@@ -54,13 +54,7 @@ export interface UseThreadDocUploadsResult {
  *  it was "off". */
 export function useThreadDocUploads(): UseThreadDocUploadsResult {
   const aui = useAui();
-  const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
   const [pendingDocs, setPendingDocs] = useState<PendingDoc[]>([]);
-  // Track scope key per chip so removeDoc dispatches the right delete
-  // (KB vs thread docs use different scope keys in the store).
-  const [chipScopeKeys, setChipScopeKeys] = useState<Record<string, string>>(
-    {},
-  );
 
   // Brand-new chats have no backend thread until the first message.
   // Initialize the local thread to mint a remoteId so RAG uploads can
@@ -105,10 +99,6 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
       };
       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 to
       // reset the index. Idempotent: both a late in-flight abort and the
@@ -186,7 +176,6 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
           scope = { kind: "thread", threadId };
           scopeKey = `thread:${threadId}`;
         }
-        setChipScopeKeys((m) => ({ ...m, [localChipId]: scopeKey as string }));
         const uploadDocument = useRagStore.getState().uploadDocument;
         const captionImages =
           useChatRuntimeStore.getState().ragCaptionImages;
@@ -298,27 +287,21 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
     [ensureThreadId],
   );
 
-  const removeDoc = useCallback(
-    (id: string) => {
-      setPendingDocs((prev) => {
-        const doc = prev.find((d) => d.id === id);
-        if (doc?.documentId) {
-          const scopeKey =
-            chipScopeKeys[id] ?? `thread:${activeThreadId ?? ""}`;
-          void useRagStore
-            .getState()
-            .deleteDocument(doc.documentId, scopeKey)
-            .catch(() => {});
-        }
-        return prev.filter((d) => d.id !== id);
-      });
-      setChipScopeKeys((m) => {
-        const { [id]: _gone, ...rest } = m;
-        return rest;
-      });
-    },
-    [activeThreadId, chipScopeKeys],
-  );
+  const removeDoc = useCallback((id: string) => {
+    // Route through the teardown thunk registered in addDoc: it aborts the
+    // upload, unsubscribes from the job SSE, releases the index slot, and
+    // deletes the backend doc using the scope key it closed over (kb vs
+    // thread). Deleting here directly leaked the slot and could target the
+    // wrong scope. Also drop the aggregate-toast entry for this file.
+    const entry = useIndexProgressStore.getState().entries[id];
+    if (entry?.cancel) {
+      void entry.cancel();
+      useIndexProgressStore.getState().remove(id);
+      return;
+    }
+    // Nothing was registered yet (no job/slot to release) — just drop the chip.
+    setPendingDocs((prev) => prev.filter((d) => d.id !== id));
+  }, []);
 
   const clearDocs = useCallback(() => setPendingDocs([]), []);
 
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 5d29fd5734..ae3888f95f 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -350,7 +350,6 @@ export function SharedComposer({
     return s.models.find((m) => m.id === checkpoint);
   });
   const aui = useAui();
-  const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
   const ragSource = useChatRuntimeStore((s) => s.ragSource);
   const setRagSource = useChatRuntimeStore((s) => s.setRagSource);
   const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
@@ -821,20 +820,19 @@ export function SharedComposer({
   }, []);
 
   const removePendingDoc = useCallback((id: string) => {
-    setPendingDocs((prev) => {
-      const doc = prev.find((d) => d.id === id);
-      if (doc?.documentId) {
-        void useRagStore
-          .getState()
-          .deleteDocument(
-            doc.documentId,
-            `thread:${activeThreadId ?? ""}`,
-          )
-          .catch(() => {});
-      }
-      return prev.filter((d) => d.id !== id);
-    });
-  }, [activeThreadId]);
+    // Route through the teardown thunk registered in addDoc: it aborts the
+    // upload, unsubscribes from the job SSE, releases the index slot, and
+    // deletes the backend doc with the scope key it closed over (kb vs
+    // thread). Deleting here directly leaked the slot and hardcoded the
+    // thread scope, mis-targeting KB-scoped docs. Also drop the toast entry.
+    const entry = useIndexProgressStore.getState().entries[id];
+    if (entry?.cancel) {
+      void entry.cancel();
+      useIndexProgressStore.getState().remove(id);
+      return;
+    }
+    setPendingDocs((prev) => prev.filter((d) => d.id !== id));
+  }, []);
 
   function clearStuckImeTimer() {
     if (stuckImeTimerRef.current) {
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 16e256510d..74ccc10a69 100644
--- a/studio/frontend/src/features/rag/stores/index-progress-store.ts
+++ b/studio/frontend/src/features/rag/stores/index-progress-store.ts
@@ -31,6 +31,7 @@ interface IndexProgressState {
   setReady: (id: string, chunks?: number) => void;
   setError: (id: string) => void;
   setCancel: (id: string, cancel: () => Promise<void> | void) => void;
+  remove: (id: string) => void;
   cancelAll: () => Promise<void>;
   clear: () => void;
 }
@@ -63,6 +64,13 @@ export const useIndexProgressStore = create<IndexProgressState>((set, get) => ({
     patch(set, id, { status: "ready", progress: 1, chunks }),
   setError: (id) => patch(set, id, { status: "error" }),
   setCancel: (id, cancel) => patch(set, id, { cancel }),
+  // Drop a single entry from the toast (e.g. the user removed one chip). The
+  // caller is responsible for tearing down that file's job/SSE/slot first.
+  remove: (id) =>
+    set((s) => {
+      const { [id]: _gone, ...rest } = s.entries;
+      return { entries: rest };
+    }),
   // Cancel every file in the batch (running, queued, finished) to restore the
   // pre-batch index state, then drop all toast entries.
   cancelAll: async () => {
diff --git a/tests/python/test_rag_retrieval.py b/tests/python/test_rag_retrieval.py
index f3ba4c62a4..61336428a0 100644
--- a/tests/python/test_rag_retrieval.py
+++ b/tests/python/test_rag_retrieval.py
@@ -16,9 +16,10 @@ def test_rrf_fuses_two_rankings():
     dense = [Hit("c", 0.9), Hit("b", 0.8), Hit("d", 0.5)]
     fused = _rrf_fuse([bm25, dense], rrf_k = 60, top_k = 3)
     ids = [h.chunk_id for h in fused]
-    # b ranks 2 in both -> highest fused score.
-    assert ids[0] == "b"
-    assert set(ids) == {"a", "b", "c"} or set(ids) == {"b", "c", "a"}
+    # c (bm25 rank2 1/63 + dense rank0 1/61 = 0.032266) narrowly beats
+    # b (rank1 in both = 2/62 = 0.032258); d falls outside top_k.
+    assert ids[0] == "c"
+    assert set(ids) == {"a", "b", "c"}
 
 
 def test_rrf_top_k_limits_output():
diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py
index e0e3bebd74..4d46f54004 100644
--- a/tests/python/test_rag_tool_handler.py
+++ b/tests/python/test_rag_tool_handler.py
@@ -53,7 +53,7 @@ def test_kb_takes_precedence_over_thread():
 
     captured = {}
 
-    def _stub_retrieve(scope, query, k):
+    def _stub_retrieve(scope, query, *args, **kwargs):
         captured["scope"] = scope
         return []
 
@@ -78,7 +78,7 @@ def test_thread_scope_when_only_thread_set():
 
     captured = {}
 
-    def _stub_retrieve(scope, query, k):
+    def _stub_retrieve(scope, query, *args, **kwargs):
         captured["scope"] = scope
         return []