Chat: display file size warning on large files (#3118)

abeatrix · web-flow · commit e59b9c9894dc · 2024-02-13T07:17:20.000-08:00
CONTEXT: https://sourcegraph.slack.com/archives/C05AGQYD528/p1707508300392389 Part of https://github.com/sourcegraph/cody/issues/2636 & https://github.com/sourcegraph/cody/issues/2965 Update chat UI to provide better file size awareness based on Tim's figma design: https://www.figma.com/file/jEsnzgsf0hNuJbqB9pKaGm/VS-Code---Cody---Commands-up-front-and-ChatGPT%2B%2B?type=design&node-id=331-103646&mode=design&t=jidOC06CYg9nSWOI-4 ![image](https://github.com/sourcegraph/cody/assets/68532117/5308fc4a-29fa-4e31-8759-14366df6ffc5) This PR includes the following changes: - Filter the `sorted results` in the `getFileContextFiles` function to remove non-files and files that are larger than 1MB - Set title as `large-file` for files that contain more tokens than we allow for file using the file size returned from `vscode.workspace.fs.stat` API - When users try to add a file using @-mentioned, they should see a warning below the file name if the file is too large - ~~Truncating prompt messages in the PromptBuilder class to include truncated context without all context being ignored due to excessive length.~~ - ~~Use the exisiting truncatePrompt function and MAX_AVAILABLE_PROMPT_LENGTH constant for truncating prompts.~~ - Display warning in at-mentions selector for large files ### Note - This PR focuses on providing information to users in the UI for files they are trying to add using `@-mention` that are too large. - ~~We are also truncating the context files to make sure large files will be truncated and used as context, instead of excluding the file.~~ Follow-up works include: - truncate large files - show users if a file was truncated ## Test plan  1. @-mentioned a large file [SimpleChatPanelProvider](https://sourcegraph.sourcegraph.com/github.com/sourcegraph/cody@916ef190bf286505c5dcb254b87c0ad142b86c5a/-/blob/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts) ![image](https://github.com/sourcegraph/cody/assets/68532117/5308fc4a-29fa-4e31-8759-14366df6ffc5) raph/cody/assets/68532117/7ea38043-c148-45d7-9a00-1d6aeb53f552)
diff --git a/vscode/src/editor/utils/editor-context.test.ts b/vscode/src/editor/utils/editor-context.test.ts
@@ -12,10 +12,37 @@ afterEach(() => {
 })
 
 describe('getFileContextFiles', () => {
+    /**
+     * Mocks the fs.stat function to return a fake stat object for the given URI.
+     * This allows tests to mock filesystem access for specific files.
+     */
+    function setFileStat(uri: vscode.Uri, isFile = true) {
+        vscode.workspace.fs.stat = vi.fn().mockImplementation(() => {
+            const relativePath = uriBasename(uri)
+            return {
+                type: isFile ? vscode.FileType.File : vscode.FileType.SymbolicLink,
+                ctime: 1,
+                mtime: 1,
+                size: 1,
+                isDirectory: () => false,
+                isFile: () => isFile,
+                isSymbolicLink: () => !isFile,
+                uri,
+                with: vi.fn(),
+                toString: vi.fn().mockReturnValue(relativePath),
+            }
+        })
+    }
+
     function setFiles(relativePaths: string[]) {
         vscode.workspace.findFiles = vi
             .fn()
             .mockResolvedValueOnce(relativePaths.map(f => testFileUri(f)))
+
+        for (const relativePath of relativePaths) {
+            const isFile = relativePath !== 'symlink'
+            setFileStat(testFileUri(relativePath), isFile)
+        }
     }
 
     async function runSearch(query: string, maxResults: number): Promise<(string | undefined)[]> {
@@ -67,6 +94,16 @@ describe('getFileContextFiles', () => {
         expect(vscode.workspace.findFiles).toBeCalledTimes(1)
     })
 
+    it('do not return non-file (e.g. symlinks) result', async () => {
+        setFiles(['symlink'])
+
+        expect(await runSearch('symlink', 5)).toMatchInlineSnapshot(`
+          []
+        `)
+
+        expect(vscode.workspace.findFiles).toBeCalledTimes(1)
+    })
+
     it('filters out ignored files', async () => {
         ignores.setActiveState(true)
         ignores.setIgnoreFiles(testFileUri(''), [
diff --git a/vscode/src/editor/utils/editor-context.ts b/vscode/src/editor/utils/editor-context.ts
@@ -14,9 +14,11 @@ import {
     type ContextFileSymbol,
     type ContextFileType,
     type SymbolKind,
+    MAX_CURRENT_FILE_TOKENS,
 } from '@sourcegraph/cody-shared'
 
 import { getOpenTabsUris, getWorkspaceSymbols } from '.'
+import { CHARS_PER_TOKEN } from '@sourcegraph/cody-shared/src/prompt/constants'
 
 const findWorkspaceFiles = async (
     cancellationToken: vscode.CancellationToken
@@ -85,7 +87,7 @@ export async function getFileContextFiles(
         threshold: -100000,
     })
 
-    // Remove ignored files and apply a penalty for segments that are in the low scoring list.
+    // Apply a penalty for segments that are in the low scoring list.
     const adjustedResults = [...results].map(result => {
         const segments = result.obj.uri.fsPath.split(path.sep)
         for (const lowScoringPathSegment of lowScoringPathSegments) {
@@ -98,20 +100,42 @@ export async function getFileContextFiles(
         }
         return result
     })
-
     // fuzzysort can return results in different order for the same query if
     // they have the same score :( So we do this hacky post-limit sorting (first
     // by score, then by path) to ensure the order stays the same.
-    const sortedResults = adjustedResults.sort((a, b) => {
-        return (
-            b.score - a.score ||
-            new Intl.Collator(undefined, { numeric: true }).compare(a.obj.uri.fsPath, b.obj.uri.fsPath)
-        )
-    })
+    const sortedResults = adjustedResults
+        .sort((a, b) => {
+            return (
+                b.score - a.score ||
+                new Intl.Collator(undefined, { numeric: true }).compare(a.obj.uri.path, b.obj.uri.path)
+            )
+        })
+        .flatMap(result => createContextFileFromUri(result.obj.uri, 'user', 'file'))
 
     // TODO(toolmantim): Add fuzzysort.highlight data to the result so we can show it in the UI
 
-    return sortedResults.flatMap(result => createContextFileFromUri(result.obj.uri, 'user', 'file'))
+    const filtered = []
+    try {
+        for (const sorted of sortedResults) {
+            // Remove file larger than 1MB and non-text files
+            // NOTE: Sourcegraph search only includes files up to 1MB
+            const fileStat = await vscode.workspace.fs.stat(sorted.uri)
+            if (fileStat.type !== vscode.FileType.File || fileStat.size > 1000000) {
+                continue
+            }
+            // Check if file contains more characters than the token limit based on fileStat.size
+            // and set the title of the result as 'large-file' for webview to display file size
+            // warning.
+            if (fileStat.size > CHARS_PER_TOKEN * MAX_CURRENT_FILE_TOKENS) {
+                sorted.title = 'large-file'
+            }
+            filtered.push(sorted)
+        }
+    } catch (error) {
+        console.log('atMention:getFileContextFiles:failed', error)
+    }
+
+    return filtered
 }
 
 export async function getSymbolContextFiles(
diff --git a/vscode/test/e2e/chat-atFile.test.ts b/vscode/test/e2e/chat-atFile.test.ts
@@ -152,15 +152,17 @@ test('@-file & @#-symbol in chat view', async ({ page, sidebar }) => {
 
     //  "ArrowLeft" / "ArrowRight" keys close the selection without altering current input.
     const noMatches = chatPanelFrame.getByRole('heading', { name: 'No matching files found' })
-    await page.keyboard.type(' @abcdefg')
+    await chatInput.type(' @abcdefg', { delay: 50 })
     await expect(chatInput).toHaveValue('Explain the @Main.java ! @abcdefgfile')
+    await noMatches.hover()
     await expect(noMatches).toBeVisible()
     await chatInput.press('ArrowLeft')
     await expect(noMatches).not.toBeVisible()
     await chatInput.press('ArrowRight')
     await expect(noMatches).not.toBeVisible()
-    await chatInput.press('?')
+    await chatInput.type('?', { delay: 50 })
     await expect(chatInput).toHaveValue('Explain the @Main.java ! @abcdefg?file')
+    await noMatches.hover()
     await expect(noMatches).toBeVisible()
     // Selection close on submit
     await chatInput.press('Enter')
diff --git a/vscode/webviews/UserContextSelector.module.css b/vscode/webviews/UserContextSelector.module.css
@@ -43,7 +43,8 @@
 
 .selection-item {
     display: flex;
-    align-items: center;
+    flex-direction: column;
+    align-items: self-start;
     font-size: inherit;
     width: 100%;
     cursor: pointer;
@@ -63,6 +64,10 @@
     overflow: hidden;
 }
 
+.warning-container {
+    margin: 0;
+}
+
 .selection-title {
     white-space: nowrap;
     text-overflow: ellipsis;
@@ -71,8 +76,7 @@
 
 .selection-description {
     margin-left: 0.25rem;
-    font-size: smaller;
-    opacity: 0.7;
+    opacity: 0.8;
     white-space: nowrap;
     text-overflow: ellipsis;
     overflow: hidden;
@@ -111,4 +115,4 @@ body[data-vscode-theme-kind='vscode-high-contrast'] .selected {
 
 .empty-symbol-search-tip i {
     vertical-align: middle;
-}
+}
diff --git a/vscode/webviews/UserContextSelector.tsx b/vscode/webviews/UserContextSelector.tsx
@@ -97,6 +97,10 @@ export const UserContextSelectorComponent: React.FunctionComponent<
                             : ''
                         const description =
                             match.type === 'file' ? undefined : displayPath(match.uri) + range
+                        const warning =
+                            match.type === 'file' && match.title === 'large-file'
+                                ? 'File too large. Type @# to choose a symbol'
+                                : undefined
                         return (
                             <React.Fragment key={`${icon}${title}${range}${description}`}>
                                 <button
@@ -105,6 +109,7 @@ export const UserContextSelectorComponent: React.FunctionComponent<
                                         styles.selectionItem,
                                         selected === i && styles.selected
                                     )}
+                                    title={title}
                                     onClick={() => onSelected(match, formInput)}
                                     type="button"
                                 >
@@ -124,6 +129,18 @@ export const UserContextSelectorComponent: React.FunctionComponent<
                                             </span>
                                         )}
                                     </span>
+                                    {warning && (
+                                        <p
+                                            className={classNames(
+                                                styles.titleAndDescriptionContainer,
+                                                styles.warningContainer
+                                            )}
+                                        >
+                                            <span className={styles.selectionDescription}>
+                                                {warning}
+                                            </span>
+                                        </p>
+                                    )}
                                 </button>
                             </React.Fragment>
                         )
