Studio: free the RAG search slot when a lookup times out or is cancelled

The bounded knowledge-base search held the sole admission slot in a detached
worker until the search returned, so a lookup that outlived its timeout (a
stalled embedding or blocked vector call) kept the slot forever and starved
every later lookup, disabling knowledge-base retrieval globally. Release the
slot from the caller when it stops waiting, exactly once, so a detached worker
finishes without re-holding it.

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:

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)