Chat: allow large file with range in @-mentions

vscode/CHANGELOG.md

@@ -18,7 +18,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 ### Fixed
 
-- Chat: Large file cannot be added via @-mention. [pull/3531](https://github.com/sourcegraph/cody/pull/3531)
+- Chat: Fixed issue where large files could not be added via @-mention. You can now @-mention line ranges within large files. [pull/3531](https://github.com/sourcegraph/cody/pull/3531) & [pull/3585](https://github.com/sourcegraph/cody/pull/3585)
 - Edit: Improved the response reliability, Edit commands should no longer occasionally produce Markdown outputs.[pull/3192](https://github.com/sourcegraph/cody/pull/3192)
 - Chat: Handle empty chat message input and prevent submission of empty messages. [pull/3554](https://github.com/sourcegraph/cody/pull/3554)
 - Chat: Warnings are now displayed correctly for large files in the @-mention file selection list. [pull/3526](https://github.com/sourcegraph/cody/pull/3526)

vscode/src/chat/context/constants.ts

@@ -4,7 +4,9 @@ export const SYMBOL_HELP_LABEL = 'Search for a symbol to include...'
 export const NO_SYMBOL_MATCHES_LABEL = 'No symbols found'
 export const NO_FILE_MATCHES_LABEL = 'No files found'
 export const NO_SYMBOL_MATCHES_HELP_LABEL = ' (language extensions may be loading)'
-export const FILE_TOO_LARGE_LABEL = 'File too large. Type @# to choose a symbol.'
+export const FILE_RANGE_TOOLTIP_LABEL = 'Type a line range to include, e.g. 5-10...'
+export const LARGE_FILE_WARNING_LABEL =
+    'File too large. Add line range with : or use @# to choose a symbol'
 
 export const QUICK_PICK_ITEM_EMPTY_INDENT_PREFIX = '\u00A0\u00A0\u00A0\u00A0\u00A0'
 export const QUICK_PICK_ITEM_CHECKED_PREFIX = '$(check)'

vscode/src/edit/input/get-input.ts

@@ -10,8 +10,8 @@ import * as vscode from 'vscode'
 
 import {
     FILE_HELP_LABEL,
-    FILE_TOO_LARGE_LABEL,
     GENERAL_HELP_LABEL,
+    LARGE_FILE_WARNING_LABEL,
     NO_FILE_MATCHES_LABEL,
     NO_SYMBOL_MATCHES_LABEL,
     SYMBOL_HELP_LABEL,
@@ -413,7 +413,9 @@ export const getInput = async (
                         label: shortLabel || key,
                         description: shortLabel ? key : undefined,
                         detail:
-                            'isTooLarge' in item && item.isTooLarge ? FILE_TOO_LARGE_LABEL : undefined,
+                            item.type === 'file' && item.isTooLarge
+                                ? LARGE_FILE_WARNING_LABEL
+                                : undefined,
                     })),
                     {
                         kind: vscode.QuickPickItemKind.Separator,
@@ -450,7 +452,7 @@ export const getInput = async (
                         unitTestInput.render(activeTitle, '')
                         return
                     case FILE_HELP_LABEL:
-                    case FILE_TOO_LARGE_LABEL:
+                    case LARGE_FILE_WARNING_LABEL:
                     case SYMBOL_HELP_LABEL:
                     case NO_FILE_MATCHES_LABEL:
                     case NO_SYMBOL_MATCHES_LABEL:

vscode/src/prompt-builder/index.ts

@@ -119,22 +119,27 @@ export class PromptBuilder {
             }
         }
         for (const v of contextItemsAndMessages) {
+            const item = contextItem(v)
             const uri = contextItem(v).uri
             if (uri.scheme === 'file' && isCodyIgnoredFile(uri)) {
-                ignored.push(contextItem(v))
+                ignored.push(item)
                 continue
             }
             const id = contextItemId(v)
             if (this.seenContext.has(id)) {
-                duplicate.push(contextItem(v))
+                duplicate.push(item)
                 continue
             }
             const contextMessage = isContextItem(v) ? renderContextItem(v) : v
             const contextLen = contextMessage
                 ? contextMessage.speaker.length + contextMessage.text.length + 3
                 : 0
             if (this.charsUsed + contextLen > effectiveCharLimit) {
-                ignored.push(contextItem(v))
+                // Marks excluded context items as too large to be displayed in UI
+                if (item.type === 'file') {
+                    item.isTooLarge = true
+                }
+                ignored.push(item)
                 limitReached = true
                 continue
             }
@@ -144,13 +149,16 @@ export class PromptBuilder {
                 this.reverseMessages.push(contextMessage)
             }
             this.charsUsed += contextLen
-            used.push(contextItem(v))
+            // Marks added file items as not too large
+            if (item.type === 'file') {
+                item.isTooLarge = false
+            }
+            used.push(item)
         }
         return {
             limitReached,
             used,
-            // Marks excluded context items as too large to be displayed in UI
-            ignored: ignored.map(c => ({ ...c, isTooLarge: true })),
+            ignored,
             duplicate,
         }
     }

vscode/test/e2e/chat-atFile.test.ts

@@ -87,8 +87,6 @@ test.extend<ExpectedEvents>({
     await chatInput.fill('Explain @mj')
     await chatPanelFrame.getByRole('option', { name: 'Main.java' }).click()
     await expect(chatInput).toHaveText('Explain @Main.java ')
-    await expect(chatInput.getByText('@Main.java')).not.toHaveClass(/context-item-mention-node/)
-    await chatInput.press('Enter')
     await expect(chatInput.getByText('@Main.java')).toHaveClass(/context-item-mention-node/)
     await chatInput.press('Enter')
     await expect(chatInput).toBeEmpty()
@@ -118,7 +116,6 @@ test.extend<ExpectedEvents>({
     await chatInput.press('ArrowDown') // second item again
     await expect(chatPanelFrame.getByRole('option', { selected: true })).toHaveText(/visualize\.go/)
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText(
         withPlatformSlashes(
             'Explain @lib/batches/env/var.go and @lib/codeintel/tools/lsif-visualize/visualize.go '
@@ -146,15 +143,13 @@ test.extend<ExpectedEvents>({
     await chatInput.pressSequentially('@Main.java', { delay: 10 })
     await expect(chatPanelFrame.getByRole('option', { name: 'Main.java' })).toBeVisible()
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText('@Main.java ')
 
     // Check pressing tab after typing a partial filename but where that complete
     // filename already exists earlier in the input.
     // https://github.com/sourcegraph/cody/issues/2243
     await chatInput.pressSequentially('and @Main.ja', { delay: 10 })
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText('@Main.java and @Main.java ')
 
     // Support @-file in mid-sentence
@@ -170,7 +165,6 @@ test.extend<ExpectedEvents>({
     await chatInput.pressSequentially('@Main', { delay: 10 })
     await expect(chatPanelFrame.getByRole('option', { name: 'Main.java' })).toBeVisible()
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText('Explain the @Main.java file')
     // Confirm the cursor is at the end of the newly added file name with space
     await page.keyboard.press('!')
@@ -215,7 +209,6 @@ test('editing a chat message with @-mention', async ({ page, sidebar }) => {
     // Send a message with an @-mention.
     await chatInput.fill('Explain @mj')
     await chatPanelFrame.getByRole('option', { name: 'Main.java' }).click()
-    await chatInput.press('Enter')
     await expect(chatInput).toHaveText('Explain @Main.java ')
     await expect(chatInput.getByText('@Main.java')).toHaveClass(/context-item-mention-node/)
     await chatInput.press('Enter')
@@ -235,7 +228,6 @@ test('editing a chat message with @-mention', async ({ page, sidebar }) => {
     await expect(chatInput).toHaveText('Explain @Main.java ')
     await chatInput.pressSequentially('and @index.ht')
     await chatPanelFrame.getByRole('option', { name: 'index.html' }).click()
-    await chatPanelFrame.getByRole('option', { name: 'index.html' }).click()
     await expect(chatInput).toHaveText('Explain @Main.java and @index.html')
     await expect(chatInput.getByText('@index.html')).toHaveClass(/context-item-mention-node/)
     await chatInput.press('Enter')
@@ -254,7 +246,6 @@ test('pressing Enter with @-mention menu open selects item, does not submit mess
     await chatInput.fill('Explain @index.htm')
     await expect(chatPanelFrame.getByRole('option', { name: 'index.html' })).toBeVisible()
     await chatInput.press('Enter')
-    await chatInput.press('Enter')
     await expect(chatInput).toHaveText('Explain @index.html')
     await expect(chatInput.getByText('@index.html')).toHaveClass(/context-item-mention-node/)
 })
@@ -269,7 +260,6 @@ test('@-mention links in transcript message', async ({ page, sidebar }) => {
     await chatInput.fill('Hello @buzz.ts')
     await chatPanelFrame.getByRole('option', { name: 'buzz.ts' }).click()
     await chatInput.press('Enter')
-    await chatInput.press('Enter')
 
     // In the transcript, the @-mention is linked, and clicking the link opens the file.
     const transcriptMessage = chatPanelFrame.getByText('Hello @buzz.ts')
@@ -292,9 +282,6 @@ test('@-mention file range', async ({ page, sidebar }) => {
     await expect(chatPanelFrame.getByRole('option', { name: 'buzz.ts Lines 2-4' })).toBeVisible()
     await chatPanelFrame.getByRole('option', { name: 'buzz.ts Lines 2-4' }).click()
     await expect(chatInput).toHaveText('@buzz.ts:2-4 ')
-    // Enter again to turn it into a token
-    await chatInput.press('Enter')
-
     // Submit the message
     await chatInput.press('Enter')
 

vscode/test/e2e/chat-edits.test.ts

@@ -105,7 +105,6 @@ test.extend<ExpectedEvents>({
     await expect(chatInput).not.toHaveText('Four')
     await expect(chatFrame.getByRole('option', { name: 'Main.java' })).toBeVisible()
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText('Explain @Main.java ')
 
     // Enter should submit the message and exit editing mode
@@ -122,7 +121,6 @@ test.extend<ExpectedEvents>({
     await expect(chatInput).toHaveText('Explain @Main.java ')
     await chatInput.type('and @vgo', { delay: 50 })
     await chatInput.press('Tab')
-    await chatInput.press('Tab')
     await expect(chatInput).toHaveText(
         withPlatformSlashes('Explain @Main.java and @lib/batches/env/var.go ')
     )

vscode/webviews/promptEditor/nodes/ContextItemMentionNode.ts

@@ -176,11 +176,7 @@ export function contextItemMentionNodeDisplayText(contextItem: SerializedContext
     // range needs to go to the start (0th character) of line 5. Also, `RangeData` is 0-indexed but
     // display ranges are 1-indexed.
     const rangeText = contextItem.range ? `:${displayLineRange(contextItem.range)}` : ''
-    if (contextItem.type === 'file' && !contextItem.isTooLarge) {
-        return `@${displayPath(URI.parse(contextItem.uri))}${rangeText}`
-    }
-    // Large file cannot be added to the context file list unless it has a range.
-    if (contextItem.type === 'file' && contextItem.isTooLarge && rangeText) {
+    if (contextItem.type === 'file') {
         return `@${displayPath(URI.parse(contextItem.uri))}${rangeText}`
     }
     if (contextItem.type === 'symbol') {

vscode/webviews/promptEditor/plugins/atMentions/OptionsList.module.css

@@ -72,7 +72,7 @@
     background-color: var(--vscode-list-hoverBackground);
 }
 .option-item.disabled {
-    opacity: 0.5;
+    opacity: 0.8;
     cursor: default;
 }
 

vscode/webviews/promptEditor/plugins/atMentions/OptionsList.tsx

@@ -11,8 +11,9 @@ import classNames from 'classnames'
 import { type FunctionComponent, useEffect, useRef } from 'react'
 import {
     FILE_HELP_LABEL,
-    FILE_TOO_LARGE_LABEL,
+    FILE_RANGE_TOOLTIP_LABEL,
     GENERAL_HELP_LABEL,
+    LARGE_FILE_WARNING_LABEL,
     NO_FILE_MATCHES_LABEL,
     NO_SYMBOL_MATCHES_HELP_LABEL,
     NO_SYMBOL_MATCHES_LABEL,
@@ -49,7 +50,9 @@ export const OptionsList: FunctionComponent<
                                 : NO_SYMBOL_MATCHES_LABEL +
                                   (mentionQuery.text.length < 3 ? NO_SYMBOL_MATCHES_HELP_LABEL : '')
                           : options.length > 0
-                              ? FILE_HELP_LABEL
+                              ? isValidLineRangeQuery(query)
+                                    ? FILE_RANGE_TOOLTIP_LABEL
+                                    : FILE_HELP_LABEL
                               : NO_FILE_MATCHES_LABEL}
                 </span>
                 <br />
@@ -87,16 +90,19 @@ const Item: FunctionComponent<{
     className?: string
 }> = ({ query, isSelected, onClick, onMouseEnter, option, className }) => {
     const item = option.item
-    const icon =
-        item.type === 'file' ? null : item.kind === 'class' ? 'symbol-structure' : 'symbol-method'
-    const title = item.title ?? (item.type === 'file' ? displayPathBasename(item.uri) : item.symbolName)
+    const isFileType = item.type === 'file'
+    const icon = isFileType ? null : item.kind === 'class' ? 'symbol-structure' : 'symbol-method'
+    const title = item.title ?? (isFileType ? displayPathBasename(item.uri) : item.symbolName)
+
     const range = getLineRangeInMention(query, item.range)
-    const dirname = displayPathDirname(item.uri)
-    const description =
-        item.type === 'file'
-            ? `${range ? `Lines ${range} · ` : ''}${dirname === '.' ? '' : dirname}`
-            : displayPath(item.uri) + `:${range}`
-    const warning = item.type === 'file' && item.isTooLarge ? FILE_TOO_LARGE_LABEL : undefined
+    const dir = displayPathDirname(item.uri)
+    const description = isFileType
+        ? `${range ? `Lines ${range} · ` : ''}${dir === '.' ? '' : dir}`
+        : `${displayPath(item.uri)}:${getLineRangeInMention(query, item.range)}`
+
+    const isLargeFile = isFileType && item.isTooLarge
+    const warning =
+        isLargeFile && !item.range && !isValidLineRangeQuery(query) ? LARGE_FILE_WARNING_LABEL : ''
 
     return (
         // biome-ignore lint/a11y/useKeyWithClickEvents:
@@ -135,6 +141,9 @@ const Item: FunctionComponent<{
     )
 }
 
+const isValidLineRangeQuery = (query: string): boolean =>
+    query.endsWith(':') || RANGE_MATCHES_REGEXP.test(query)
+
 /**
  * Gets the display line range from the query string.
  */

vscode/webviews/promptEditor/plugins/atMentions/atMentions.tsx

@@ -111,22 +111,24 @@ export default function MentionsPlugin(): JSX.Element | null {
                 if (!currentInputText) {
                     return
                 }
-                // On first selection, add the selected option as text.
-                // This allows users to autocomplete the file path, and provide them with
-                // the options to make additional changes, e.g. add range, before inserting the mention.
-                const textNode = $createContextItemTextNode(selectedOption.item)
-                if (!currentInputText.endsWith(textNode.__text) && !currentInputText.startsWith('@#')) {
+
+                const selectedItem = selectedOption.item
+                const isLargeFile = selectedItem.type === 'file' && selectedItem.isTooLarge
+                // When selecting a large file without range, add the selected option as text node with : at the end.
+                // This allows users to autocomplete the file path, and provide them with the options to add range.
+                if (isLargeFile && !selectedItem.range) {
+                    const textNode = $createContextItemTextNode(selectedItem)
                     nodeToReplace.replace(textNode)
-                    textNode.select()
-                    closeMenu()
-                    return
+                    const colonNode = $createTextNode(':')
+                    textNode.insertAfter(colonNode)
+                    colonNode.select()
+                } else {
+                    const mentionNode = $createContextItemMentionNode(selectedItem)
+                    nodeToReplace.replace(mentionNode)
+                    const spaceNode = $createTextNode(' ')
+                    mentionNode.insertAfter(spaceNode)
+                    spaceNode.select()
                 }
-
-                const mentionNode = $createContextItemMentionNode(selectedOption.item)
-                nodeToReplace?.replace(mentionNode)
-                const spaceAfter = $createTextNode(' ')
-                mentionNode.insertAfter(spaceAfter)
-                spaceAfter.select()
                 closeMenu()
             })
         },

vscode/webviews/promptEditor/plugins/atMentions/chatContextClient.tsx

@@ -1,4 +1,4 @@
-import type { ContextItem } from '@sourcegraph/cody-shared'
+import { type ContextItem, displayPath } from '@sourcegraph/cody-shared'
 import { type FunctionComponent, createContext, useContext, useEffect, useState } from 'react'
 import { getVSCodeAPI } from '../../../utils/VSCodeApi'
 import { LINE_RANGE_REGEXP, RANGE_MATCHES_REGEXP, parseLineRangeInMention } from './atMentions'
@@ -66,11 +66,20 @@ export function useChatContextItems(query: string | null): ContextItem[] | undef
             return
         }
 
-        // If user is typing a line range, no need to fetch new chat context items.
-        if (results?.length && RANGE_MATCHES_REGEXP.test(query)) {
-            if (!LINE_RANGE_REGEXP.test(query)) {
-                return
-            }
+        // If the query ends with a colon, we will reuse current results but remove the range.
+        if (query.endsWith(':')) {
+            const selected = results?.find(r => displayPath(r.uri) === query.slice(0, -1))
+            setResults(
+                selected
+                    ? [{ ...selected, range: undefined }]
+                    : results?.map(r => ({ ...r, range: undefined }))
+            )
+            return
+        }
+
+        // If user is typing a line range, fetch new chat context items only if there are no results
+        if (results?.length && RANGE_MATCHES_REGEXP.test(query) && !LINE_RANGE_REGEXP.test(query)) {
+            return
         }
 
         // Track if the query changed since this request was sent (which would make our results