diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 6ebda4fe98..c257b04c7a 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -3357,18 +3357,34 @@ def _search_knowledge_base_with_budget( return "Error: knowledge base search cancelled." if deadline is not None and time.monotonic() >= deadline: return "Error: knowledge base search timed out." - if cancel_event is not None and cancel_event.is_set(): + + # Release the admission slot exactly once, whether the search finishes or the caller stops + # waiting on timeout/cancel. Freeing it as soon as the caller gives up keeps a slow or hung + # retrieval from holding the sole slot forever and starving every later lookup; the detached + # worker then finishes without touching the slot. + _slot_lock = threading.Lock() + _slot_released = False + + def release_slot() -> None: + nonlocal _slot_released + with _slot_lock: + if _slot_released: + return + _slot_released = True _RAG_SEARCH_SLOT.release() + + if cancel_event is not None and cancel_event.is_set(): + release_slot() return "Error: knowledge base search cancelled." if deadline is not None and time.monotonic() >= deadline: - _RAG_SEARCH_SLOT.release() + release_slot() return "Error: knowledge base search timed out." if timeout is None and cancel_event is None: try: return _search_knowledge_base(arguments, rag_scope) finally: - _RAG_SEARCH_SLOT.release() + release_slot() result: queue.Queue = queue.Queue(maxsize = 1) @@ -3378,17 +3394,19 @@ def _search_knowledge_base_with_budget( except BaseException as exc: result.put((False, exc)) finally: - _RAG_SEARCH_SLOT.release() + release_slot() try: threading.Thread(target = search, name = "rag-tool-search", daemon = True).start() except Exception: - _RAG_SEARCH_SLOT.release() + release_slot() raise while True: if cancel_event is not None and cancel_event.is_set(): + release_slot() return "Error: knowledge base search cancelled." if deadline is not None and time.monotonic() >= deadline: + release_slot() return "Error: knowledge base search timed out." wait = 0.05 if deadline is not None: diff --git a/studio/backend/tests/test_rag_retrieval.py b/studio/backend/tests/test_rag_retrieval.py index e630711c10..d0149357f7 100644 --- a/studio/backend/tests/test_rag_retrieval.py +++ b/studio/backend/tests/test_rag_retrieval.py @@ -243,6 +243,34 @@ def test_knowledge_search_honors_cancellation_and_timeout(monkeypatch): tools._RAG_SEARCH_SLOT.release() +def test_timed_out_search_frees_slot_for_next_lookup(monkeypatch): + # A search that outlives its timeout must not keep holding the sole RAG slot, or every later + # lookup would starve. The caller frees the slot when it stops waiting; the detached worker + # finishes later without re-holding it. + from core.inference import tools + + started = threading.Event() + release = threading.Event() + + def stalled_search(arguments, rag_scope): + started.set() + release.wait() + return "late" + + monkeypatch.setattr(tools, "_search_knowledge_base", stalled_search) + try: + timed_out = tools._search_knowledge_base_with_budget( + {"query": "q"}, {"kb_id": "a"}, timeout = 1 + ) + assert "timed out" in timed_out.lower() + assert started.is_set() + # The worker is still stalled, but the slot must be free for the next lookup. + assert tools._RAG_SEARCH_SLOT.acquire(timeout = 1) + tools._RAG_SEARCH_SLOT.release() + finally: + release.set() + + def test_search_for_autoinject_gates_on_dense_score(rag_conn, bow_embeddings, monkeypatch): _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3)