feat(ask_sb): Add search context into ask sourcebot toolbar (#397)

* new context selector

* ui nits

* move search context fetch to server

* feedback

* search context for chat suggestion, nits

* type nit

* fix minor ui nit

packages/web/src/actions.ts

@@ -1862,11 +1862,16 @@ export const getSearchContexts = async (domain: string) => sew(() =>
                 where: {
                     orgId: org.id,
                 },
+                include: {
+                    repos: true,
+                },
             });
 
             return searchContexts.map((context) => ({
+                id: context.id,
                 name: context.name,
                 description: context.description ?? undefined,
+                repoNames: context.repos.map((repo) => repo.name),
             }));
         }, /* minRequiredRole = */ OrgRole.GUEST), /* allowAnonymousAccess = */ true
     ));

packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx

@@ -3,15 +3,17 @@
 import { ResizablePanel } from '@/components/ui/resizable';
 import { ChatThread } from '@/features/chat/components/chatThread';
 import { LanguageModelInfo, SBChatMessage, SET_CHAT_STATE_QUERY_PARAM, SetChatStatePayload } from '@/features/chat/types';
-import { RepositoryQuery } from '@/lib/types';
+import { RepositoryQuery, SearchContextQuery } from '@/lib/types';
 import { CreateUIMessage } from 'ai';
 import { useRouter, useSearchParams } from 'next/navigation';
 import { useEffect, useState } from 'react';
 import { useChatId } from '../../useChatId';
+import { ContextItem } from '@/features/chat/components/chatBox/contextSelector';
 
 interface ChatThreadPanelProps {
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     order: number;
     messages: SBChatMessage[];
     isChatReadonly: boolean;
@@ -20,6 +22,7 @@ interface ChatThreadPanelProps {
 export const ChatThreadPanel = ({
     languageModels,
     repos,
+    searchContexts,
     order,
     messages,
     isChatReadonly,
@@ -31,8 +34,31 @@ export const ChatThreadPanel = ({
     const searchParams = useSearchParams();
     const [inputMessage, setInputMessage] = useState<CreateUIMessage<SBChatMessage> | undefined>(undefined);
 
-    // Use the last user's last message to determine what repos we should select by default.
+    // Use the last user's last message to determine what repos and contexts we should select by default.
-    const [selectedRepos, setSelectedRepos] = useState<string[]>(messages.findLast((message) => message.role === "user")?.metadata?.selectedRepos ?? []);
+    const lastUserMessage = messages.findLast((message) => message.role === "user");
+    const defaultSelectedRepos = lastUserMessage?.metadata?.selectedRepos ?? [];
+    const defaultSelectedContexts = lastUserMessage?.metadata?.selectedContexts ?? [];
+    
+    const [selectedItems, setSelectedItems] = useState<ContextItem[]>([
+        ...defaultSelectedRepos.map(repoName => {
+            const repoInfo = repos.find(r => r.repoName === repoName);
+            return {
+                type: 'repo' as const,
+                value: repoName,
+                name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
+                codeHostType: repoInfo?.codeHostType || ''
+            };
+        }),
+        ...defaultSelectedContexts.map(contextName => {
+            const context = searchContexts.find(c => c.name === contextName);
+            return {
+                type: 'context' as const,
+                value: contextName,
+                name: contextName,
+                repoCount: context?.repoNames.length || 0
+            };
+        })
+    ]);
 
     useEffect(() => {
         const setChatState = searchParams.get(SET_CHAT_STATE_QUERY_PARAM);
@@ -41,9 +67,28 @@ export const ChatThreadPanel = ({
         }
 
         try {
-            const { inputMessage, selectedRepos } = JSON.parse(setChatState) as SetChatStatePayload;
+            const { inputMessage, selectedRepos, selectedContexts } = JSON.parse(setChatState) as SetChatStatePayload;
             setInputMessage(inputMessage);
-            setSelectedRepos(selectedRepos);
+            setSelectedItems([
+                ...selectedRepos.map(repoName => {
+                    const repoInfo = repos.find(r => r.repoName === repoName);
+                    return {
+                        type: 'repo' as const,
+                        value: repoName,
+                        name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
+                        codeHostType: repoInfo?.codeHostType || ''
+                    };
+                }),
+                ...selectedContexts.map(contextName => {
+                    const context = searchContexts.find(c => c.name === contextName);
+                    return {
+                        type: 'context' as const,
+                        value: contextName,
+                        name: contextName,  
+                        repoCount: context?.repoNames.length || 0
+                    };
+                })
+            ]);
         } catch {
             console.error('Invalid message in URL');
         }
@@ -52,7 +97,7 @@ export const ChatThreadPanel = ({
         const newSearchParams = new URLSearchParams(searchParams.toString());
         newSearchParams.delete(SET_CHAT_STATE_QUERY_PARAM);
         router.replace(`?${newSearchParams.toString()}`, { scroll: false });
-    }, [searchParams, router]);
+    }, [searchParams, router, repos, searchContexts]);
 
     return (
         <ResizablePanel
@@ -67,8 +112,9 @@ export const ChatThreadPanel = ({
                     inputMessage={inputMessage}
                     languageModels={languageModels}
                     repos={repos}
-                    selectedRepos={selectedRepos}
+                    searchContexts={searchContexts}
-                    onSelectedReposChange={setSelectedRepos}
+                    selectedItems={selectedItems}
+                    onSelectedItemsChange={setSelectedItems}
                     isChatReadonly={isChatReadonly}
                 />
             </div>

packages/web/src/app/[domain]/chat/[id]/page.tsx

@@ -1,4 +1,4 @@
-import { getRepos } from '@/actions';
+import { getRepos, getSearchContexts } from '@/actions';
 import { getUserChatHistory, getConfiguredLanguageModelsInfo, getChatInfo } from '@/features/chat/actions';
 import { ServiceErrorException } from '@/lib/serviceError';
 import { isServiceError } from '@/lib/utils';
@@ -22,6 +22,7 @@ interface PageProps {
 export default async function Page({ params }: PageProps) {
     const languageModels = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(params.domain);
+    const searchContexts = await getSearchContexts(params.domain);
     const chatInfo = await getChatInfo({ chatId: params.id }, params.domain);
     const session = await auth();
     const chatHistory = session ? await getUserChatHistory(params.domain) : [];
@@ -34,6 +35,10 @@ export default async function Page({ params }: PageProps) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     if (isServiceError(chatInfo)) {
         if (chatInfo.statusCode === StatusCodes.NOT_FOUND) {
             return notFound();
@@ -74,6 +79,7 @@ export default async function Page({ params }: PageProps) {
                 <ChatThreadPanel
                     languageModels={languageModels}
                     repos={indexedRepos}
+                    searchContexts={searchContexts}
                     messages={messages}
                     order={2}
                     isChatReadonly={isReadonly}

packages/web/src/app/[domain]/chat/components/newChatPanel.tsx

@@ -6,29 +6,32 @@ import { ChatBoxToolbar } from "@/features/chat/components/chatBox/chatBoxToolba
 import { CustomSlateEditor } from "@/features/chat/customSlateEditor";
 import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
 import { LanguageModelInfo } from "@/features/chat/types";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { useCallback, useState } from "react";
 import { Descendant } from "slate";
 import { useLocalStorage } from "usehooks-ts";
+import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
 
 interface NewChatPanelProps {
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     order: number;
 }
 
 export const NewChatPanel = ({
     languageModels,
     repos,
+    searchContexts,
     order,
 }: NewChatPanelProps) => {
-    const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
+    const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
     const { createNewChatThread, isLoading } = useCreateNewChatThread();
-    const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
+    const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
 
     const onSubmit = useCallback((children: Descendant[]) => {
-        createNewChatThread(children, selectedRepos);
+        createNewChatThread(children, selectedItems);
-    }, [createNewChatThread, selectedRepos]);
+    }, [createNewChatThread, selectedItems]);
 
 
     return (
@@ -47,17 +50,19 @@ export const NewChatPanel = ({
                             preferredSuggestionsBoxPlacement="bottom-start"
                             isRedirecting={isLoading}
                             languageModels={languageModels}
-                            selectedRepos={selectedRepos}
+                            selectedItems={selectedItems}
-                            onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                            searchContexts={searchContexts}
+                            onContextSelectorOpenChanged={setIsContextSelectorOpen}
                         />
                         <div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
                             <ChatBoxToolbar
                                 languageModels={languageModels}
                                 repos={repos}
-                                selectedRepos={selectedRepos}
+                                searchContexts={searchContexts}
-                                onSelectedReposChange={setSelectedRepos}
+                                selectedItems={selectedItems}
-                                isRepoSelectorOpen={isRepoSelectorOpen}
+                                onSelectedItemsChange={setSelectedItems}
-                                onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                                isContextSelectorOpen={isContextSelectorOpen}
+                                onContextSelectorOpenChanged={setIsContextSelectorOpen}
                             />
                         </div>
                     </CustomSlateEditor>

packages/web/src/app/[domain]/chat/page.tsx

@@ -1,4 +1,4 @@
-import { getRepos } from "@/actions";
+import { getRepos, getSearchContexts } from "@/actions";
 import { getUserChatHistory, getConfiguredLanguageModelsInfo } from "@/features/chat/actions";
 import { ServiceErrorException } from "@/lib/serviceError";
 import { isServiceError } from "@/lib/utils";
@@ -18,6 +18,7 @@ interface PageProps {
 export default async function Page({ params }: PageProps) {
     const languageModels = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(params.domain);
+    const searchContexts = await getSearchContexts(params.domain);
     const session = await auth();
     const chatHistory = session ? await getUserChatHistory(params.domain) : [];
 
@@ -29,6 +30,10 @@ export default async function Page({ params }: PageProps) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     const indexedRepos = repos.filter((repo) => repo.indexedAt !== undefined);
 
     return (
@@ -48,6 +53,7 @@ export default async function Page({ params }: PageProps) {
                 <AnimatedResizableHandle />
                 <NewChatPanel
                     languageModels={languageModels}
+                    searchContexts={searchContexts}
                     repos={indexedRepos}
                     order={2}
                 />

packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx

@@ -8,7 +8,7 @@ import { LanguageModelInfo } from "@/features/chat/types";
 import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
 import { resetEditor } from "@/features/chat/utils";
 import { useDomain } from "@/hooks/useDomain";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { getDisplayTime } from "@/lib/utils";
 import { BrainIcon, FileIcon, LucideIcon, SearchIcon } from "lucide-react";
 import Link from "next/link";
@@ -16,6 +16,7 @@ import { ReactNode, useCallback, useEffect, useRef, useState } from "react";
 import { ReactEditor, useSlate } from "slate-react";
 import { SearchModeSelector, SearchModeSelectorProps } from "./toolbar";
 import { useLocalStorage } from "usehooks-ts";
+import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
 
 // @todo: we should probably rename this to a different type since it sort-of clashes
 // with the Suggestion system we have built into the chat box.
@@ -109,6 +110,7 @@ interface AgenticSearchProps {
     searchModeSelectorProps: SearchModeSelectorProps;
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     chatHistory: {
         id: string;
         createdAt: Date;
@@ -120,15 +122,16 @@ export const AgenticSearch = ({
     searchModeSelectorProps,
     languageModels,
     repos,
+    searchContexts,
     chatHistory,
 }: AgenticSearchProps) => {
     const [selectedSuggestionType, _setSelectedSuggestionType] = useState<SuggestionType | undefined>(undefined);
     const { createNewChatThread, isLoading } = useCreateNewChatThread();
     const dropdownRef = useRef<HTMLDivElement>(null);
     const editor = useSlate();
-    const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
+    const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
     const domain = useDomain();
-    const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
+    const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
 
     const setSelectedSuggestionType = useCallback((type: SuggestionType | undefined) => {
         _setSelectedSuggestionType(type);
@@ -158,13 +161,14 @@ export const AgenticSearch = ({
             >
                 <ChatBox
                     onSubmit={(children) => {
-                        createNewChatThread(children, selectedRepos);
+                        createNewChatThread(children, selectedItems);
                     }}
                     className="min-h-[50px]"
                     isRedirecting={isLoading}
                     languageModels={languageModels}
-                    selectedRepos={selectedRepos}
+                    selectedItems={selectedItems}
-                    onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                    searchContexts={searchContexts}
+                    onContextSelectorOpenChanged={setIsContextSelectorOpen}
                 />
                 <Separator />
                 <div className="relative">
@@ -172,10 +176,11 @@ export const AgenticSearch = ({
                         <ChatBoxToolbar
                             languageModels={languageModels}
                             repos={repos}
-                            selectedRepos={selectedRepos}
+                            searchContexts={searchContexts}
-                            onSelectedReposChange={setSelectedRepos}
+                            selectedItems={selectedItems}
-                            isRepoSelectorOpen={isRepoSelectorOpen}
+                            onSelectedItemsChange={setSelectedItems}
-                            onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                            isContextSelectorOpen={isContextSelectorOpen}
+                            onContextSelectorOpenChanged={setIsContextSelectorOpen}
                         />
                         <SearchModeSelector
                             {...searchModeSelectorProps}
@@ -201,7 +206,7 @@ export const AgenticSearch = ({
                                         setSelectedSuggestionType(undefined);
 
                                         if (openRepoSelector) {
-                                            setIsRepoSelectorOpen(true);
+                                            setIsContextSelectorOpen(true);
                                         } else {
                                             ReactEditor.focus(editor);
                                         }

packages/web/src/app/[domain]/components/homepage/index.tsx

@@ -2,7 +2,7 @@
 
 import { SourcebotLogo } from "@/app/components/sourcebotLogo";
 import { LanguageModelInfo } from "@/features/chat/types";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { useHotkeys } from "react-hotkeys-hook";
 import { AgenticSearch } from "./agenticSearch";
 import { PreciseSearch } from "./preciseSearch";
@@ -13,6 +13,7 @@ import { useCallback, useState } from "react";
 
 interface HomepageProps {
     initialRepos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     languageModels: LanguageModelInfo[];
     chatHistory: {
         id: string;
@@ -25,6 +26,7 @@ interface HomepageProps {
 
 export const Homepage = ({
     initialRepos,
+    searchContexts,
     languageModels,
     chatHistory,
     initialSearchMode,
@@ -82,6 +84,7 @@ export const Homepage = ({
                         }}
                         languageModels={languageModels}
                         repos={initialRepos}
+                        searchContexts={searchContexts}
                         chatHistory={chatHistory}
                     />
                 </CustomSlateEditor>

packages/web/src/app/[domain]/page.tsx

@@ -1,4 +1,4 @@
-import { getRepos } from "@/actions";
+import { getRepos, getSearchContexts } from "@/actions";
 import { Footer } from "@/app/components/footer";
 import { getOrgFromDomain } from "@/data/org";
 import { getConfiguredLanguageModelsInfo, getUserChatHistory } from "@/features/chat/actions";
@@ -22,12 +22,17 @@ export default async function Home({ params: { domain } }: { params: { domain: s
 
     const models = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(domain);
+    const searchContexts = await getSearchContexts(domain);
     const chatHistory = session ? await getUserChatHistory(domain) : [];
 
     if (isServiceError(repos)) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     if (isServiceError(chatHistory)) {
         throw new ServiceErrorException(chatHistory);
     }
@@ -52,6 +57,7 @@ export default async function Home({ params: { domain } }: { params: { domain: s
 
             <Homepage
                 initialRepos={indexedRepos}
+                searchContexts={searchContexts}
                 languageModels={models}
                 chatHistory={chatHistory}
                 initialSearchMode={initialSearchMode}

packages/web/src/app/api/(server)/chat/route.ts

@@ -64,11 +64,12 @@ export async function POST(req: Request) {
         return serviceErrorResponse(schemaValidationError(parsed.error));
     }
 
-    const { messages, id, selectedRepos, languageModelId } = parsed.data;
+    const { messages, id, selectedRepos, selectedContexts, languageModelId } = parsed.data;
     const response = await chatHandler({
         messages,
         id,
         selectedRepos,
+        selectedContexts,
         languageModelId,
     }, domain);
 
@@ -93,10 +94,11 @@ interface ChatHandlerProps {
     messages: SBChatMessage[];
     id: string;
     selectedRepos: string[];
+    selectedContexts?: string[];
     languageModelId: string;
 }
 
-const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandlerProps, domain: string) => sew(async () =>
+const chatHandler = ({ messages, id, selectedRepos, selectedContexts, languageModelId }: ChatHandlerProps, domain: string) => sew(async () =>
     withAuth((userId) =>
         withOrgMembership(userId, domain, async ({ org }) => {
             const chat = await prisma.chat.findUnique({
@@ -186,13 +188,34 @@ const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandl
 
                     const startTime = new Date();
 
+                    // Expand search contexts to repos
+                    let expandedRepos = [...selectedRepos];
+                    if (selectedContexts && selectedContexts.length > 0) {
+                        const searchContexts = await prisma.searchContext.findMany({
+                            where: {
+                                orgId: org.id,
+                                name: { in: selectedContexts }
+                            },
+                            include: {
+                                repos: true
+                            }
+                        });
+                        
+                        const contextRepos = searchContexts.flatMap(context => 
+                            context.repos.map(repo => repo.name)
+                        );
+                        
+                        // Combine and deduplicate repos
+                        expandedRepos = Array.from(new Set([...selectedRepos, ...contextRepos]));
+                    }
+
                     const researchStream = await createAgentStream({
                         model,
                         providerOptions,
                         headers,
                         inputMessages: messageHistory,
                         inputSources: sources,
-                        selectedRepos,
+                        selectedRepos: expandedRepos,
                         onWriteSource: (source) => {
                             writer.write({
                                 type: 'data-source',

packages/web/src/features/chat/components/chatBox/chatBox.tsx

@@ -18,6 +18,8 @@ import { Suggestion } from "./types";
 import { useSuggestionModeAndQuery } from "./useSuggestionModeAndQuery";
 import { useSuggestionsData } from "./useSuggestionsData";
 import { useToast } from "@/components/hooks/use-toast";
+import { ContextItem } from "./contextSelector";
+import { SearchContextQuery } from "@/lib/types";
 
 interface ChatBoxProps {
     onSubmit: (children: Descendant[], editor: CustomEditor) => void;
@@ -27,8 +29,9 @@ interface ChatBoxProps {
     isRedirecting?: boolean;
     isGenerating?: boolean;
     languageModels: LanguageModelInfo[];
-    selectedRepos: string[];
+    selectedItems: ContextItem[];
-    onRepoSelectorOpenChanged: (isOpen: boolean) => void;
+    searchContexts: SearchContextQuery[];
+    onContextSelectorOpenChanged: (isOpen: boolean) => void;
 }
 
 export const ChatBox = ({
@@ -39,8 +42,9 @@ export const ChatBox = ({
     isRedirecting,
     isGenerating,
     languageModels,
-    selectedRepos,
+    selectedItems,
-    onRepoSelectorOpenChanged,
+    searchContexts,
+    onContextSelectorOpenChanged,
 }: ChatBoxProps) => {
     const suggestionsBoxRef = useRef<HTMLDivElement>(null);
     const [index, setIndex] = useState(0);
@@ -49,7 +53,20 @@ export const ChatBox = ({
     const { suggestions, isLoading } = useSuggestionsData({
         suggestionMode,
         suggestionQuery,
-        selectedRepos,
+        selectedRepos: selectedItems.map((item) => {
+            if (item.type === 'repo') {
+                return [item.value];
+            }
+
+            if (item.type === 'context') {
+                const context = searchContexts.find((context) => context.name === item.value);
+                if (context) {
+                    return context.repoNames;
+                }
+            }
+
+            return [];
+        }).flat(),
     });
     const { selectedLanguageModel } = useSelectedLanguageModel({
         initialLanguageModel: languageModels.length > 0 ? languageModels[0] : undefined,
@@ -113,7 +130,7 @@ export const ChatBox = ({
             }
         }
 
-        if (selectedRepos.length === 0) {
+        if (selectedItems.length === 0) {
             return {
                 isSubmitDisabled: true,
                 isSubmitDisabledReason: "no-repos-selected",
@@ -137,7 +154,7 @@ export const ChatBox = ({
         editor.children,
         isRedirecting,
         isGenerating,
-        selectedRepos.length,
+        selectedItems.length,
         selectedLanguageModel,
     ])
 
@@ -145,17 +162,17 @@ export const ChatBox = ({
         if (isSubmitDisabled) {
             if (isSubmitDisabledReason === "no-repos-selected") {
                 toast({
-                    description: "⚠️ One or more repositories must be selected.",
+                    description: "⚠️ One or more repositories or search contexts must be selected.",
                     variant: "destructive",
                 });
-                onRepoSelectorOpenChanged(true);
+                onContextSelectorOpenChanged(true);
             }
 
             return;
         }
 
         _onSubmit(editor.children, editor);
-    }, [_onSubmit, editor, isSubmitDisabled, isSubmitDisabledReason, toast, onRepoSelectorOpenChanged]);
+    }, [_onSubmit, editor, isSubmitDisabled, isSubmitDisabledReason, toast, onContextSelectorOpenChanged]);
 
     const onInsertSuggestion = useCallback((suggestion: Suggestion) => {
         switch (suggestion.type) {
@@ -322,7 +339,7 @@ export const ChatBox = ({
                                 <TooltipContent>
                                     <div className="flex flex-row items-center">
                                         <TriangleAlertIcon className="h-4 w-4 text-warning mr-1" />
-                                        <span className="text-destructive">One or more repositories must be selected.</span>
+                                        <span className="text-destructive">One or more repositories or search contexts must be selected.</span>
                                     </div>
                                 </TooltipContent>
                             )}

packages/web/src/features/chat/components/chatBox/chatBoxToolbar.tsx

@@ -5,31 +5,33 @@ import { Button } from "@/components/ui/button";
 import { Separator } from "@/components/ui/separator";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { LanguageModelInfo } from "@/features/chat/types";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { AtSignIcon } from "lucide-react";
 import { useCallback } from "react";
 import { useHotkeys } from "react-hotkeys-hook";
 import { ReactEditor, useSlate } from "slate-react";
 import { useSelectedLanguageModel } from "../../useSelectedLanguageModel";
 import { LanguageModelSelector } from "./languageModelSelector";
-import { RepoSelector } from "./repoSelector";
+import { ContextSelector, type ContextItem } from "./contextSelector";
 
 export interface ChatBoxToolbarProps {
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
-    selectedRepos: string[];
+    searchContexts: SearchContextQuery[];
-    onSelectedReposChange: (repos: string[]) => void;
+    selectedItems: ContextItem[];
-    isRepoSelectorOpen: boolean;
+    onSelectedItemsChange: (items: ContextItem[]) => void;
-    onRepoSelectorOpenChanged: (isOpen: boolean) => void;
+    isContextSelectorOpen: boolean;
+    onContextSelectorOpenChanged: (isOpen: boolean) => void;
 }
 
 export const ChatBoxToolbar = ({
     languageModels,
     repos,
-    selectedRepos,
+    searchContexts,
-    onSelectedReposChange,
+    selectedItems,
-    isRepoSelectorOpen,
+    onSelectedItemsChange,
-    onRepoSelectorOpenChanged,
+    isContextSelectorOpen,
+    onContextSelectorOpenChanged,
 }: ChatBoxToolbarProps) => {
     const editor = useSlate();
     
@@ -76,17 +78,18 @@ export const ChatBoxToolbar = ({
             <Separator orientation="vertical" className="h-3 mx-1" />
             <Tooltip>
                 <TooltipTrigger asChild>
-                    <RepoSelector
+                    <ContextSelector
                         className="bg-inherit w-fit h-6 min-h-6"
-                        repos={repos.map((repo) => repo.repoName)}
+                        repos={repos}
-                        selectedRepos={selectedRepos}
+                        searchContexts={searchContexts}
-                        onSelectedReposChange={onSelectedReposChange}
+                        selectedItems={selectedItems}
-                        isOpen={isRepoSelectorOpen}
+                        onSelectedItemsChange={onSelectedItemsChange}
-                        onOpenChanged={onRepoSelectorOpenChanged}
+                        isOpen={isContextSelectorOpen}
+                        onOpenChanged={onContextSelectorOpenChanged}
                     />
                 </TooltipTrigger>
                 <TooltipContent side="bottom">
-                    <span>Repositories to scope conversation to.</span>
+                    <span>Search contexts and repositories to scope conversation to.</span>
                 </TooltipContent>
             </Tooltip>
             {languageModels.length > 0 && (

packages/web/src/features/chat/components/chatBox/contextSelector.tsx

@@ -0,0 +1,279 @@
+// Adapted from: web/src/components/ui/multi-select.tsx
+
+import * as React from "react";
+import {
+    CheckIcon,
+    ChevronDown,
+    FolderIcon,
+    LayersIcon,
+    LibraryBigIcon,
+} from "lucide-react";
+import Image from "next/image";
+
+import { cn, getCodeHostIcon } from "@/lib/utils";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
+import { Button } from "@/components/ui/button";
+import { Badge } from "@/components/ui/badge";
+import {
+    Popover,
+    PopoverContent,
+    PopoverTrigger,
+} from "@/components/ui/popover";
+import {
+    Command,
+    CommandEmpty,
+    CommandGroup,
+    CommandInput,
+    CommandItem,
+    CommandList,
+    CommandSeparator,
+} from "@/components/ui/command";
+
+export type RepoContextItem = {
+    type: 'repo';
+    value: string;
+    name: string;
+    codeHostType: string;
+}
+
+export type SearchContextItem = {
+    type: 'context';
+    value: string;
+    name: string;
+    repoCount: number;
+}
+
+export type ContextItem = RepoContextItem | SearchContextItem;
+
+interface ContextSelectorProps extends React.ButtonHTMLAttributes<HTMLButtonElement> {
+    repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
+    selectedItems: ContextItem[];
+    onSelectedItemsChange: (items: ContextItem[]) => void;
+    className?: string;
+    isOpen: boolean;
+    onOpenChanged: (isOpen: boolean) => void;
+}
+
+export const ContextSelector = React.forwardRef<
+    HTMLButtonElement,
+    ContextSelectorProps
+>(
+    (
+        {
+            repos,
+            searchContexts,
+            onSelectedItemsChange,
+            className,
+            selectedItems,
+            isOpen,
+            onOpenChanged,
+            ...props
+        },
+        ref
+    ) => {
+        const scrollContainerRef = React.useRef<HTMLDivElement>(null);
+        const scrollPosition = React.useRef<number>(0);
+
+        const handleInputKeyDown = (
+            event: React.KeyboardEvent<HTMLInputElement>
+        ) => {
+            if (event.key === "Enter") {
+                onOpenChanged(true);
+            } else if (event.key === "Backspace" && !event.currentTarget.value) {
+                const newSelectedItems = [...selectedItems];
+                newSelectedItems.pop();
+                onSelectedItemsChange(newSelectedItems);
+            }
+        };
+
+        const toggleItem = (item: ContextItem) => {
+            // Store current scroll position before state update
+            if (scrollContainerRef.current) {
+                scrollPosition.current = scrollContainerRef.current.scrollTop;
+            }
+
+            const isSelected = selectedItems.some(
+                (selected) => selected.type === item.type && selected.value === item.value
+            );
+            
+            const newSelectedItems = isSelected
+                ? selectedItems.filter(
+                    (selected) => !(selected.type === item.type && selected.value === item.value)
+                  )
+                : [...selectedItems, item];
+            onSelectedItemsChange(newSelectedItems);
+        };
+
+        const handleClear = () => {
+            onSelectedItemsChange([]);
+        };
+
+        const handleTogglePopover = () => {
+            onOpenChanged(!isOpen);
+        };
+
+        const allItems = React.useMemo(() => {
+            const contextItems: ContextItem[] = searchContexts.map(context => ({
+                type: 'context' as const,
+                value: context.name,
+                name: context.name,
+                repoCount: context.repoNames.length
+            }));
+            
+            const repoItems: ContextItem[] = repos.map(repo => ({
+                type: 'repo' as const,
+                value: repo.repoName,
+                name: repo.repoDisplayName || repo.repoName.split('/').pop() || repo.repoName,
+                codeHostType: repo.codeHostType,
+            }));
+            
+            return [...contextItems, ...repoItems];
+        }, [repos, searchContexts]);
+
+        const sortedItems = React.useMemo(() => {
+            return allItems
+                .map((item) => ({
+                    item,
+                    isSelected: selectedItems.some(
+                        (selected) => selected.type === item.type && selected.value === item.value
+                    )
+                }))
+                .sort((a, b) => {
+                    // Selected items first
+                    if (a.isSelected && !b.isSelected) return -1;
+                    if (!a.isSelected && b.isSelected) return 1;
+                    // Then contexts before repos
+                    if (a.item.type === 'context' && b.item.type === 'repo') return -1;
+                    if (a.item.type === 'repo' && b.item.type === 'context') return 1;
+                    return 0;
+                })
+        }, [allItems, selectedItems]);
+
+        // Restore scroll position after re-render
+        React.useEffect(() => {
+            if (scrollContainerRef.current && scrollPosition.current > 0) {
+                scrollContainerRef.current.scrollTop = scrollPosition.current;
+            }
+        }, [sortedItems]);
+
+        return (
+            <Popover
+                open={isOpen}
+                onOpenChange={onOpenChanged}
+            >
+                <PopoverTrigger asChild>
+                    <Button
+                        ref={ref}
+                        {...props}
+                        onClick={handleTogglePopover}
+                        className={cn(
+                            "flex p-1 rounded-md items-center justify-between bg-inherit h-6",
+                            className
+                        )}
+                    >
+                        <div className="flex items-center justify-between w-full mx-auto">
+                            <LayersIcon className="h-4 w-4 text-muted-foreground mr-1" />
+                            <span
+                                className={cn("text-sm text-muted-foreground mx-1 font-medium")}
+                            >
+                                {
+                                    selectedItems.length === 0 ? `Select context` :
+                                        selectedItems.length === 1 ? selectedItems[0].name :
+                                            `${selectedItems.length} selected`
+                                }
+                            </span>
+                            <ChevronDown className="h-4 cursor-pointer text-muted-foreground ml-2" />
+                        </div>
+                    </Button>
+                </PopoverTrigger>
+                <PopoverContent
+                    className="w-auto p-0"
+                    align="start"
+                    onEscapeKeyDown={() => onOpenChanged(false)}
+                >
+                    <Command>
+                        <CommandInput
+                            placeholder="Search contexts and repos..."
+                            onKeyDown={handleInputKeyDown}
+                        />
+                        <CommandList ref={scrollContainerRef}>
+                            <CommandEmpty>No results found.</CommandEmpty>
+                            <CommandGroup>
+                                {sortedItems.map(({ item, isSelected }) => {
+                                    return (
+                                        <CommandItem
+                                            key={`${item.type}-${item.value}`}
+                                            onSelect={() => toggleItem(item)}
+                                            className="cursor-pointer"
+                                        >
+                                            <div
+                                                className={cn(
+                                                    "mr-2 flex h-4 w-4 items-center justify-center rounded-sm border border-primary",
+                                                    isSelected
+                                                        ? "bg-primary text-primary-foreground"
+                                                        : "opacity-50 [&_svg]:invisible"
+                                                )}
+                                            >
+                                                <CheckIcon className="h-4 w-4" />
+                                            </div>
+                                            <div className="flex items-center gap-2 flex-1">
+                                                {item.type === 'context' ? (
+                                                    <LibraryBigIcon className="h-4 w-4 text-muted-foreground" />
+                                                ) : (
+                                                    // Render code host icon for repos
+                                                    (() => {
+                                                        const codeHostIcon = item.codeHostType ? getCodeHostIcon(item.codeHostType) : null;
+                                                        return codeHostIcon ? (
+                                                            <Image 
+                                                                src={codeHostIcon.src} 
+                                                                alt={`${item.codeHostType} icon`}
+                                                                width={16}
+                                                                height={16}
+                                                                className={cn("h-4 w-4", codeHostIcon.className)}
+                                                            />
+                                                        ) : (
+                                                            <FolderIcon className="h-4 w-4 text-muted-foreground" />
+                                                        );
+                                                    })()
+                                                )}
+                                                <div className="flex flex-col flex-1">
+                                                    <div className="flex items-center gap-2">
+                                                        <span className="font-medium">
+                                                            {item.name}
+                                                        </span>
+                                                        {item.type === 'context' && (
+                                                            <Badge 
+                                                                variant="default" 
+                                                                className="text-[10px] px-1.5 py-0 h-4 bg-primary text-primary-foreground"
+                                                            >
+                                                                {item.repoCount} repo{item.repoCount === 1 ? '' : 's'}
+                                                            </Badge>
+                                                        )}
+                                                    </div>
+                                                </div>
+                                            </div>
+                                        </CommandItem>
+                                    );
+                                })}
+                            </CommandGroup>
+                        </CommandList>
+                        {selectedItems.length > 0 && (
+                            <>
+                                <CommandSeparator />
+                                <CommandItem
+                                    onSelect={handleClear}
+                                    className="flex-1 justify-center cursor-pointer"
+                                >
+                                    Clear
+                                </CommandItem>
+                            </>
+                        )}
+                    </Command>
+                </PopoverContent>
+            </Popover>
+        );
+    }
+);
+
+ContextSelector.displayName = "ContextSelector";

packages/web/src/features/chat/components/chatBox/repoSelector.tsx

@@ -1,191 +0,0 @@
-// Adapted from: web/src/components/ui/multi-select.tsx
-
-import * as React from "react";
-import {
-    CheckIcon,
-    ChevronDown,
-    BookMarkedIcon,
-} from "lucide-react";
-
-import { cn } from "@/lib/utils";
-import { Button } from "@/components/ui/button";
-import {
-    Popover,
-    PopoverContent,
-    PopoverTrigger,
-} from "@/components/ui/popover";
-import {
-    Command,
-    CommandEmpty,
-    CommandGroup,
-    CommandInput,
-    CommandItem,
-    CommandList,
-    CommandSeparator,
-} from "@/components/ui/command";
-
-interface RepoSelectorProps extends React.ButtonHTMLAttributes<HTMLButtonElement> {
-    repos: string[];
-    selectedRepos: string[];
-    onSelectedReposChange: (repos: string[]) => void;
-    className?: string;
-    isOpen: boolean;
-    onOpenChanged: (isOpen: boolean) => void;
-}
-
-export const RepoSelector = React.forwardRef<
-    HTMLButtonElement,
-    RepoSelectorProps
->(
-    (
-        {
-            repos,
-            onSelectedReposChange,
-            className,
-            selectedRepos,
-            isOpen,
-            onOpenChanged,
-            ...props
-        },
-        ref
-    ) => {
-        const scrollContainerRef = React.useRef<HTMLDivElement>(null);
-        const scrollPosition = React.useRef<number>(0);
-
-        const handleInputKeyDown = (
-            event: React.KeyboardEvent<HTMLInputElement>
-        ) => {
-            if (event.key === "Enter") {
-                onOpenChanged(true);
-            } else if (event.key === "Backspace" && !event.currentTarget.value) {
-                const newSelectedRepos = [...selectedRepos];
-                newSelectedRepos.pop();
-                onSelectedReposChange(newSelectedRepos);
-            }
-        };
-
-        const toggleRepo = (repo: string) => {
-            // Store current scroll position before state update
-            if (scrollContainerRef.current) {
-                scrollPosition.current = scrollContainerRef.current.scrollTop;
-            }
-
-            const newSelectedValues = selectedRepos.includes(repo)
-                ? selectedRepos.filter((value) => value !== repo)
-                : [...selectedRepos, repo];
-            onSelectedReposChange(newSelectedValues);
-        };
-
-        const handleClear = () => {
-            onSelectedReposChange([]);
-        };
-
-        const handleTogglePopover = () => {
-            onOpenChanged(!isOpen);
-        };
-
-        const sortedRepos = React.useMemo(() => {
-            return repos
-                .map((repo) => ({
-                    repo,
-                    isSelected: selectedRepos.includes(repo)
-                }))
-                .sort((a, b) => {
-                    if (a.isSelected && !b.isSelected) return -1;
-                    if (!a.isSelected && b.isSelected) return 1;
-                    return 0;
-                })
-        }, [repos, selectedRepos]);
-
-        // Restore scroll position after re-render
-        React.useEffect(() => {
-            if (scrollContainerRef.current && scrollPosition.current > 0) {
-                scrollContainerRef.current.scrollTop = scrollPosition.current;
-            }
-        }, [sortedRepos]);
-
-        return (
-            <Popover
-                open={isOpen}
-                onOpenChange={onOpenChanged}
-            >
-                <PopoverTrigger asChild>
-                    <Button
-                        ref={ref}
-                        {...props}
-                        onClick={handleTogglePopover}
-                        className={cn(
-                            "flex p-1 rounded-md items-center justify-between bg-inherit h-6",
-                            className
-                        )}
-                    >
-                        <div className="flex items-center justify-between w-full mx-auto">
-                            <BookMarkedIcon className="h-4 w-4 text-muted-foreground mr-1" />
-                            <span
-                                className={cn("text-sm text-muted-foreground mx-1 font-medium")}
-                            >
-                                {
-                                    selectedRepos.length === 0 ? `Select a repo` :
-                                        selectedRepos.length === 1 ? `${selectedRepos[0].split('/').pop()}` :
-                                            `${selectedRepos.length} repo${selectedRepos.length === 1 ? '' : 's'}`
-                                }
-                            </span>
-                            <ChevronDown className="h-4 cursor-pointer text-muted-foreground ml-2" />
-                        </div>
-                    </Button>
-                </PopoverTrigger>
-                <PopoverContent
-                    className="w-auto p-0"
-                    align="start"
-                    onEscapeKeyDown={() => onOpenChanged(false)}
-                >
-                    <Command>
-                        <CommandInput
-                            placeholder="Search repos..."
-                            onKeyDown={handleInputKeyDown}
-                        />
-                        <CommandList ref={scrollContainerRef}>
-                            <CommandEmpty>No results found.</CommandEmpty>
-                            <CommandGroup>
-                                {sortedRepos.map(({ repo, isSelected }) => {
-                                    return (
-                                        <CommandItem
-                                            key={repo}
-                                            onSelect={() => toggleRepo(repo)}
-                                            className="cursor-pointer"
-                                        >
-                                            <div
-                                                className={cn(
-                                                    "mr-2 flex h-4 w-4 items-center justify-center rounded-sm border border-primary",
-                                                    isSelected
-                                                        ? "bg-primary text-primary-foreground"
-                                                        : "opacity-50 [&_svg]:invisible"
-                                                )}
-                                            >
-                                                <CheckIcon className="h-4 w-4" />
-                                            </div>
-                                            <span>{repo}</span>
-                                        </CommandItem>
-                                    );
-                                })}
-                            </CommandGroup>
-                        </CommandList>
-                        {selectedRepos.length > 0 && (
-                            <>
-                                <CommandSeparator />
-                                <CommandItem
-                                    onSelect={handleClear}
-                                    className="flex-1 justify-center cursor-pointer"
-                                >
-                                    Clear
-                                </CommandItem>
-                            </>
-                        )}
-                    </Command>
-                </PopoverContent>
-            </Popover>
-        );
-    }
-);
-
-RepoSelector.displayName = "RepoSelector";

packages/web/src/features/chat/components/chatThread/chatThread.tsx

@@ -12,7 +12,7 @@ import { useChat } from '@ai-sdk/react';
 import { CreateUIMessage, DefaultChatTransport } from 'ai';
 import { ArrowDownIcon } from 'lucide-react';
 import { useNavigationGuard } from 'next-navigation-guard';
-import { Fragment, useCallback, useEffect, useRef, useState } from 'react';
+import { Fragment, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { Descendant } from 'slate';
 import { useMessagePairs } from '../../useMessagePairs';
 import { useSelectedLanguageModel } from '../../useSelectedLanguageModel';
@@ -22,7 +22,8 @@ import { ChatThreadListItem } from './chatThreadListItem';
 import { ErrorBanner } from './errorBanner';
 import { useRouter } from 'next/navigation';
 import { usePrevious } from '@uidotdev/usehooks';
-import { RepositoryQuery } from '@/lib/types';
+import { RepositoryQuery, SearchContextQuery } from '@/lib/types';
+import { ContextItem } from '../chatBox/contextSelector';
 
 type ChatHistoryState = {
     scrollOffset?: number;
@@ -34,8 +35,9 @@ interface ChatThreadProps {
     inputMessage?: CreateUIMessage<SBChatMessage>;
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
-    selectedRepos: string[];
+    searchContexts: SearchContextQuery[];
-    onSelectedReposChange: (repos: string[]) => void;
+    selectedItems: ContextItem[];
+    onSelectedItemsChange: (items: ContextItem[]) => void;
     isChatReadonly: boolean;
 }
 
@@ -45,8 +47,9 @@ export const ChatThread = ({
     inputMessage,
     languageModels,
     repos,
-    selectedRepos,
+    searchContexts,
-    onSelectedReposChange,
+    selectedItems,
+    onSelectedItemsChange,
     isChatReadonly,
 }: ChatThreadProps) => {
     const domain = useDomain();
@@ -57,7 +60,13 @@ export const ChatThread = ({
     const [isAutoScrollEnabled, setIsAutoScrollEnabled] = useState(false);
     const { toast } = useToast();
     const router = useRouter();
-    const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
+    const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
+
+    const { selectedRepos, selectedContexts } = useMemo(() => {
+        const repos = selectedItems.filter(item => item.type === 'repo').map(item => item.value);
+        const contexts = selectedItems.filter(item => item.type === 'context').map(item => item.value);
+        return { selectedRepos: repos, selectedContexts: contexts };
+    }, [selectedItems]);
 
     // Initial state is from attachments that exist in in the chat history.
     const [sources, setSources] = useState<Source[]>(
@@ -114,10 +123,11 @@ export const ChatThread = ({
         _sendMessage(message, {
             body: {
                 selectedRepos,
+                selectedContexts,
                 languageModelId: selectedLanguageModel.model,
             } satisfies AdditionalChatRequestParams,
         });
-    }, [_sendMessage, selectedLanguageModel, selectedRepos, toast]);
+    }, [_sendMessage, selectedLanguageModel, toast, selectedRepos, selectedContexts]);
 
 
     const messagePairs = useMessagePairs(messages);
@@ -233,13 +243,13 @@ export const ChatThread = ({
         const text = slateContentToString(children);
         const mentions = getAllMentionElements(children);
 
-        const message = createUIMessage(text, mentions.map(({ data }) => data), selectedRepos);
+        const message = createUIMessage(text, mentions.map(({ data }) => data), selectedRepos, selectedContexts);
         sendMessage(message);
 
         setIsAutoScrollEnabled(true);
 
         resetEditor(editor);
-    }, [sendMessage, selectedRepos]);
+    }, [sendMessage, selectedRepos, selectedContexts]);
 
     return (
         <>
@@ -317,17 +327,19 @@ export const ChatThread = ({
                             isGenerating={status === "streaming" || status === "submitted"}
                             onStop={stop}
                             languageModels={languageModels}
-                            selectedRepos={selectedRepos}
+                            selectedItems={selectedItems}
-                            onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                            searchContexts={searchContexts}
+                            onContextSelectorOpenChanged={setIsContextSelectorOpen}
                         />
                         <div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
                             <ChatBoxToolbar
                                 languageModels={languageModels}
                                 repos={repos}
-                                selectedRepos={selectedRepos}
+                                searchContexts={searchContexts}
-                                onSelectedReposChange={onSelectedReposChange}
+                                selectedItems={selectedItems}
-                                isRepoSelectorOpen={isRepoSelectorOpen}
+                                onSelectedItemsChange={onSelectedItemsChange}
-                                onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                                isContextSelectorOpen={isContextSelectorOpen}
+                                onContextSelectorOpenChanged={setIsContextSelectorOpen}
                             />
                         </div>
                     </CustomSlateEditor>

packages/web/src/features/chat/types.ts

@@ -51,6 +51,7 @@ export const sbChatMessageMetadataSchema = z.object({
         userId: z.string(),
     })).optional(),
     selectedRepos: z.array(z.string()).optional(),
+    selectedContexts: z.array(z.string()).optional(),
     traceId: z.string().optional(),
 });
 
@@ -139,6 +140,7 @@ export const SET_CHAT_STATE_QUERY_PARAM = 'setChatState';
 export type SetChatStatePayload = {
     inputMessage: CreateUIMessage<SBChatMessage>;
     selectedRepos: string[];
+    selectedContexts: string[];
 }
 
 
@@ -156,5 +158,6 @@ export type LanguageModelInfo = {
 export const additionalChatRequestParamsSchema = z.object({
     languageModelId: z.string(),
     selectedRepos: z.array(z.string()),
+    selectedContexts: z.array(z.string()),
 });
 export type AdditionalChatRequestParams = z.infer<typeof additionalChatRequestParamsSchema>;

packages/web/src/features/chat/useCreateNewChatThread.ts

@@ -11,6 +11,7 @@ import { createChat } from "./actions";
 import { isServiceError } from "@/lib/utils";
 import { createPathWithQueryParams } from "@/lib/utils";
 import { SET_CHAT_STATE_QUERY_PARAM, SetChatStatePayload } from "./types";
+import { ContextItem } from "./components/chatBox/contextSelector";
 
 export const useCreateNewChatThread = () => {
     const domain = useDomain();
@@ -18,10 +19,15 @@ export const useCreateNewChatThread = () => {
     const { toast } = useToast();
     const router = useRouter();
 
-    const createNewChatThread = useCallback(async (children: Descendant[], selectedRepos: string[]) => {
+    const createNewChatThread = useCallback(async (children: Descendant[], selectedItems: ContextItem[]) => {
         const text = slateContentToString(children);
         const mentions = getAllMentionElements(children);
-        const inputMessage = createUIMessage(text, mentions.map((mention) => mention.data), selectedRepos);
+        
+        // Extract repos and contexts from selectedItems
+        const selectedRepos = selectedItems.filter(item => item.type === 'repo').map(item => item.value);
+        const selectedContexts = selectedItems.filter(item => item.type === 'context').map(item => item.value);
+        
+        const inputMessage = createUIMessage(text, mentions.map((mention) => mention.data), selectedRepos, selectedContexts);
 
         setIsLoading(true);
         const response = await createChat(domain);
@@ -37,6 +43,7 @@ export const useCreateNewChatThread = () => {
             [SET_CHAT_STATE_QUERY_PARAM, JSON.stringify({
                 inputMessage,
                 selectedRepos,
+                selectedContexts,
             } satisfies SetChatStatePayload)],
         );
 

packages/web/src/features/chat/utils.ts

@@ -172,7 +172,7 @@ export const addLineNumbers = (source: string, lineOffset = 1) => {
     return source.split('\n').map((line, index) => `${index + lineOffset}:${line}`).join('\n');
 }
 
-export const createUIMessage = (text: string, mentions: MentionData[], selectedRepos: string[]): CreateUIMessage<SBChatMessage> => {
+export const createUIMessage = (text: string, mentions: MentionData[], selectedRepos: string[], selectedContexts: string[]): CreateUIMessage<SBChatMessage> => {
     // Converts applicable mentions into sources.
     const sources: Source[] = mentions
         .map((mention) => {
@@ -206,6 +206,7 @@ export const createUIMessage = (text: string, mentions: MentionData[], selectedR
         ],
         metadata: {
             selectedRepos,
+            selectedContexts,
         },
     }
 }

packages/web/src/lib/schemas.ts

@@ -28,6 +28,13 @@ export const repositoryQuerySchema = z.object({
     repoIndexingStatus: z.nativeEnum(RepoIndexingStatus),
 });
 
+export const searchContextQuerySchema = z.object({
+    id: z.number(),
+    name: z.string(),
+    description: z.string().optional(),
+    repoNames: z.array(z.string()),
+});
+
 export const verifyCredentialsRequestSchema = z.object({
     email: z.string().email(),
     password: z.string().min(8),

packages/web/src/lib/types.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { getVersionResponseSchema, repositoryQuerySchema } from "./schemas";
+import { getVersionResponseSchema, repositoryQuerySchema, searchContextQuerySchema } from "./schemas";
 import { tenancyModeSchema } from "@/env.mjs";
 
 export type KeymapType = "default" | "vim";
@@ -26,3 +26,4 @@ export type NewsItem = {
 
 export type TenancyMode = z.infer<typeof tenancyModeSchema>;
 export type RepositoryQuery = z.infer<typeof repositoryQuerySchema>;
+export type SearchContextQuery = z.infer<typeof searchContextQuerySchema>;