HKUDS · mtsui-cmyk · May 19, 2026
diff --git a/web/app/(workspace)/chat/[[...sessionId]]/page.tsx b/web/app/(workspace)/chat/[[...sessionId]]/page.tsx
@@ -90,6 +90,10 @@ import {
   type OutlineItem,
   type ResearchSource,
 } from "@/lib/research-types";
+import {
+  applyChatRetrievalMode,
+  type ChatRetrievalMode,
+} from "@/lib/chat-retrieval-mode";
 import { listKnowledgeBases } from "@/lib/knowledge-api";
 import { listLLMOptions, type LLMOption } from "@/lib/llm-options";
 import { downloadChatMarkdown } from "@/lib/chat-export";
@@ -467,6 +471,8 @@ export default function ChatPage() {
   const [toolMenuOpen, setToolMenuOpen] = useState(false);
   const [spaceMenuOpen, setSpaceMenuOpen] = useState(false);
   const [kbMenuOpen, setKbMenuOpen] = useState(false);
+  const [chatRetrievalMode, setChatRetrievalMode] =
+    useState<ChatRetrievalMode>("auto");
   const [selectedNotebookRecords, setSelectedNotebookRecords] = useState<
     SelectedRecord[]
   >([]);
@@ -1299,6 +1305,18 @@ export default function ChatPage() {
       if (isVisualizeMode) config = buildVisualizeWSConfig(visualizeConfig);
       if (isResearchMode) config = buildResearchWSConfig(researchConfig);
 
+      const retrievalSelection =
+        activeCap.value === ""
+          ? applyChatRetrievalMode(
+              state.enabledTools,
+              state.knowledgeBases,
+              chatRetrievalMode,
+            )
+          : {
+              enabledTools: [...state.enabledTools],
+              knowledgeBases: [...state.knowledgeBases],
+            };
+
       const skillsPayload = skillsAutoMode ? ["auto"] : [...selectedSkills];
       const memoryPayload = [...memoryReferencesPayload];
       const messageContent =
@@ -1326,7 +1344,42 @@ export default function ChatPage() {
         config,
         notebookReferencesPayload,
         historyReferencesPayload,
-        { bookReferences: bookReferencesPayload },
+        {
+          bookReferences: bookReferencesPayload,
+          requestSnapshotOverride: {
+            content: messageContent,
+            capability: state.activeCapability,
+            enabledTools: retrievalSelection.enabledTools,
+            knowledgeBases: retrievalSelection.knowledgeBases,
+            language: state.language,
+            ...(extraAttachments.length
+              ? { attachments: extraAttachments }
+              : {}),
+            ...(config && Object.keys(config).length > 0 ? { config } : {}),
+            ...(notebookReferencesPayload.length
+              ? { notebookReferences: notebookReferencesPayload }
+              : {}),
+            ...(historyReferencesPayload.length
+              ? { historyReferences: historyReferencesPayload }
+              : {}),
+            ...(questionNotebookReferencesPayload.length
+              ? {
+                  questionNotebookReferences:
+                    questionNotebookReferencesPayload,
+                }
+              : {}),
+            ...(bookReferencesPayload.length
+              ? { bookReferences: bookReferencesPayload }
+              : {}),
+            ...(skillsPayload.length ? { skills: skillsPayload } : {}),
+            ...(memoryPayload.length
+              ? { memoryReferences: memoryPayload }
+              : {}),
+            ...(state.llmSelection
+              ? { llmSelection: state.llmSelection }
+              : {}),
+          },
+        },
         questionNotebookReferencesPayload,
         skillsPayload,
         memoryPayload,
@@ -1343,6 +1396,7 @@ export default function ChatPage() {
     },
     [
       attachments,
+      activeCap.value,
       bookReferencesPayload,
       historyReferencesPayload,
       isMathAnimatorMode,
@@ -1352,6 +1406,7 @@ export default function ChatPage() {
       mathAnimatorConfig,
       memoryReferencesPayload,
       notebookReferencesPayload,
+      chatRetrievalMode,
       questionNotebookReferencesPayload,
       quizConfig,
       quizPdf,
@@ -1365,7 +1420,12 @@ export default function ChatPage() {
       skillsAutoMode,
       sendMessage,
       shouldAutoScrollRef,
+      state.activeCapability,
+      state.enabledTools,
+      state.knowledgeBases,
       state.isStreaming,
+      state.language,
+      state.llmSelection,
       t,
       visualizeConfig,
     ],
@@ -1710,6 +1770,8 @@ export default function ChatPage() {
           onSetToolMenuOpen={setToolMenuOpen}
           onSetSpaceMenuOpen={setSpaceMenuOpen}
           onSetKbMenuOpen={setKbMenuOpen}
+          chatRetrievalMode={chatRetrievalMode}
+          onSetChatRetrievalMode={setChatRetrievalMode}
           onToggleAutoCap={noopToggleAutoCap}
           onToggleKB={handleToggleKB}
           onSelectLLM={setLLMSelection}

diff --git a/web/components/chat/home/ChatComposer.tsx b/web/components/chat/home/ChatComposer.tsx
@@ -39,6 +39,7 @@ import type { ResearchSource } from "@/lib/research-types";
 import ChatSpaceMenu from "@/components/chat/space/ChatSpaceMenu";
 import type { SpaceMemoryFile } from "@/lib/space-items";
 import type { SelectedBookReference } from "@/lib/book-references";
+import type { ChatRetrievalMode } from "@/lib/chat-retrieval-mode";
 import ModelSelector from "./ModelSelector";
 
 type SpaceSelectionCounts = {
@@ -122,6 +123,7 @@ export default memo(function ChatComposer({
   skillsAutoMode,
   selectedMemoryFiles,
   selectedKnowledgeBases,
+  chatRetrievalMode,
   isStreaming,
   isResearchMode,
   isMathAnimatorMode,
@@ -138,6 +140,7 @@ export default memo(function ChatComposer({
   onSetToolMenuOpen,
   onSetSpaceMenuOpen,
   onSetKbMenuOpen,
+  onSetChatRetrievalMode,
   onToggleAutoCap,
   onToggleKB,
   onSelectLLM,
@@ -208,6 +211,7 @@ export default memo(function ChatComposer({
   skillsAutoMode: boolean;
   selectedMemoryFiles: SpaceMemoryFile[];
   selectedKnowledgeBases: string[];
+  chatRetrievalMode: ChatRetrievalMode;
   isStreaming: boolean;
   isResearchMode: boolean;
   isMathAnimatorMode: boolean;
@@ -238,6 +242,7 @@ export default memo(function ChatComposer({
   onSetToolMenuOpen: (open: boolean | ((prev: boolean) => boolean)) => void;
   onSetSpaceMenuOpen: (open: boolean | ((prev: boolean) => boolean)) => void;
   onSetKbMenuOpen: (open: boolean | ((prev: boolean) => boolean)) => void;
+  onSetChatRetrievalMode: (mode: ChatRetrievalMode) => void;
   onToggleAutoCap: (cap: string) => void;
   onToggleKB: (name: string) => void;
   onSelectLLM: (selection: LLMSelection | null) => void;
@@ -312,6 +317,9 @@ export default memo(function ChatComposer({
     [onAddFiles],
   );
 
+  const activeCapabilityKey = activeCap.value || "chat";
+  const showRetrievalMode =
+    activeCapabilityKey === "chat" && selectedKnowledgeBases.length > 0;
   useEffect(() => {
     if (!hasMessages) textareaRef.current?.focus();
   }, [hasMessages]);
@@ -933,6 +941,39 @@ export default memo(function ChatComposer({
                   )}
                 </div>
 
+                {showRetrievalMode ? (
+                  <div
+                    className="ml-1 inline-flex items-center rounded-lg border border-[var(--border)] bg-[var(--background)]/65 p-0.5"
+                    title={t("Retrieval mode")}
+                    aria-label={t("Retrieval mode")}
+                  >
+                    {(
+                      [
+                        ["auto", t("Auto")],
+                        ["kb_only", t("KB")],
+                        ["kb_web", t("KB + Web")],
+                        ["off", t("Off")],
+                      ] as const
+                    ).map(([mode, label]) => {
+                      const active = chatRetrievalMode === mode;
+                      return (
+                        <button
+                          key={mode}
+                          type="button"
+                          onClick={() => onSetChatRetrievalMode(mode)}
+                          className={`rounded-md px-2 py-1 text-[10px] font-medium transition-colors ${
+                            active
+                              ? "bg-[var(--primary)] text-white"
+                              : "text-[var(--muted-foreground)] hover:text-[var(--foreground)]"
+                          }`}
+                        >
+                          {label}
+                        </button>
+                      );
+                    })}
+                  </div>
+                ) : null}
+
                 <div className="relative flex items-center gap-0.5">
                   <button
                     ref={spaceBtnRef}

diff --git a/web/components/chat/home/ChatMessages.tsx b/web/components/chat/home/ChatMessages.tsx
@@ -30,6 +30,10 @@ import type {
 import { apiUrl } from "@/lib/api";
 import { docIconFor } from "@/lib/doc-attachments";
 import { extractMathAnimatorResult } from "@/lib/math-animator-types";
+import {
+  extractRetrievalSummary,
+  type RetrievalSummary,
+} from "@/lib/retrieval-summary";
 import { extractQuizQuestions } from "@/lib/quiz-types";
 import { extractVisualizeResult } from "@/lib/visualize-types";
 import type { StreamEvent } from "@/lib/unified-ws";
@@ -170,6 +174,11 @@ const AssistantMessage = memo(function AssistantMessage({
     return extractVisualizeResult(resultEvent.metadata);
   }, [msg.capability, resultEvent]);
 
+  const retrievalSummary = useMemo(() => {
+    if (msg.capability !== "chat" || !resultEvent) return null;
+    return extractRetrievalSummary(resultEvent.metadata);
+  }, [msg.capability, resultEvent]);
+
   // Auto turns have a fundamentally different layout (interleaved
   // thinking + collapsed delegation cards + final synthesis). Each sub-
   // capability's internal trace is rendered INSIDE its delegation card —
@@ -202,6 +211,9 @@ const AssistantMessage = memo(function AssistantMessage({
         isStreaming={isStreaming}
         content={msg.content}
       />
+      {retrievalSummary ? (
+        <RetrievalSummaryCard summary={retrievalSummary} />
+      ) : null}
       {isStreaming && onAnswerNow ? (
         <AnswerNowRow onAnswerNow={onAnswerNow} />
       ) : null}
@@ -237,6 +249,62 @@ const AssistantMessage = memo(function AssistantMessage({
 
 AssistantMessage.displayName = "AssistantMessage";
 
+const RetrievalSummaryCard = memo(function RetrievalSummaryCard({
+  summary,
+}: {
+  summary: RetrievalSummary;
+}) {
+  const { t } = useTranslation();
+
+  return (
+    <div className="mb-3 rounded-xl border border-[var(--border)] bg-[var(--card)] px-3 py-2.5">
+      <div className="flex flex-wrap items-center gap-x-3 gap-y-1 text-[12px] text-[var(--muted-foreground)]">
+        <span className="font-medium text-[var(--foreground)]">
+          {t("Retrieved from")}
+        </span>
+        <span>
+          {summary.knowledgeBases.length} {t("Knowledge Bases")}
+        </span>
+        <span>
+          {summary.totalFiles} {t("files")}
+        </span>
+        <span>
+          {summary.totalChunks} {t("chunks")}
+        </span>
+      </div>
+      <div className="mt-2 flex flex-col gap-2">
+        {summary.knowledgeBases.map((item) => (
+          <div key={item.kbName} className="flex flex-col gap-1">
+            <div className="flex flex-wrap items-center gap-x-2 gap-y-1 text-[12px]">
+              <span className="rounded-full bg-[var(--muted)] px-2 py-0.5 font-medium text-[var(--foreground)]">
+                {item.kbName}
+              </span>
+              <span className="text-[var(--muted-foreground)]">
+                {item.files.length} {t("files")} · {item.chunkCount}{" "}
+                {t("chunks")}
+              </span>
+            </div>
+            {item.files.length > 0 ? (
+              <div className="flex flex-wrap gap-1.5">
+                {item.files.map((file) => (
+                  <span
+                    key={`${item.kbName}:${file}`}
+                    className="rounded-full border border-[var(--border)] px-2 py-0.5 text-[11px] text-[var(--muted-foreground)]"
+                  >
+                    {file}
+                  </span>
+                ))}
+              </div>
+            ) : null}
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+});
+
+RetrievalSummaryCard.displayName = "RetrievalSummaryCard";
+
 /**
  * Inline "Answer now" affordance shown alongside the active assistant turn.
  * Lives outside the trace panel so it is visible as soon as the turn starts

diff --git a/web/lib/chat-retrieval-mode.ts b/web/lib/chat-retrieval-mode.ts
@@ -0,0 +1,45 @@
+export type ChatRetrievalMode = "auto" | "kb_only" | "kb_web" | "off";
+
+export interface AppliedRetrievalSelection {
+  enabledTools: string[];
+  knowledgeBases: string[];
+}
+
+function uniqueStrings(values: string[]): string[] {
+  return Array.from(new Set(values));
+}
+
+export function applyChatRetrievalMode(
+  enabledTools: string[],
+  knowledgeBases: string[],
+  mode: ChatRetrievalMode,
+): AppliedRetrievalSelection {
+  const normalizedTools = uniqueStrings(enabledTools);
+  const normalizedKbs = uniqueStrings(knowledgeBases);
+
+  if (mode === "auto") {
+    return {
+      enabledTools: normalizedTools,
+      knowledgeBases: normalizedKbs,
+    };
+  }
+
+  if (mode === "kb_only") {
+    return {
+      enabledTools: normalizedTools.filter((tool) => tool !== "web_search"),
+      knowledgeBases: normalizedKbs,
+    };
+  }
+
+  if (mode === "kb_web") {
+    return {
+      enabledTools: uniqueStrings([...normalizedTools, "web_search"]),
+      knowledgeBases: normalizedKbs,
+    };
+  }
+
+  return {
+    enabledTools: normalizedTools.filter((tool) => tool !== "web_search"),
+    knowledgeBases: [],
+  };
+}