Studio: RAG-as-tool composer button (Phase 4)

Promotes RAG to a first-class composer toggle alongside Think / Web
Search / Code, with tool-use semantics on local models that support
tools and a pre-fetch fallback on external providers. The model
decides when to call `search_knowledge_base` on local inference; on
external providers retrieval still fires before each message (the
existing pre-fetch path), gated on the same button.

Backend
- core/rag/tool.py (new): search_knowledge_base handler + JSON-schema
  tool spec. Resolves scope (kb_id wins over thread_id) from the
  request's rag_scope, runs retrieve_hybrid + optional rerank, then
  hydrates filename / page_number / text from sqlite and formats as
  numbered Markdown citations ('[1] file.pdf (page 5): ...') for the
  LLM to cite. Empty scope returns a user-facing hint; empty results
  return a clear no-match message instead of an empty string.
- core/inference/tools.py: SEARCH_KNOWLEDGE_BASE_TOOL added to
  ALL_TOOLS (lazy import keeps tools.py importable on inference
  paths that never touch RAG). execute_tool() gains a tool_context
  parameter that carries per-request extras the LLM doesn't see
  (currently just rag_scope). The new 'search_knowledge_base' branch
  dispatches to the handler with scope unpacked from tool_context.
- core/inference/llama_cpp.py + safetensors_agentic.py +
  orchestrator.py: thread tool_context through generate_chat_completion_
  with_tools / run_safetensors_tool_loop / execute_tool. Both local
  backends (GGUF llama-server and safetensors agentic) carry the same
  context object.
- models/inference.py: ChatCompletionRequest gains optional
  rag_scope: dict ({kb_id?, thread_id?, enable_rerank?, default_top_k?,
  reranker_model?}). Ignored unless 'search_knowledge_base' is in
  enabled_tools.
- routes/inference.py: both the GGUF and safetensors call sites for
  generate_chat_completion_with_tools forward payload.rag_scope into
  tool_context.

Frontend
- chat-runtime-store.ts: global ragToolEnabled boolean + setter +
  CHAT_RAG_TOOL_ENABLED_KEY localStorage, mirroring toolsEnabled /
  codeToolsEnabled. Settings-hydration migration auto-flips
  ragToolEnabled=true for pre-Phase-4 users who already had ragSource
  set, so existing RAG users don't silently lose retrieval on upgrade.
- shared-composer.tsx: new 'RAG' pill button after Images (uses
  lucide BookOpenIcon, composer-pill-btn style, data-active toggle).
  Disabled when no model is loaded. Toggling on from ragSource='off'
  auto-flips source to 'thread' so the sidebar lands ready-to-go.
- chat-adapter.ts:
  * The existing pre-fetch block is now gated on ragToolEnabled AND
    only fires when the tool path isn't viable (external provider OR
    local model without tool-use support). Tool-capable local models
    skip pre-fetch and let the LLM decide.
  * The local-model body assembly adds 'search_knowledge_base' to
    enabled_tools and packs ragSource + enableRerank + ragTopK into a
    rag_scope object the backend tool handler consumes.
- chat-settings-sheet.tsx: entire Retrieval CollapsibleSection is
  wrapped in {ragToolEnabled && ...} so it hides when the button is
  off — the button is now the single on/off control. The 'Off'
  option is removed from the Source dropdown (the button handles
  that). Default open when shown so settings are one click away.

Tests
- test_rag_tool_handler.py: handler covers empty query, missing
  scope, kb_id > thread_id precedence, thread-only path, citation
  formatting (numbered + page numbers + unknown source); tool spec
  shape (function/name/required); execute_tool dispatch with and
  without tool_context; ALL_TOOLS includes the new spec without
  dropping the existing ones.

Verification scope
- Local GGUF with tools: toggle button on, upload doc, ask about
  doc content → assistant emits a search_knowledge_base tool call
  card (rendered by the existing ToolFallback component since no
  custom UI exists yet — that's a v2 nice-to-have).
- External provider (Anthropic / OpenAI / etc.): same button, same
  UX, but uses the pre-fetch path under the hood.
- Migration: pre-existing ragSource != off → button initializes ON
  so retrieval keeps working.

studio/backend/core/inference/llama_cpp.py

@@ -4428,6 +4428,7 @@ class LlamaCppBackend:
         auto_heal_tool_calls: bool = True,
         tool_call_timeout: int = 300,
         session_id: Optional[str] = None,
+        tool_context: Optional[dict] = None,
     ) -> Generator[dict, None, None]:
         """
         Agentic loop: let the model call tools, execute them, and continue.
@@ -5080,6 +5081,7 @@ class LlamaCppBackend:
                             cancel_event = cancel_event,
                             timeout = _effective_timeout,
                             session_id = session_id,
+                            tool_context = tool_context,
                         )
 
                     yield {

studio/backend/core/inference/orchestrator.py

@@ -838,6 +838,7 @@ class InferenceOrchestrator:
         auto_heal_tool_calls: bool = True,
         tool_call_timeout: int = 300,
         session_id: Optional[str] = None,
+        tool_context: Optional[dict] = None,
         use_adapter: Optional[Union[bool, str]] = None,
         **_unused,
     ):
@@ -895,6 +896,7 @@ class InferenceOrchestrator:
             max_tool_iterations = max_tool_iterations,
             tool_call_timeout = tool_call_timeout,
             session_id = session_id,
+            tool_context = tool_context,
         )
 
     def generate_with_adapter_control(

studio/backend/core/inference/safetensors_agentic.py

@@ -105,6 +105,7 @@ def run_safetensors_tool_loop(
     max_tool_iterations: int = 25,
     tool_call_timeout: int = 300,
     session_id: Optional[str] = None,
+    tool_context: Optional[dict] = None,
 ) -> Generator[dict, None, None]:
     """Drive an agentic tool loop on top of a cumulative-text generator.
 
@@ -340,6 +341,7 @@ def run_safetensors_tool_loop(
                             cancel_event = cancel_event,
                             timeout = eff_timeout,
                             session_id = session_id,
+                            tool_context = tool_context,
                         )
                     except Exception as exc:
                         logger.exception("Tool %s raised: %s", tool_name, exc)

studio/backend/core/inference/tools.py

@@ -502,7 +502,20 @@ TERMINAL_TOOL = {
     },
 }
 
-ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
+# Lazy import — keeps studio.db init lazy so tools.py doesn't pull in
+# the whole rag stack on inference paths that never see RAG.
+def _get_rag_tool_spec():
+    from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL
+
+    return SEARCH_KNOWLEDGE_BASE_TOOL
+
+
+# RAG_SEARCH_TOOL is included in ALL_TOOLS; routes/inference.py filters
+# the list against payload.enabled_tools so each request only sees the
+# tools the frontend explicitly enabled. When the RAG button is off
+# the tool name won't be in enabled_tools and the LLM will never see
+# the spec.
+ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()]
 
 
 _TIMEOUT_UNSET = object()
@@ -514,12 +527,17 @@ def execute_tool(
     cancel_event = None,
     timeout: int | None = _TIMEOUT_UNSET,
     session_id: str | None = None,
+    tool_context: dict | None = None,
 ) -> str:
     """Execute a tool by name with the given arguments. Returns result as a string.
 
     ``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
     unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
     ``session_id``: optional thread/session ID for per-conversation sandbox isolation.
+    ``tool_context``: optional per-request extras the LLM does not see (RAG scope,
+    future per-tool overrides). Keys consumed:
+      - ``rag_scope``: ``{kb_id?, thread_id?, enable_rerank?, default_top_k?,
+        reranker_model?}`` — consumed by ``search_knowledge_base``.
     """
     logger.info(
         f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
@@ -539,6 +557,19 @@ def execute_tool(
         return _bash_exec(
             arguments.get("command", ""), cancel_event, effective_timeout, session_id
         )
+    if name == "search_knowledge_base":
+        from core.rag.tool import search_knowledge_base
+
+        scope = (tool_context or {}).get("rag_scope") or {}
+        return search_knowledge_base(
+            query = arguments.get("query", ""),
+            top_k = arguments.get("top_k"),
+            scope_kb_id = scope.get("kb_id"),
+            scope_thread_id = scope.get("thread_id"),
+            enable_rerank = bool(scope.get("enable_rerank")),
+            reranker_model = scope.get("reranker_model"),
+            default_top_k = int(scope.get("default_top_k") or 5),
+        )
     return f"Unknown tool: {name}"
 
 

studio/backend/core/rag/tool.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""`search_knowledge_base` tool — RAG retrieval surfaced to the LLM.
+
+Invoked from `core/inference/tools.execute_tool` when the local model
+emits a `search_knowledge_base` call. The handler runs the existing
+hybrid retrieval, hydrates chunk text + filename + page number from
+sqlite, and returns a Markdown-with-numbered-citations string that
+the LLM consumes as the tool-result message.
+
+Scope (`kb_id` / `thread_id`) is not exposed as a tool argument — it
+comes from the chat-completions request body (`rag_scope`) so the
+LLM doesn't need to know about KB UUIDs.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+SEARCH_KNOWLEDGE_BASE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "search_knowledge_base",
+        "description": (
+            "Search the user's attached documents for information relevant to "
+            "the user's question. Call this when the user references content "
+            "from their docs, asks fact-heavy questions, or needs grounded "
+            "citations. Returns numbered chunks with source filenames; cite "
+            "them in your reply as [1], [2], etc."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": (
+                        "A focused search query — phrase it as the question "
+                        "you want answered, not as a keyword list."
+                    ),
+                },
+                "top_k": {
+                    "type": "integer",
+                    "minimum": 1,
+                    "maximum": 20,
+                    "description": (
+                        "How many chunks to retrieve (default 5). Higher = "
+                        "more grounding, more tokens."
+                    ),
+                },
+            },
+            "required": ["query"],
+        },
+    },
+}
+
+
+def _format_hits_for_llm(hits: list[Any]) -> str:
+    """Render hits as numbered Markdown citations for the LLM.
+
+    Empty results produce a one-line message rather than an empty
+    string — the model needs to know the search ran but found nothing
+    so it can fall back to its own knowledge or ask the user.
+    """
+    if not hits:
+        return (
+            "No matching chunks were found in the attached documents. "
+            "Either nothing in this scope is relevant, or no documents "
+            "have been ingested yet."
+        )
+    lines: list[str] = []
+    for index, hit in enumerate(hits, start = 1):
+        name = hit.get("filename") or "unknown source"
+        page = hit.get("page_number")
+        suffix = f" (page {page})" if page is not None else ""
+        text = (hit.get("text") or "").strip()
+        lines.append(f"[{index}] {name}{suffix}: {text}")
+    return "\n\n".join(lines)
+
+
+def search_knowledge_base(
+    *,
+    query: str,
+    top_k: int | None = None,
+    scope_kb_id: str | None = None,
+    scope_thread_id: str | None = None,
+    enable_rerank: bool = False,
+    reranker_model: str | None = None,
+    default_top_k: int = 5,
+) -> str:
+    """Execute the RAG search and return a tool-result string.
+
+    `kb_id` takes precedence over `thread_id` when both are set —
+    matches the create/upload contract that a document belongs to one
+    or the other, never both.
+    """
+    if not query or not query.strip():
+        return "Error: empty query."
+
+    if not scope_kb_id and not scope_thread_id:
+        return (
+            "No knowledge base or thread documents are configured for "
+            "retrieval. Ask the user to upload a document or select a "
+            "knowledge base in the chat settings."
+        )
+
+    from core.rag import retrieval
+    from core.rag.vector_store import kb_scope, thread_scope
+    from storage.studio_db import get_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
+
+    if enable_rerank:
+        from utils.rag.config import RAG_RERANK_CANDIDATE_K
+
+        candidate_k = max(k, RAG_RERANK_CANDIDATE_K)
+    else:
+        candidate_k = k
+
+    try:
+        hits = retrieval.retrieve_hybrid(scope, query.strip(), k = candidate_k)
+    except Exception as exc:  # noqa: BLE001
+        logger.exception("search_knowledge_base retrieval failed")
+        return f"Error: retrieval failed ({type(exc).__name__})."
+
+    chunk_ids = [h.chunk_id for h in hits]
+    lookup: dict[str, dict] = {}
+    if chunk_ids:
+        placeholders = ",".join("?" for _ in chunk_ids)
+        with get_connection() as conn:
+            rows = conn.execute(
+                f"""
+                SELECT c.id AS chunk_id, c.text, c.page_number,
+                       c.kind, d.filename
+                FROM rag_chunks c
+                JOIN rag_documents d ON d.id = c.document_id
+                WHERE c.id IN ({placeholders})
+                """,
+                chunk_ids,
+            ).fetchall()
+        for row in rows:
+            lookup[row["chunk_id"]] = dict(row)
+
+    if enable_rerank and hits:
+        from core.rag import reranker
+
+        pairs = [
+            (hit, lookup[hit.chunk_id]["text"])
+            for hit in hits
+            if hit.chunk_id in lookup
+        ]
+        try:
+            hits = reranker.rerank(
+                query.strip(),
+                pairs,
+                model_name = reranker_model,
+                top_k = k,
+            )
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("rerank failed in search_knowledge_base: %s", exc)
+            hits = hits[:k]
+    else:
+        hits = hits[:k]
+
+    # Image-kind hits don't carry LLM-friendly text — skip them. The
+    # paired caption (linked_chunk_id) usually surfaces separately.
+    formatted = [
+        lookup[hit.chunk_id]
+        for hit in hits
+        if hit.chunk_id in lookup and lookup[hit.chunk_id].get("kind") != "image"
+    ]
+    return _format_hits_for_llm(formatted)

studio/backend/models/inference.py

@@ -686,6 +686,16 @@ class ChatCompletionRequest(BaseModel):
         None,
         description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
     )
+    rag_scope: Optional[dict] = Field(
+        None,
+        description = (
+            "[x-unsloth] Per-request context the `search_knowledge_base` tool "
+            "consumes when the LLM invokes it. Shape: "
+            "{kb_id?: str, thread_id?: str, enable_rerank?: bool, "
+            "default_top_k?: int, reranker_model?: str}. Ignored unless "
+            "'search_knowledge_base' is in enabled_tools."
+        ),
+    )
     cancel_id: Optional[str] = Field(
         None,
         description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.",

studio/backend/routes/inference.py

@@ -2483,6 +2483,11 @@ async def openai_chat_completions(
                     if payload.tool_call_timeout is not None
                     else 300,
                     session_id = payload.session_id,
+                    tool_context = (
+                        {"rag_scope": payload.rag_scope}
+                        if payload.rag_scope
+                        else None
+                    ),
                 )
 
             _tool_sentinel = object()
@@ -2950,6 +2955,11 @@ async def openai_chat_completions(
 
         def sf_generate_with_tools():
             return backend.generate_chat_completion_with_tools(
+                tool_context = (
+                    {"rag_scope": payload.rag_scope}
+                    if payload.rag_scope
+                    else None
+                ),
                 messages = _sf_chat_messages,
                 tools = _sf_tools_to_use,
                 system_prompt = _sf_system_prompt or "",

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -993,8 +993,22 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       // prepend it as a system-role block. Failures are logged but
       // don't break the chat — better to answer without context than
       // to drop a message the user just sent.
+      //
+      // Phase 4: pre-fetch only fires when the new RAG button is on AND
+      // we can't register `search_knowledge_base` as a real tool —
+      // i.e., external providers or local models that don't expose
+      // tool-use. Local models with tool support take the tool-call
+      // path further down (see `enabled_tools` assembly), and the LLM
+      // decides per turn whether to invoke retrieval.
       const ragSource = runtime.ragSource;
-      if (ragSource.kind !== "off") {
+      const ragToolEnabled = runtime.ragToolEnabled;
+      const ragToolPathTaken =
+        ragToolEnabled && supportsTools && !isExternalRequest;
+      if (
+        ragToolEnabled
+        && ragSource.kind !== "off"
+        && !ragToolPathTaken
+      ) {
         const lastUser = [...outboundMessages]
           .reverse()
           .find((m) => m.role === "user");
@@ -1617,13 +1631,36 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
             ...(supportsPreserveThinking
               ? { preserve_thinking: preserveThinking }
               : {}),
-            ...(supportsTools && (toolsEnabled || codeToolsEnabled)
+            ...(supportsTools
+              && (toolsEnabled || codeToolsEnabled || ragToolPathTaken)
               ? {
                   enable_tools: true,
                   enabled_tools: [
                     ...(toolsEnabled ? ["web_search"] : []),
                     ...(codeToolsEnabled ? ["python", "terminal"] : []),
+                    ...(ragToolPathTaken ? ["search_knowledge_base"] : []),
                   ],
+                  // Phase 4: per-request RAG context the backend's
+                  // `search_knowledge_base` handler reads when the LLM
+                  // invokes the tool. Only sent when the tool path
+                  // is taken — external providers fall through to the
+                  // pre-fetch block above.
+                  ...(ragToolPathTaken
+                    ? {
+                        rag_scope: {
+                          kb_id:
+                            ragSource.kind === "kb"
+                              ? ragSource.kbId
+                              : null,
+                          thread_id:
+                            ragSource.kind === "thread"
+                              ? (resolvedThreadId ?? null)
+                              : null,
+                          enable_rerank: runtime.enableRerank,
+                          default_top_k: runtime.ragTopK,
+                        },
+                      }
+                    : {}),
                   auto_heal_tool_calls:
                     useChatRuntimeStore.getState().autoHealToolCalls,
                   max_tool_calls_per_message:

studio/frontend/src/features/chat/chat-settings-sheet.tsx

@@ -441,6 +441,7 @@ export function ChatSettingsPanel({
   const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
   const ragSource = useChatRuntimeStore((s) => s.ragSource);
   const setRagSource = useChatRuntimeStore((s) => s.setRagSource);
+  const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled);
   const enableRerank = useChatRuntimeStore((s) => s.enableRerank);
   const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank);
   const ragTopK = useChatRuntimeStore((s) => s.ragTopK);
@@ -1246,7 +1247,8 @@ export function ChatSettingsPanel({
           </CollapsibleSection>
         ) : null}
 
-        <CollapsibleSection label="Retrieval" defaultOpen={false}>
+        {ragToolEnabled ? (
+        <CollapsibleSection label="Retrieval" defaultOpen={true}>
           <div className="flex flex-col gap-3 pt-1">
             <div className="flex flex-col gap-1.5">
               <label className="text-[12px] font-medium text-muted-foreground">
@@ -1265,11 +1267,6 @@ export function ChatSettingsPanel({
                   align="start"
                   className="w-[var(--radix-dropdown-menu-trigger-width)]"
                 >
-                  <DropdownMenuItem
-                    onSelect={() => setRagSource({ kind: "off" })}
-                  >
-                    Off
-                  </DropdownMenuItem>
                   <DropdownMenuItem
                     onSelect={() => setRagSource({ kind: "thread" })}
                   >
@@ -1527,6 +1524,7 @@ export function ChatSettingsPanel({
             </div>
           </div>
         </CollapsibleSection>
+        ) : null}
 
         <CollapsibleSection label="System Prompt" defaultOpen={true}>
           <button

studio/frontend/src/features/chat/shared-composer.tsx

@@ -21,7 +21,7 @@ import { isTauri } from "@/lib/api-base";
 import { isMultimodalResponse } from "./types/api";
 import { getImageInputUnavailableReason } from "./utils/image-input-support";
 import { useAui } from "@assistant-ui/react";
-import { ArrowUpIcon, FileTextIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
+import { ArrowUpIcon, BookOpenIcon, FileTextIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
 import { useRagStore } from "@/features/rag/stores/rag-store";
 import { subscribeToJobEvents } from "@/features/rag/api/rag-api";
 import { toast } from "@/lib/toast";
@@ -368,6 +368,8 @@ export function SharedComposer({
   const setImageToolsEnabled = useChatRuntimeStore(
     (s) => s.setImageToolsEnabled,
   );
+  const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled);
+  const setRagToolEnabled = useChatRuntimeStore((s) => s.setRagToolEnabled);
   const lastOpenRouterChosenModel = useChatRuntimeStore(
     (s) => s.lastOpenRouterChosenModel,
   );
@@ -1275,6 +1277,34 @@ export function SharedComposer({
               <span>Images</span>
             </button>
           )}
+          {/* RAG: master switch for retrieval. On local models with
+              tool-use support, registers `search_knowledge_base` as a
+              tool the LLM can call. On external providers, falls back
+              to the pre-fetch path. The sidebar Retrieval section
+              configures the source / top-K / reranker; the button is
+              the only on/off control. */}
+          <button
+            type="button"
+            disabled={!modelLoaded}
+            onClick={() => {
+              const next = !ragToolEnabled;
+              setRagToolEnabled(next);
+              if (next && ragSource.kind === "off") {
+                setRagSource({ kind: "thread" });
+              }
+            }}
+            className="composer-pill-btn"
+            data-active={ragToolEnabled && modelLoaded ? "true" : "false"}
+            aria-label={ragToolEnabled ? "Disable RAG" : "Enable RAG"}
+            title={
+              ragToolEnabled
+                ? "RAG on — the model can search your attached documents"
+                : "Enable RAG — let the model search your documents"
+            }
+          >
+            <BookOpenIcon className="size-3.5" />
+            <span>RAG</span>
+          </button>
         </div>
         <div className="flex items-center gap-1">
           {dictationSupported && (

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -26,6 +26,7 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
 export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
 export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
 export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled";
+export const CHAT_RAG_TOOL_ENABLED_KEY = "unsloth_chat_rag_tool_enabled";
 
 // External provider selection is encoded into `params.checkpoint` as
 // `external::<providerId>::<modelId>`. PersistedChatSettings deliberately
@@ -264,6 +265,7 @@ type ChatRuntimeStore = {
    */
   supportsBuiltinImageGeneration: boolean;
   toolsEnabled: boolean;
+  ragToolEnabled: boolean;
   codeToolsEnabled: boolean;
   imageToolsEnabled: boolean;
   toolStatus: string | null;
@@ -344,6 +346,7 @@ type ChatRuntimeStore = {
   setRagSource: (source: RagSource) => void;
   setEnableRerank: (value: boolean) => void;
   setRagTopK: (value: number) => void;
+  setRagToolEnabled: (value: boolean) => void;
 };
 
 type PersistedChatSettings = Awaited<
@@ -578,6 +581,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   supportsBuiltinCodeExecution: false,
   supportsBuiltinImageGeneration: false,
   toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
+  // Phase 4: RAG button defaults off. Migration nudge happens after
+  // settings hydration, when the persisted ragSource becomes visible —
+  // see hydratePersistedSettings.
+  ragToolEnabled: loadBool(CHAT_RAG_TOOL_ENABLED_KEY, false),
   codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
   imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
   toolStatus: null,
@@ -630,6 +637,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
             ),
             ...getHydratedSettingsState(settings, state, hydrationVersions),
           };
+          // Phase 4 migration: pre-existing users with ragSource set
+          // before the RAG button shipped should keep getting RAG —
+          // auto-flip ragToolEnabled so the button starts ON for them.
+          // The CHAT_RAG_TOOL_ENABLED_KEY localStorage write makes the
+          // migration stick across reloads.
+          const hydratedRagSource =
+            (nextState.ragSource as RagSource | undefined) ?? state.ragSource;
+          if (
+            hydratedRagSource &&
+            hydratedRagSource.kind !== "off" &&
+            !state.ragToolEnabled
+          ) {
+            nextState.ragToolEnabled = true;
+            saveBool(CHAT_RAG_TOOL_ENABLED_KEY, true);
+          }
           return nextState;
         });
       } catch {
@@ -831,6 +853,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       }
       return { toolsEnabled };
     }),
+  setRagToolEnabled: (ragToolEnabled) =>
+    set(() => {
+      saveBool(CHAT_RAG_TOOL_ENABLED_KEY, ragToolEnabled);
+      return { ragToolEnabled };
+    }),
   setCodeToolsEnabled: (codeToolsEnabled) =>
     set(() => {
       saveBool(CHAT_CODE_TOOLS_ENABLED_KEY, codeToolsEnabled);

tests/python/test_rag_tool_handler.py

@@ -0,0 +1,193 @@
+"""Unit tests for the `search_knowledge_base` tool handler (Phase 4)."""
+
+import sys
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
+if str(STUDIO_BACKEND) not in sys.path:
+    sys.path.insert(0, str(STUDIO_BACKEND))
+
+
+def _make_hit(chunk_id: str):
+    """Minimal stand-in for retrieval.Hit — just needs .chunk_id."""
+    class _Hit:
+        pass
+    h = _Hit()
+    h.chunk_id = chunk_id
+    h.score = 1.0
+    h.kind = "text"
+    h.document_id = None
+    h.chunk_index = 0
+    return h
+
+
+def test_empty_query_returns_error():
+    from core.rag.tool import search_knowledge_base
+
+    result = search_knowledge_base(query = "", scope_thread_id = "t-1")
+    assert result.startswith("Error:")
+    assert "empty" in result.lower()
+
+
+def test_missing_scope_returns_user_facing_hint():
+    from core.rag.tool import search_knowledge_base
+
+    result = search_knowledge_base(
+        query = "anything",
+        scope_kb_id = None,
+        scope_thread_id = None,
+    )
+    assert "No knowledge base" in result
+    assert "thread documents" in result
+
+
+def test_kb_takes_precedence_over_thread():
+    """When both kb_id and thread_id are passed, kb_id wins."""
+    from core.rag import tool
+
+    captured = {}
+
+    def _stub_retrieve(scope, query, k):
+        captured["scope"] = scope
+        return []
+
+    with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
+                       "retrieve_hybrid",
+                       _stub_retrieve):
+        result = tool.search_knowledge_base(
+            query = "x",
+            scope_kb_id = "kb-abc",
+            scope_thread_id = "thread-xyz",
+        )
+
+    assert captured["scope"].startswith("kb_")
+    assert "kb-abc" in captured["scope"]
+    assert "thread" not in captured["scope"].split("kb_")[1]
+
+
+def test_thread_scope_when_only_thread_set():
+    from core.rag import tool
+
+    captured = {}
+
+    def _stub_retrieve(scope, query, k):
+        captured["scope"] = scope
+        return []
+
+    with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
+                       "retrieve_hybrid",
+                       _stub_retrieve):
+        tool.search_knowledge_base(
+            query = "x",
+            scope_thread_id = "thread-xyz",
+        )
+
+    assert captured["scope"].startswith("thread_")
+
+
+def test_empty_results_message_is_user_facing():
+    from core.rag.tool import _format_hits_for_llm
+
+    result = _format_hits_for_llm([])
+    assert "No matching chunks" in result
+
+
+def test_format_hits_produces_numbered_citations():
+    from core.rag.tool import _format_hits_for_llm
+
+    hits = [
+        {"filename": "alpha.pdf", "page_number": 3, "text": "first body"},
+        {"filename": "beta.md", "page_number": None, "text": "second body"},
+    ]
+    result = _format_hits_for_llm(hits)
+    assert "[1] alpha.pdf (page 3): first body" in result
+    assert "[2] beta.md: second body" in result
+    # Each hit on its own paragraph so the LLM can cite cleanly.
+    assert "\n\n" in result
+
+
+def test_format_hits_handles_unknown_source():
+    from core.rag.tool import _format_hits_for_llm
+
+    hits = [{"filename": None, "page_number": None, "text": "orphan"}]
+    result = _format_hits_for_llm(hits)
+    assert "[1] unknown source: orphan" in result
+
+
+def test_tool_spec_shape_is_openai_compatible():
+    from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL
+
+    assert SEARCH_KNOWLEDGE_BASE_TOOL["type"] == "function"
+    fn = SEARCH_KNOWLEDGE_BASE_TOOL["function"]
+    assert fn["name"] == "search_knowledge_base"
+    assert "query" in fn["parameters"]["required"]
+    assert "top_k" in fn["parameters"]["properties"]
+    # Description should hint at when to call so the LLM picks it up
+    # appropriately. Don't lock the exact wording.
+    assert "documents" in fn["description"].lower()
+
+
+def test_execute_tool_dispatches_to_search_knowledge_base():
+    """tools.execute_tool should route 'search_knowledge_base' correctly."""
+    from core.inference import tools
+
+    called = {}
+
+    def _stub(*, query, top_k = None, scope_kb_id = None, scope_thread_id = None,
+              enable_rerank = False, reranker_model = None, default_top_k = 5):
+        called["query"] = query
+        called["top_k"] = top_k
+        called["scope_kb_id"] = scope_kb_id
+        called["scope_thread_id"] = scope_thread_id
+        called["enable_rerank"] = enable_rerank
+        called["default_top_k"] = default_top_k
+        return "stub-result"
+
+    with patch("core.rag.tool.search_knowledge_base", _stub):
+        result = tools.execute_tool(
+            "search_knowledge_base",
+            {"query": "hello", "top_k": 7},
+            tool_context = {
+                "rag_scope": {
+                    "kb_id": "kb-1",
+                    "enable_rerank": True,
+                    "default_top_k": 3,
+                }
+            },
+        )
+    assert result == "stub-result"
+    assert called["query"] == "hello"
+    assert called["top_k"] == 7
+    assert called["scope_kb_id"] == "kb-1"
+    assert called["scope_thread_id"] is None
+    assert called["enable_rerank"] is True
+    assert called["default_top_k"] == 3
+
+
+def test_execute_tool_handles_missing_tool_context():
+    """tool_context=None should still dispatch without crashing."""
+    from core.inference import tools
+
+    def _stub(*, query, **_kwargs):
+        return f"got: {query}"
+
+    with patch("core.rag.tool.search_knowledge_base", _stub):
+        result = tools.execute_tool(
+            "search_knowledge_base",
+            {"query": "ping"},
+            tool_context = None,
+        )
+    assert result == "got: ping"
+
+
+def test_all_tools_includes_rag():
+    from core.inference.tools import ALL_TOOLS
+
+    names = [t["function"]["name"] for t in ALL_TOOLS]
+    assert "search_knowledge_base" in names
+    assert "web_search" in names  # regression — we shouldn't have removed the others
+    assert "python" in names