Chat: fix at-mention token size (#3526)

vscode/CHANGELOG.md

@@ -16,6 +16,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 - Chat: Large file cannot be added via @-mention. [pull/3531](https://github.com/sourcegraph/cody/pull/3531)
 - 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)
 
 ### Changed
 

vscode/src/chat/chat-view/SimpleChatModel.ts

@@ -4,6 +4,7 @@ import {
     type ChatMessage,
     type ContextItem,
     type Message,
+    ModelProvider,
     type SerializedChatInteraction,
     type SerializedChatTranscript,
     errorToChatError,
@@ -15,13 +16,17 @@ import type { Repo } from '../../context/repo-fetcher'
 import { getChatPanelTitle } from './chat-helpers'
 
 export class SimpleChatModel {
+    // The maximum number of characters available in the model's context window.
+    public readonly maxChars: number
     constructor(
         public modelID: string,
         private messages: ChatMessage[] = [],
         public readonly sessionID: string = new Date(Date.now()).toUTCString(),
         private customChatTitle?: string,
         private selectedRepos?: Repo[]
-    ) {}
+    ) {
+        this.maxChars = ModelProvider.getMaxCharsByModel(this.modelID)
+    }
 
     public isEmpty(): boolean {
         return this.messages.length === 0

vscode/src/chat/chat-view/SimpleChatPanelProvider.ts

@@ -46,12 +46,13 @@ import {
 } from '../../services/utils/codeblock-action-tracker'
 import { openExternalLinks, openLocalFileWithRange } from '../../services/utils/workspace-action'
 import { TestSupport } from '../../test-support'
-import { countGeneratedCode } from '../utils'
+import { countGeneratedCode, getContextWindowLimitInBytes } from '../utils'
 
 import type { Span } from '@opentelemetry/api'
 import { captureException } from '@sentry/core'
 import type { ContextItemWithContent } from '@sourcegraph/cody-shared/src/codebase-context/messages'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
+import { ANSWER_TOKENS, tokensToChars } from '@sourcegraph/cody-shared/src/prompt/constants'
 import { recordErrorToSpan, tracer } from '@sourcegraph/cody-shared/src/tracing'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { Repo } from '../../context/repo-fetcher'
@@ -596,12 +597,20 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
                 })
             },
         }
+        // Use the number of characters left in the chat model as the limit
+        // for adding user context files to the chat.
+        const contextLimit = getContextWindowLimitInBytes(
+            [...this.chatModel.getMessages()],
+            // Minus the character limit reserved for the answer token
+            this.chatModel.maxChars - tokensToChars(ANSWER_TOKENS)
+        )
 
         try {
             const items = await getChatContextItemsForMention(
                 query,
                 cancellation.token,
-                scopedTelemetryRecorder
+                scopedTelemetryRecorder,
+                contextLimit
             )
             if (cancellation.token.isCancellationRequested) {
                 return
@@ -781,7 +790,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
     ): Promise<Message[]> {
         const { prompt, newContextUsed, newContextIgnored } = await prompter.makePrompt(
             this.chatModel,
-            ModelProvider.getMaxCharsByModel(this.chatModel.modelID)
+            this.chatModel.maxChars
         )
 
         // Update UI based on prompt construction

vscode/src/chat/context/chatContext.ts

@@ -19,7 +19,9 @@ export async function getChatContextItemsForMention(
     telemetryRecorder?: {
         empty: () => void
         withType: (type: MentionQuery['type']) => void
-    }
+    },
+    // The number of characters left in current context window.
+    maxChars?: number
 ): Promise<ContextItem[]> {
     const mentionQuery = parseMentionQuery(query)
 
@@ -35,12 +37,12 @@ export async function getChatContextItemsForMention(
     const MAX_RESULTS = 20
     switch (mentionQuery.type) {
         case 'empty':
-            return getOpenTabsContextFile()
+            return getOpenTabsContextFile(maxChars)
         case 'symbol':
             // It would be nice if the VS Code symbols API supports cancellation, but it doesn't
             return getSymbolContextFiles(mentionQuery.text, MAX_RESULTS)
         case 'file':
-            return getFileContextFiles(mentionQuery.text, MAX_RESULTS, cancellationToken)
+            return getFileContextFiles(mentionQuery.text, MAX_RESULTS, maxChars)
         case 'url':
             return (await isURLContextFeatureFlagEnabled())
                 ? getURLContextItems(

vscode/src/chat/utils.ts

@@ -1,4 +1,4 @@
-import type { AuthStatus } from '@sourcegraph/cody-shared'
+import type { AuthStatus, ChatMessage } from '@sourcegraph/cody-shared'
 import { defaultAuthStatus, unauthenticatedStatus } from './protocol'
 
 /**
@@ -82,3 +82,35 @@ export const countGeneratedCode = (text: string): { lineCount: number; charCount
     }
     return count
 }
+
+/**
+ * Counts the total number of bytes used in a list of chat messages.
+ *
+ * This function is exported and can be used to calculate the byte usage
+ * of chat messages for storage/bandwidth purposes.
+ *
+ * @param messages - The list of chat messages to count bytes for
+ * @returns The total number of bytes used in the messages
+ */
+export function countBytesInChatMessages(messages: ChatMessage[]): number {
+    if (messages.length === 0) {
+        return 0
+    }
+    return messages.reduce((acc, msg) => acc + msg.speaker.length + (msg.text?.length || 0) + 3, 0)
+}
+
+/**
+ * Gets the context window limit in bytes for chat messages, taking into
+ * account the maximum allowed character count. Returns 0 if the used bytes
+ * exceeds the limit.
+ * @param messages - The chat messages
+ * @param maxChars - The maximum allowed character count
+ * @returns The context window limit in bytes
+ */
+export function getContextWindowLimitInBytes(messages: ChatMessage[], maxChars: number): number {
+    const used = countBytesInChatMessages(messages)
+    if (used > maxChars) {
+        return 0
+    }
+    return maxChars - used
+}

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

@@ -1,6 +1,6 @@
 import type { ContextItem, MentionQuery } from '@sourcegraph/cody-shared'
-import * as vscode from 'vscode'
 
+import { DEFAULT_FAST_MODEL_CHARS_LIMIT } from '@sourcegraph/cody-shared/src/prompt/constants'
 import { getFileContextFiles, getSymbolContextFiles } from '../../editor/utils/editor-context'
 import { getLabelForContextItem } from './utils'
 
@@ -29,11 +29,10 @@ export async function getMatchingContext(
     }
 
     if (mentionQuery.type === 'file') {
-        const cancellation = new vscode.CancellationTokenSource()
         const fileResults = await getFileContextFiles(
             mentionQuery.text,
             MAX_FUZZY_RESULTS,
-            cancellation.token
+            DEFAULT_FAST_MODEL_CHARS_LIMIT
         )
         return fileResults.map(result => ({
             key: getLabelForContextItem(result),

vscode/src/editor/utils/editor-context.test.ts

@@ -41,11 +41,7 @@ describe('getFileContextFiles', () => {
     }
 
     async function runSearch(query: string, maxResults: number): Promise<(string | undefined)[]> {
-        const results = await getFileContextFiles(
-            query,
-            maxResults,
-            new vscode.CancellationTokenSource().token
-        )
+        const results = await getFileContextFiles(query, maxResults)
 
         return results.map(f => uriBasename(f.uri))
     }
@@ -158,13 +154,14 @@ describe('filterLargeFiles', () => {
         expect(filtered).toEqual([])
     })
 
-    it('sets isTooLarge for files exceeding token limit', async () => {
+    it('sets isTooLarge for files exceeding token limit but under 1MB', async () => {
         const largeTextFile: ContextItemFile = {
             uri: vscode.Uri.file('/large-text.txt'),
             type: 'file',
         }
+        const oneByteOverTokenLimit = MAX_CURRENT_FILE_TOKENS * CHARS_PER_TOKEN + 1
         vscode.workspace.fs.stat = vi.fn().mockResolvedValueOnce({
-            size: MAX_CURRENT_FILE_TOKENS * CHARS_PER_TOKEN + 1,
+            size: oneByteOverTokenLimit,
             type: vscode.FileType.File,
         } as vscode.FileStat)
 

vscode/src/editor/utils/editor-context.ts

@@ -56,7 +56,7 @@ const throttledFindFiles = throttle(() => findWorkspaceFiles(), 10000)
 export async function getFileContextFiles(
     query: string,
     maxResults: number,
-    token: vscode.CancellationToken
+    charsLimit?: number
 ): Promise<ContextItemFile[]> {
     if (!query.trim()) {
         return []
@@ -116,7 +116,7 @@ export async function getFileContextFiles(
 
     // TODO(toolmantim): Add fuzzysort.highlight data to the result so we can show it in the UI
 
-    return await filterLargeFiles(sortedResults)
+    return await filterLargeFiles(sortedResults, charsLimit)
 }
 
 export async function getSymbolContextFiles(
@@ -182,11 +182,12 @@ export async function getSymbolContextFiles(
  * Gets context files for each open editor tab in VS Code.
  * Filters out large files over 1MB to avoid expensive parsing.
  */
-export async function getOpenTabsContextFile(): Promise<ContextItemFile[]> {
+export async function getOpenTabsContextFile(charsLimit?: number): Promise<ContextItemFile[]> {
     return await filterLargeFiles(
         getOpenTabsUris()
             .filter(uri => !isCodyIgnoredFile(uri))
-            .flatMap(uri => createContextFileFromUri(uri, ContextItemSource.User, 'file'))
+            .flatMap(uri => createContextFileFromUri(uri, ContextItemSource.User, 'file')),
+        charsLimit
     )
 }
 
@@ -253,7 +254,10 @@ function createContextFileRange(selectionRange: vscode.Range): ContextItem['rang
  * Filters the given context files to remove files larger than 1MB and non-text files.
  * Sets {@link ContextItemFile.isTooLarge} for files contains more characters than the token limit.
  */
-export async function filterLargeFiles(contextFiles: ContextItemFile[]): Promise<ContextItemFile[]> {
+export async function filterLargeFiles(
+    contextFiles: ContextItemFile[],
+    charsLimit = CHARS_PER_TOKEN * MAX_CURRENT_FILE_TOKENS
+): Promise<ContextItemFile[]> {
     const filtered = []
     for (const cf of contextFiles) {
         // Remove file larger than 1MB and non-text files
@@ -262,13 +266,13 @@ export async function filterLargeFiles(contextFiles: ContextItemFile[]): Promise
             stat => stat,
             error => undefined
         )
-        if (fileStat?.type !== vscode.FileType.File || fileStat?.size > 1000000) {
+        if (cf.type !== 'file' || 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 {@link ContextItemFile.isTooLarge} for webview to display file size
         // warning.
-        if (fileStat.size > CHARS_PER_TOKEN * MAX_CURRENT_FILE_TOKENS) {
+        if (fileStat.size > charsLimit) {
             cf.isTooLarge = true
         }
         filtered.push(cf)