Chat: fix at-mention token size

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ContextItem.kt

@@ -31,6 +31,7 @@ data class ContextItemFile(
   val source: ContextItemSource? = null, // Oneof: embeddings, user, keyword, editor, filename, search, unified, selection, terminal
   val type: TypeEnum? = null, // Oneof: file
   val isTooLarge: Boolean? = null,
+  val size: Int? = null,
 ) : ContextItem() {
 
   enum class TypeEnum {

lib/shared/src/codebase-context/messages.ts

@@ -84,6 +84,8 @@ export interface ContextItemFile extends ContextItemCommon {
      * Whether the file is too large to be included as context.
      */
     isTooLarge?: boolean
+
+    size?: number
 }
 
 /**

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,27 @@ import type { Repo } from '../../context/repo-fetcher'
 import { getChatPanelTitle } from './chat-helpers'
 
 export class SimpleChatModel {
+    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 get charsLeft(): number {
+        let used = 0
+        for (const msg of this.messages) {
+            if (used > this.maxChars) {
+                return 0
+            }
+            used += msg.speaker.length + (msg.text?.length || 0) + 3
+        }
+        return this.maxChars - used
+    }
 
     public isEmpty(): boolean {
         return this.messages.length === 0

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

@@ -52,6 +52,7 @@ 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 } 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'
@@ -601,7 +602,8 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
             const items = await getChatContextItemsForMention(
                 query,
                 cancellation.token,
-                scopedTelemetryRecorder
+                scopedTelemetryRecorder,
+                this.chatModel.charsLeft - ANSWER_TOKENS
             )
             if (cancellation.token.isCancellationRequested) {
                 return
@@ -779,8 +781,10 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         prompter: IPrompter,
         sendTelemetry?: (contextSummary: any) => void
     ): Promise<Message[]> {
-        const maxChars = ModelProvider.getMaxCharsByModel(this.chatModel.modelID)
-        const { prompt, newContextUsed } = await prompter.makePrompt(this.chatModel, maxChars)
+        const { prompt, newContextUsed } = await prompter.makePrompt(
+            this.chatModel,
+            this.chatModel.maxChars
+        )
 
         // Update UI based on prompt construction
         this.chatModel.setLastMessageContext(newContextUsed)

vscode/src/chat/context/chatContext.ts

@@ -19,7 +19,8 @@ export async function getChatContextItemsForMention(
     telemetryRecorder?: {
         empty: () => void
         withType: (type: MentionQuery['type']) => void
-    }
+    },
+    charsLimit?: number
 ): Promise<ContextItem[]> {
     const mentionQuery = parseMentionQuery(query)
 
@@ -35,12 +36,12 @@ export async function getChatContextItemsForMention(
     const MAX_RESULTS = 20
     switch (mentionQuery.type) {
         case 'empty':
-            return getOpenTabsContextFile()
+            return getOpenTabsContextFile(charsLimit)
         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, cancellationToken, charsLimit)
         case 'url':
             return (await isURLContextFeatureFlagEnabled())
                 ? getURLContextItems(

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

@@ -158,13 +158,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)
 
@@ -174,6 +175,7 @@ describe('filterLargeFiles', () => {
             type: 'file',
             uri: largeTextFile.uri,
             isTooLarge: true,
+            size: oneByteOverTokenLimit,
         })
     })
 })

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

@@ -56,7 +56,8 @@ const throttledFindFiles = throttle(() => findWorkspaceFiles(), 10000)
 export async function getFileContextFiles(
     query: string,
     maxResults: number,
-    token: vscode.CancellationToken
+    cancellationToken: vscode.CancellationToken,
+    charsLimit?: number
 ): Promise<ContextItemFile[]> {
     if (!query.trim()) {
         return []
@@ -116,7 +117,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 +183,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 +255,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 +267,14 @@ 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) {
+        cf.size = fileStat.size
+        if (fileStat.size > charsLimit) {
             cf.isTooLarge = true
         }
         filtered.push(cf)