From 38e22c15d7e8441c4bc2014826802d27dbf871a0 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 26 May 2026 21:24:31 +0400
Subject: [PATCH] Studio: bump llama-server prefill timeout to 300s + warm RAG
 embedder when RAG turns on

---
 studio/backend/core/inference/llama_cpp.py    |  5 +++-
 studio/backend/routes/rag.py                  | 26 +++++++++++++++++++
 .../chat/stores/chat-runtime-store.ts         | 21 ++++++++++++++-
 .../frontend/src/features/rag/api/rag-api.ts  |  7 +++++
 4 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index ba438d9ea6..0531dd0e93 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -4203,9 +4203,12 @@ class LlamaCppBackend:
             # without triggering a retry storm.  Cancel during both
             # prefill and streaming is handled by the watcher thread
             # which closes the response, unblocking any httpx read.
+            # 300 s headroom for large models (30B+) re-prefilling after
+            # a tool call that returned a long result (e.g. RAG chunks
+            # with images) — prior 120 s was tripping on Gemma-4-31B.
             prefill_timeout = httpx.Timeout(
                 connect = 30,
-                read = 120.0,
+                read = 300.0,
                 write = 10,
                 pool = 10,
             )
diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py
index 022e5ddb16..38ef6b14e4 100644
--- a/studio/backend/routes/rag.py
+++ b/studio/backend/routes/rag.py
@@ -435,6 +435,32 @@ def get_rag_defaults(
     return _load_rag_defaults()
 
 
+@router.post("/warmup")
+def warmup_rag_embedder(
+    current_subject: str = Depends(get_current_subject),
+) -> dict:
+    """Preload the configured default embedder so the first retrieval is warm.
+
+    Called from the frontend when the user enables the RAG pill — moves the
+    cold-load latency (Qwen3-VL-Embedding-2B is ~4 GB) out of the first
+    chat-completion path, where a 30s+ load can race the llama-server
+    prefill timeout.
+    """
+    from utils.rag.config import resolve_embedder
+
+    defaults = _load_rag_defaults()
+    model_name = defaults.embedding_model or resolve_embedder(
+        defaults.mode,
+        defaults.chunking_strategy,
+    )
+    try:
+        embeddings.get_embedder(model_name)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("RAG warmup failed for %s: %s", model_name, exc)
+        return {"ok": False, "model": model_name, "error": str(exc)}
+    return {"ok": True, "model": model_name}
+
+
 @router.put("/defaults", response_model = RagDefaults)
 def set_rag_defaults(
     payload: UpdateRagDefaultsRequest,
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 34c4948dac..35b2f5ab09 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -657,6 +657,17 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
             nextState.ragToolEnabled = true;
             saveBool(CHAT_RAG_TOOL_ENABLED_KEY, true);
           }
+          // After hydration, if RAG is on (persisted or just migrated),
+          // warm the embedder so the first message doesn't pay the
+          // cold load inline.
+          if (
+            nextState.ragToolEnabled === true ||
+            (nextState.ragToolEnabled === undefined && state.ragToolEnabled)
+          ) {
+            void import("@/features/rag/api/rag-api")
+              .then((m) => m.warmupRagEmbedder())
+              .catch(() => {});
+          }
           return nextState;
         });
       } catch {
@@ -869,8 +880,16 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       return { toolsEnabled };
     }),
   setRagToolEnabled: (ragToolEnabled) =>
-    set(() => {
+    set((state) => {
       saveBool(CHAT_RAG_TOOL_ENABLED_KEY, ragToolEnabled);
+      // Warmup on off→on transitions: kick the backend to preload the
+      // embedder so the user's first RAG-using message doesn't pay the
+      // cold-start (~30s for Qwen3-VL-Embedding-2B) inline. Fire-and-forget.
+      if (ragToolEnabled && !state.ragToolEnabled) {
+        void import("@/features/rag/api/rag-api")
+          .then((m) => m.warmupRagEmbedder())
+          .catch(() => {});
+      }
       return { ragToolEnabled };
     }),
   setCodeToolsEnabled: (codeToolsEnabled) =>
diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts
index 4671985645..9f56e20ac9 100644
--- a/studio/frontend/src/features/rag/api/rag-api.ts
+++ b/studio/frontend/src/features/rag/api/rag-api.ts
@@ -304,6 +304,13 @@ export async function setRagDefaults(
   return parseJsonOrThrow<RagDefaults>(response);
 }
 
+/** Preload the configured embedder on the backend. Long-running (cold
+ *  load can take 30s+). Fire-and-forget: failure is non-fatal because
+ *  the first real query will lazy-load again. */
+export async function warmupRagEmbedder(): Promise<void> {
+  await authFetch("/api/rag/warmup", { method: "POST" });
+}
+
 // --- Search ---
 
 export async function search(req: SearchRequest): Promise<SearchHit[]> {