feat(ask): start chat from selected code range

packages/web/src/app/(app)/components/editorContextMenu.tsx

@@ -2,6 +2,7 @@
 
 import { useToast } from "@/components/hooks/use-toast";
 import { Button } from "@/components/ui/button";
+import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
 import useCaptureEvent from "@/hooks/useCaptureEvent";
 import { createPathWithQueryParams } from "@/lib/utils";
 import { autoPlacement, computePosition, offset, shift, VirtualElement } from "@floating-ui/react";
@@ -28,6 +29,7 @@ export const EditorContextMenu = ({
     const ref = useRef<HTMLDivElement>(null);
     const { toast } = useToast();
     const captureEvent = useCaptureEvent();
+    const { createChatFromSource } = useCreateNewChatThread();
     useEffect(() => {
         if (selection.empty) {
             ref.current?.classList.add('hidden');
@@ -126,19 +128,47 @@ export const EditorContextMenu = ({
         )
     }, [selection.from, selection.to, repoName, revisionName, path, toast, captureEvent, view]);
 
+    const onAskSourcebot = useCallback(() => {
+        if (selection.empty) {
+            return;
+        }
+
+        const startLine = view.state.doc.lineAt(selection.from).number;
+        const endLine = view.state.doc.lineAt(selection.to - 1).number;
+
+        void createChatFromSource({
+            type: 'file',
+            repo: repoName,
+            path,
+            name: path.split('/').pop() ?? path,
+            revision: revisionName,
+            range: { startLine, endLine },
+        });
+    }, [createChatFromSource, path, repoName, revisionName, selection, view]);
+
     return (
-        <div
-            ref={ref}
-            className="absolute z-10 flex flex-col gap-2 bg-background border border-gray-300 dark:border-gray-700 rounded-md shadow-lg p-2"
-        >
-            <Button
-                variant="ghost"
-                size="sm"
-                onClick={onCopyLinkToSelection}
-            >
-                <Link2Icon className="h-4 w-4 mr-1" />
-                Share selection
-            </Button>
-        </div>
+<div
+  ref={ref}
+  className="absolute z-10 flex items-center gap-1 rounded-lg border border-border bg-popover p-1 shadow-xl"
+>
+  <Button
+    variant="secondary"
+    size="sm"
+    onClick={onCopyLinkToSelection}
+    className="h-8 px-3 hover:bg-black/50"
+  >
+    <Link2Icon className="mr-2 h-4 w-4" />
+    Share selection
+  </Button>
+
+  <Button
+    variant="secondary"
+    size="sm"
+    onClick={onAskSourcebot}
+    className="h-8 px-3 hover:bg-black/50"
+  >
+    Ask SourceBot
+  </Button>
+</div>
     )
-}
+}

packages/web/src/ee/features/chat/agent.test.ts

@@ -121,7 +121,7 @@ vi.mock('ai', async (importOriginal) => {
     };
 });
 
-const { createMessageStream } = await import('./agent');
+const { createMessageStream, sliceFileSourceForPrompt } = await import('./agent');
 const { getPromptCacheStrategy } = await import('./promptCaching');
 
 // Strategies reused across the prompt-caching tests below.
@@ -135,6 +135,33 @@ const listReposInput = {
     direction: 'asc',
 } as const;
 
+describe('sliceFileSourceForPrompt', () => {
+    test('slices an inclusive range and preserves its original line offset', () => {
+        expect(sliceFileSourceForPrompt('one\ntwo\nthree\nfour', {
+            startLine: 2,
+            endLine: 3,
+        })).toEqual({
+            source: 'two\nthree',
+            lineOffset: 2,
+        });
+    });
+
+    test('keeps the full file when there is no selected range', () => {
+        expect(sliceFileSourceForPrompt('one\ntwo', undefined)).toEqual({
+            source: 'one\ntwo',
+            lineOffset: 1,
+        });
+    });
+
+    test.each([
+        { startLine: 0, endLine: 1 },
+        { startLine: 3, endLine: 2 },
+        { startLine: 1, endLine: 4 },
+    ])('ignores invalid ranges safely', (range) => {
+        expect(sliceFileSourceForPrompt('one\ntwo\nthree', range)).toBeUndefined();
+    });
+});
+
 const dynamicApprovalRespondedPart = {
     type: 'dynamic-tool',
     toolName: 'mcp_linear__save_issue',

packages/web/src/ee/features/chat/agent.ts

@@ -21,7 +21,7 @@ import {
 import { randomUUID } from "crypto";
 import _dedent from "dedent";
 import { ANSWER_TAG, FILE_REFERENCE_PREFIX } from "@/features/chat/constants";
-import { Source } from "@/features/chat/types";
+import { FileSource, Source } from "@/features/chat/types";
 import { addLineNumbers, fileReferenceToString, formatAttachmentsForPrompt, getAnswerPartFromAssistantMessage, getTurnProgressState, getUserMessageAttachments, getUserMessageText } from "@/features/chat/utils";
 import { createTools } from "./tools";
 import { getConnectedMcpClients } from "@/ee/features/chat/mcp/mcpClientFactory";
@@ -551,6 +551,29 @@ interface AgentOptions {
     orgId?: number;
 }
 
+export const sliceFileSourceForPrompt = (
+    source: string,
+    range: FileSource['range'],
+): { source: string; lineOffset: number } | undefined => {
+    if (!range) {
+        return { source, lineOffset: 1 };
+    }
+
+    const lines = source.split('\n');
+    if (
+        range.startLine < 1 ||
+        range.endLine < range.startLine ||
+        range.endLine > lines.length
+    ) {
+        return undefined;
+    }
+
+    return {
+        source: lines.slice(range.startLine - 1, range.endLine).join('\n'),
+        lineOffset: range.startLine,
+    };
+};
+
 const createAgentStream = async ({
     model,
     promptCacheStrategy,
@@ -591,12 +614,20 @@ const createAgentStream = async ({
                 return undefined;
             }
 
+            const selectedSource = sliceFileSourceForPrompt(fileSource.source, source.range);
+            if (!selectedSource) {
+                logger.warn(`Ignoring invalid selected range for ${source.repo}:${source.path}`);
+                return undefined;
+            }
+
             return {
                 path: fileSource.path,
-                source: fileSource.source,
+                source: selectedSource.source,
                 repo: fileSource.repo,
                 language: fileSource.language,
                 revision: source.revision,
+                lineOffset: selectedSource.lineOffset,
+                range: source.range,
             };
         }))
     ).filter((source) => source !== undefined);
@@ -898,6 +929,8 @@ const createPrompt = ({
         repo: string;
         language: string;
         revision: string;
+        range?: FileSource['range'];
+        lineOffset: number;
     }[],
     repos: string[],
     mcpToolRegistry: McpToolRegistryEntry[],
@@ -999,8 +1032,8 @@ const createPrompt = ({
         <files>
         The user has mentioned the following files, which are automatically included for analysis.
 
-        ${files.map(file => `<file path="${file.path}" repository="${file.repo}" language="${file.language}" revision="${file.revision}">
-            ${addLineNumbers(file.source)}
+        ${files.map(file => `<file path="${file.path}" repository="${file.repo}" language="${file.language}" revision="${file.revision}"${file.range ? ` selected_lines="${file.range.startLine}-${file.range.endLine}"` : ''}>
+            ${addLineNumbers(file.source, file.lineOffset)}
             </file>`).join('\n\n')}
         </files>
         `);

packages/web/src/features/chat/useCreateNewChatThread.test.tsx

@@ -0,0 +1,114 @@
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, expect, test, vi } from "vitest";
+import type { Source } from "./types";
+import { DISABLED_MCP_SERVER_IDS_LOCAL_STORAGE_KEY } from "./constants";
+
+const mocks = vi.hoisted(() => ({
+    createChat: vi.fn(),
+    createUIMessage: vi.fn(),
+    push: vi.fn(),
+    setChatState: vi.fn(),
+    toast: vi.fn(),
+}));
+
+vi.mock("next/navigation", () => ({
+    useRouter: () => ({ push: mocks.push }),
+}));
+
+vi.mock("@/components/hooks/use-toast", () => ({
+    useToast: () => ({ toast: mocks.toast }),
+}));
+
+vi.mock("./actions", () => ({
+    createChat: mocks.createChat,
+}));
+
+vi.mock("@/lib/utils", () => ({
+    isServiceError: () => false,
+    createPathWithQueryParams: (path: string) => path,
+}));
+
+vi.mock("./utils", () => ({
+    createUIMessage: mocks.createUIMessage,
+    getAllMentionElements: () => [],
+    slateContentToString: () => "",
+}));
+
+vi.mock("usehooks-ts", () => ({
+    useSessionStorage: () => [null, mocks.setChatState],
+}));
+
+const { useCreateNewChatThread } = await import("./useCreateNewChatThread");
+
+afterEach(() => {
+    window.localStorage.clear();
+    vi.clearAllMocks();
+});
+
+test("createChatFromSource preserves disabled MCP servers from local storage", async () => {
+    const disabledMcpServerIds = ["linear", "github"];
+    window.localStorage.setItem(
+        DISABLED_MCP_SERVER_IDS_LOCAL_STORAGE_KEY,
+        JSON.stringify(disabledMcpServerIds),
+    );
+    mocks.createChat.mockResolvedValue({ id: "chat-1" });
+    mocks.createUIMessage.mockReturnValue({ id: "initial-message" });
+    const { result } = renderHook(() => useCreateNewChatThread());
+    const source: Source = {
+        type: "file",
+        repo: "github.com/sourcebot-dev/sourcebot",
+        path: "packages/web/src/auth.ts",
+        name: "auth.ts",
+        revision: "main",
+    };
+
+    await act(async () => {
+        await result.current.createChatFromSource(source);
+    });
+
+    expect(mocks.createUIMessage).toHaveBeenCalledWith(
+        "Explain this selected code.",
+        [],
+        [],
+        disabledMcpServerIds,
+        [],
+        [source],
+    );
+    expect(mocks.setChatState).toHaveBeenCalledWith({
+        inputMessage: { id: "initial-message" },
+        selectedSearchScopes: [],
+        disabledMcpServerIds,
+    });
+});
+
+test("createChatFromSource ignores duplicate calls while chat creation is pending", async () => {
+    let resolveCreateChat: ((value: { id: string }) => void) | undefined;
+    mocks.createChat.mockImplementation(() => new Promise((resolve) => {
+        resolveCreateChat = resolve;
+    }));
+    mocks.createUIMessage.mockReturnValue({ id: "initial-message" });
+    const { result } = renderHook(() => useCreateNewChatThread());
+    const source: Source = {
+        type: "file",
+        repo: "github.com/sourcebot-dev/sourcebot",
+        path: "packages/web/src/auth.ts",
+        name: "auth.ts",
+        revision: "main",
+    };
+
+    let firstCreate: Promise<void> | undefined;
+    let secondCreate: Promise<void> | undefined;
+    act(() => {
+        firstCreate = result.current.createChatFromSource(source);
+        secondCreate = result.current.createChatFromSource(source);
+    });
+
+    expect(mocks.createChat).toHaveBeenCalledTimes(1);
+
+    resolveCreateChat?.({ id: "chat-1" });
+    await act(async () => {
+        await Promise.all([firstCreate, secondCreate]);
+    });
+
+    expect(mocks.push).toHaveBeenCalledTimes(1);
+});