diff --git a/packages/web/src/app/(app)/components/editorContextMenu.tsx b/packages/web/src/app/(app)/components/editorContextMenu.tsx index e64b1d13d..dd136b577 100644 --- a/packages/web/src/app/(app)/components/editorContextMenu.tsx +++ b/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(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 ( -
- -
+
+ + + +
) -} \ No newline at end of file +} diff --git a/packages/web/src/ee/features/chat/agent.test.ts b/packages/web/src/ee/features/chat/agent.test.ts index 1432cbe9f..af45b70a0 100644 --- a/packages/web/src/ee/features/chat/agent.test.ts +++ b/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', diff --git a/packages/web/src/ee/features/chat/agent.ts b/packages/web/src/ee/features/chat/agent.ts index eccdf389e..b336005e4 100644 --- a/packages/web/src/ee/features/chat/agent.ts +++ b/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 = ({ The user has mentioned the following files, which are automatically included for analysis. - ${files.map(file => ` - ${addLineNumbers(file.source)} + ${files.map(file => ` + ${addLineNumbers(file.source, file.lineOffset)} `).join('\n\n')} `); diff --git a/packages/web/src/features/chat/useCreateNewChatThread.test.tsx b/packages/web/src/features/chat/useCreateNewChatThread.test.tsx new file mode 100644 index 000000000..880155ec6 --- /dev/null +++ b/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 | undefined; + let secondCreate: Promise | 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); +}); diff --git a/packages/web/src/features/chat/useCreateNewChatThread.ts b/packages/web/src/features/chat/useCreateNewChatThread.ts index ef5018ccf..baa768a3b 100644 --- a/packages/web/src/features/chat/useCreateNewChatThread.ts +++ b/packages/web/src/features/chat/useCreateNewChatThread.ts @@ -1,6 +1,6 @@ 'use client'; -import { useCallback, useState } from "react"; +import { useCallback, useRef, useState } from "react"; import { Descendant } from "slate"; import { createUIMessage, getAllMentionElements } from "./utils"; import { slateContentToString } from "./utils"; @@ -9,17 +9,34 @@ import { useRouter } from "next/navigation"; import { createChat } from "./actions"; import { isServiceError } from "@/lib/utils"; import { createPathWithQueryParams } from "@/lib/utils"; -import { AttachmentData, SearchScope, SetChatStatePayload } from "./types"; +import { AttachmentData, SearchScope, SetChatStatePayload, Source } from "./types"; import { DISABLED_MCP_SERVER_IDS_LOCAL_STORAGE_KEY, SELECTED_SEARCH_SCOPES_LOCAL_STORAGE_KEY, SET_CHAT_STATE_SESSION_STORAGE_KEY } from "./constants"; import { useSessionStorage } from "usehooks-ts"; +const getStoredDisabledMcpServerIds = (): string[] => { + try { + const stored = window.localStorage.getItem(DISABLED_MCP_SERVER_IDS_LOCAL_STORAGE_KEY); + if (stored) { + return JSON.parse(stored) as string[]; + } + } catch { /* fall through to [] */ } + + return []; +} + export const useCreateNewChatThread = () => { const [isLoading, setIsLoading] = useState(false); + const createInFlightRef = useRef(false); const { toast } = useToast(); const router = useRouter(); const [, setChatState] = useSessionStorage(SET_CHAT_STATE_SESSION_STORAGE_KEY, null); const createNewChatThread = useCallback(async (children: Descendant[], overrideSearchScopes?: SearchScope[], overrideDisabledMcpServerIds?: string[], attachments: AttachmentData[] = []) => { + if (createInFlightRef.current) { + return; + } + createInFlightRef.current = true; + const text = slateContentToString(children); const mentions = getAllMentionElements(children); @@ -31,16 +48,8 @@ export const useCreateNewChatThread = () => { } } catch { /* fall through to [] */ } - let storedDisabledMcpServerIds: string[] = []; - try { - const stored = window.localStorage.getItem(DISABLED_MCP_SERVER_IDS_LOCAL_STORAGE_KEY); - if (stored) { - storedDisabledMcpServerIds = JSON.parse(stored) as string[]; - } - } catch { /* fall through to [] */ } - const selectedSearchScopes = overrideSearchScopes ?? storedScopes; - const disabledMcpServerIds = overrideDisabledMcpServerIds ?? storedDisabledMcpServerIds; + const disabledMcpServerIds = overrideDisabledMcpServerIds ?? getStoredDisabledMcpServerIds(); const inputMessage = createUIMessage(text, mentions.map((mention) => mention.data), selectedSearchScopes, disabledMcpServerIds, attachments); setIsLoading(true); @@ -50,6 +59,7 @@ export const useCreateNewChatThread = () => { description: `❌ Failed to create chat. Reason: ${response.message}` }); setIsLoading(false); + createInFlightRef.current = false; return; } @@ -64,8 +74,45 @@ export const useCreateNewChatThread = () => { router.push(url); }, [router, toast, setChatState]); + const createChatFromSource = useCallback(async (source: Source) => { + if (createInFlightRef.current) { + return; + } + createInFlightRef.current = true; + + const disabledMcpServerIds = getStoredDisabledMcpServerIds(); + const inputMessage = createUIMessage( + 'Explain this selected code.', + [], + [], + disabledMcpServerIds, + [], + [source], + ); + + setIsLoading(true); + const response = await createChat({ source: 'sourcebot-web-client' }); + if (isServiceError(response)) { + toast({ + description: `❌ Failed to create chat. Reason: ${response.message}`, + }); + setIsLoading(false); + createInFlightRef.current = false; + return; + } + + setChatState({ + inputMessage, + selectedSearchScopes: [], + disabledMcpServerIds, + }); + + router.push(`/chat/${response.id}`); + }, [router, setChatState, toast]); + return { createNewChatThread, + createChatFromSource, isLoading, }; } diff --git a/packages/web/src/features/chat/utils.test.ts b/packages/web/src/features/chat/utils.test.ts index bf598a25e..d1ad754ad 100644 --- a/packages/web/src/features/chat/utils.test.ts +++ b/packages/web/src/features/chat/utils.test.ts @@ -769,6 +769,29 @@ test('repairReferences handles malformed inline code blocks', () => { }); describe('createUIMessage', () => { + test('includes an explicit ranged file source', () => { + const result = createUIMessage('Explain this selected code.', [], [], [], [], [{ + type: 'file', + repo: 'github.com/sourcebot-dev/sourcebot', + path: 'packages/web/src/auth.ts', + name: 'auth.ts', + revision: 'main', + range: { startLine: 12, endLine: 30 }, + }]); + + expect(result.parts).toContainEqual({ + type: 'data-source', + data: { + type: 'file', + repo: 'github.com/sourcebot-dev/sourcebot', + path: 'packages/web/src/auth.ts', + name: 'auth.ts', + revision: 'main', + range: { startLine: 12, endLine: 30 }, + }, + }); + }); + test('includes disabledMcpServerIds in metadata when provided', () => { const result = createUIMessage('hello', [], [], ['server1', 'server2']); diff --git a/packages/web/src/features/chat/utils.ts b/packages/web/src/features/chat/utils.ts index 23cd08590..908ee7c55 100644 --- a/packages/web/src/features/chat/utils.ts +++ b/packages/web/src/features/chat/utils.ts @@ -200,7 +200,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[], selectedSearchScopes: SearchScope[], disabledMcpServerIds: string[] = [], attachments: AttachmentData[] = []): CreateUIMessage => { +export const createUIMessage = (text: string, mentions: MentionData[], selectedSearchScopes: SearchScope[], disabledMcpServerIds: string[] = [], attachments: AttachmentData[] = [], explicitSources: Source[] = []): CreateUIMessage => { // Converts applicable mentions into sources. const sources: Source[] = mentions .map((mention) => { @@ -218,6 +218,7 @@ export const createUIMessage = (text: string, mentions: MentionData[], selectedS return undefined; }) .filter((source) => source !== undefined); + sources.push(...explicitSources); const commandInvocation = createCommandInvocationData( text, mentions.filter((mention) => mention.type === 'command'), diff --git a/packages/web/src/features/tools/types.ts b/packages/web/src/features/tools/types.ts index 9c221a409..d71356120 100644 --- a/packages/web/src/features/tools/types.ts +++ b/packages/web/src/features/tools/types.ts @@ -6,6 +6,10 @@ const fileSourceSchema = z.object({ path: z.string(), name: z.string(), revision: z.string(), + range: z.object({ + startLine: z.number().int().positive(), + endLine: z.number().int().positive(), + }).refine(({ startLine, endLine }) => endLine >= startLine).optional(), }); export type FileSource = z.infer;