Context: Update code template for selection (#4123)

agent/src/__tests__/example-ts/.cody/commands.json

@@ -28,7 +28,7 @@
       }
     },
     "countDirFiles": {
-      "prompt": "How many file context have I shared with you?",
+      "prompt": "How many file context have I shared with you? Reply single number. Skip preamble.",
       "context": {
         "currentDir": true
       }

lib/shared/src/editor/displayPath.test.ts

@@ -2,12 +2,12 @@ import { afterEach, beforeEach, describe, expect, test } from 'vitest'
 import { URI } from 'vscode-uri'
 
 import { isWindows } from '../common/platform'
-
 import {
     type DisplayPathEnvInfo,
     displayPath,
     displayPathBasename,
     displayPathDirname,
+    displayPathWithLines,
     setDisplayPathEnvInfo,
     uriHasPrefix,
 } from './displayPath'
@@ -413,3 +413,68 @@ describe('setDisplayPathEnvInfo', () => {
         }).toThrowError('no environment info for displayPath')
     })
 })
+
+describe('displayPathWithLines', () => {
+    const testCases = [
+        ...DISPLAY_PATH_TEST_CASES.flatMap(({ name, tests: { nonWindows, windows, all } }) => [
+            ...(nonWindows
+                ? [
+                      {
+                          name: `nonWindows: ${name}`,
+                          envInfo: { ...nonWindows.envInfo, isWindows: false },
+                          cases: nonWindows.cases,
+                      },
+                  ]
+                : []),
+            ...(windows
+                ? [
+                      {
+                          name: `windows: ${name}`,
+                          envInfo: { ...windows.envInfo, isWindows: true },
+                          cases: windows.cases,
+                      },
+                  ]
+                : []),
+            ...(all
+                ? [
+                      {
+                          name: `all nonWindows: ${name}`,
+                          envInfo: { ...all.envInfo, isWindows: false },
+                          cases: all.cases,
+                      },
+                      {
+                          name: `all windows: ${name}`,
+                          envInfo: { ...all.envInfo, isWindows: true },
+                          cases: all.cases,
+                      },
+                  ]
+                : []),
+        ]),
+    ]
+
+    for (const { name, envInfo, cases } of testCases) {
+        if (envInfo.isWindows === false) {
+            describe.skipIf(isWindows())(name, () => {
+                for (const { input, expected } of cases) {
+                    const range = { start: { line: 0, character: 5 }, end: { line: 99, character: 0 } }
+                    test(`${input.fsPath} -> ${expected}:1-99`, () => {
+                        expect(withEnvInfo(envInfo, () => displayPathWithLines(input, range))).toBe(
+                            `${expected}:1-99`
+                        )
+                    })
+                }
+            })
+        } else {
+            describe(name, () => {
+                for (const { input, expected } of cases) {
+                    const range = { start: { line: 0, character: 5 }, end: { line: 99, character: 0 } }
+                    test(`${input.fsPath} -> ${expected}:1-99`, () => {
+                        expect(withEnvInfo(envInfo, () => displayPathWithLines(input, range))).toBe(
+                            `${expected}:1-99`
+                        )
+                    })
+                }
+            })
+        }
+    }
+})

lib/shared/src/editor/displayPath.ts

@@ -1,6 +1,7 @@
 import { URI } from 'vscode-uri'
 
 import { pathFunctionsForURI, posixFilePaths, windowsFilePaths } from '../common/path'
+import { type RangeData, displayLineRange } from '../common/range'
 
 /**
  * Convert an absolute URI to a (possibly shorter) path to display to the user. The display path is
@@ -27,6 +28,18 @@ export function displayPath(location: URI): string {
     return typeof result === 'string' ? result : result.toString()
 }
 
+/**
+ * Displays the path of a URI {@link displayPath} with zero-based line values extracted from range.
+ * Example: `src/foo/bar.ts:5-10`
+ *
+ * @param location - The URI to display the path for.
+ * @param range - The line range data to display.
+ * @returns The formatted path with the start and end lines of the range appended to it.
+ */
+export function displayPathWithLines(location: URI, range: RangeData): string {
+    return `${displayPath(location)}:${displayLineRange(range)}`
+}
+
 /**
  * Dirname of the location's display path, to display to the user. Similar to
  * `dirname(displayPath(location))`, but it uses the right path separators in `dirname` ('\' for

lib/shared/src/prompt/prompt-string.ts

@@ -5,11 +5,12 @@ import type { ContextItem } from '../codebase-context/messages'
 import type { ContextFiltersProvider } from '../cody-ignore/context-filters-provider'
 import type { TerminalOutputArguments } from '../commands/types'
 import { markdownCodeBlockLanguageIDForFilename } from '../common/languages'
+import type { RangeData } from '../common/range'
 import type { AutocompleteContextSnippet, DocumentContext } from '../completions/types'
 import type { ConfigGetter } from '../configuration'
 import type { ActiveTextEditorDiagnostic } from '../editor'
 import { createGitDiff } from '../editor/create-git-diff'
-import { displayPath } from '../editor/displayPath'
+import { displayPath, displayPathWithLines } from '../editor/displayPath'
 import { getEditorInsertSpaces, getEditorTabSize } from '../editor/utils'
 import { logDebug } from '../logger'
 import { telemetryRecorder } from '../telemetry-v2/singleton'
@@ -237,6 +238,11 @@ export class PromptString {
         return internal_createPromptString(displayPath(uri), [uri])
     }
 
+    public static fromDisplayPathLineRange(uri: vscode.Uri, range?: RangeData) {
+        const pathToDisplay = range ? displayPathWithLines(uri, range) : displayPath(uri)
+        return internal_createPromptString(pathToDisplay, [uri])
+    }
+
     public static fromDocumentText(document: vscode.TextDocument, range?: vscode.Range): PromptString {
         return internal_createPromptString(document.getText(range), [document.uri])
     }

lib/shared/src/prompt/templates.ts

@@ -1,5 +1,6 @@
 import type { URI } from 'vscode-uri'
 
+import type { RangeData } from '../common/range'
 import type { ActiveTextEditorDiagnostic } from '../editor'
 import { displayPath } from '../editor/displayPath'
 import { PromptString, ps } from './prompt-string'
@@ -48,29 +49,17 @@ export function populateTerminalOutputContextTemplate(output: string): string {
     return COMMAND_OUTPUT_TEMPLATE + output
 }
 
-const SELECTED_CODE_CONTEXT_TEMPLATE = ps`My selected {languageName} code from file \`{filePath}\`:
-<selected>
-{code}
-</selected>`
-
-const SELECTED_CODE_CONTEXT_TEMPLATE_WITH_REPO = ps`My selected {languageName} code from file \`{filePath}\` in \`{repoName}\` repository:
-<selected>
-{code}
-</selected>`
+const SELECTED_CODE_CONTEXT_TEMPLATE = ps`My selected code from @{filePath}{codebase}:\n\`\`\`\n{code}\`\`\``
 
 export function populateCurrentSelectedCodeContextTemplate(
     code: PromptString,
     fileUri: URI,
+    range?: RangeData,
     repoName?: PromptString
 ): PromptString {
-    return (
-        repoName
-            ? SELECTED_CODE_CONTEXT_TEMPLATE_WITH_REPO.replace('{repoName}', repoName)
-            : SELECTED_CODE_CONTEXT_TEMPLATE
-    )
-        .replace('{code}', code)
-        .replaceAll('{filePath}', PromptString.fromDisplayPath(fileUri))
-        .replace('{languageName}', PromptString.fromMarkdownCodeBlockLanguageIDForFilename(fileUri))
+    return SELECTED_CODE_CONTEXT_TEMPLATE.replace('{code}', code)
+        .replace('{codebase}', repoName ? ps` in \`{repoName}\` repository` : ps``)
+        .replaceAll('{filePath}', PromptString.fromDisplayPathLineRange(fileUri, range))
 }
 
 const DIRECTORY_FILE_LIST_TEMPLATE =
@@ -88,9 +77,12 @@ export function populateListOfFilesContextTemplate(fileList: string, fileUri?: U
 export function populateContextTemplateFromText(
     templateText: PromptString,
     content: PromptString,
-    fileUri: URI
+    fileUri: URI,
+    range?: RangeData
 ): PromptString {
-    return templateText.replace('{fileName}', PromptString.fromDisplayPath(fileUri)).concat(content)
+    return templateText
+        .replace('{displayPath}', PromptString.fromDisplayPathLineRange(fileUri, range))
+        .concat(content)
 }
 
 const FILE_IMPORTS_TEMPLATE = ps`{fileName} has imported the following: `

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

@@ -89,7 +89,7 @@ import { InitDoer } from './InitDoer'
 import { SimpleChatModel, prepareChatMessage } from './SimpleChatModel'
 import { getChatPanelTitle, openFile } from './chat-helpers'
 import { getEnhancedContext } from './context'
-import { DefaultPrompter, type IPrompter } from './prompt'
+import { DefaultPrompter } from './prompt'
 
 interface SimpleChatPanelProviderOptions {
     config: ChatPanelConfig
@@ -494,6 +494,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
                     userContextFiles || [],
                     inputText
                 )
+
                 span.setAttribute('strategy', this.config.useContext)
                 const prompter = new DefaultPrompter(
                     userContextItems,
@@ -594,7 +595,8 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
                 'user',
                 contextFiles,
                 editorState,
-                addEnhancedContext
+                addEnhancedContext,
+                'chat'
             )
         } catch {
             this.postError(new Error('Failed to edit prompt'), 'transcript')
@@ -850,7 +852,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
      * Constructs the prompt and updates the UI with the context used in the prompt.
      */
     private async buildPrompt(
-        prompter: IPrompter,
+        prompter: DefaultPrompter,
         sendTelemetry?: (contextSummary: any, privateContextStats?: any) => void
     ): Promise<Message[]> {
         const { prompt, context } = await prompter.makePrompt(

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

@@ -15,6 +15,7 @@ import {
     wrapInActiveSpan,
 } from '@sourcegraph/cody-shared'
 
+import { getContextFileFromSelection } from '../../commands/context/selection'
 import type { RemoteSearch } from '../../context/remote-search'
 import type { VSCodeEditor } from '../../editor/vscode-editor'
 import type { ContextRankingController } from '../../local-context/context-ranking'
@@ -286,32 +287,6 @@ const userAttentionRegexps: RegExp[] = [
     /have\s+open/,
 ]
 
-function getCurrentSelectionContext(editor: VSCodeEditor): ContextItem[] {
-    const selection = editor.getActiveTextEditorSelection()
-    if (!selection?.selectedText) {
-        return []
-    }
-    let range: vscode.Range | undefined
-    if (selection.selectionRange) {
-        range = new vscode.Range(
-            selection.selectionRange.start.line,
-            selection.selectionRange.start.character,
-            selection.selectionRange.end.line,
-            selection.selectionRange.end.character
-        )
-    }
-
-    return [
-        {
-            type: 'file',
-            content: selection.selectedText,
-            uri: selection.fileUri,
-            range,
-            source: ContextItemSource.Selection,
-        },
-    ]
-}
-
 function getVisibleEditorContext(editor: VSCodeEditor): ContextItem[] {
     return wrapInActiveSpan('chat.context.visibleEditorContext', () => {
         const visible = editor.getActiveTextEditorVisibleContent()
@@ -340,10 +315,9 @@ async function getPriorityContext(
 ): Promise<ContextItem[]> {
     return wrapInActiveSpan('chat.context.priority', async () => {
         const priorityContext: ContextItem[] = []
-        const selectionContext = getCurrentSelectionContext(editor)
-        if (selectionContext.length > 0) {
-            priorityContext.push(...selectionContext)
-        } else if (needsUserAttentionContext(text)) {
+        const selectionContext = await getContextFileFromSelection()
+        priorityContext.push(...selectionContext)
+        if (needsUserAttentionContext(text)) {
             // Query refers to current editor
             priorityContext.push(...getVisibleEditorContext(editor))
         } else if (needsReadmeContext(editor, text)) {

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

@@ -9,9 +9,7 @@ import {
     isDefined,
     wrapInActiveSpan,
 } from '@sourcegraph/cody-shared'
-
 import { logDebug } from '../../log'
-
 import { PromptBuilder } from '../../prompt-builder'
 import type { SimpleChatModel } from './SimpleChatModel'
 
@@ -28,11 +26,7 @@ interface PromptInfo {
     }
 }
 
-export interface IPrompter {
-    makePrompt(chat: SimpleChatModel, codyApiVersion: number): Promise<PromptInfo>
-}
-
-export class DefaultPrompter implements IPrompter {
+export class DefaultPrompter {
     constructor(
         private explicitContext: ContextItemWithContent[],
         private getEnhancedContext?: (query: PromptString) => Promise<ContextItem[]>
@@ -72,7 +66,7 @@ export class DefaultPrompter implements IPrompter {
             // Counter for context items categorized by source
             const ignoredContext = { user: 0, enhanced: 0, transcript: 0 }
 
-            // Add context from new user-specified context items, e.g. @-mentions, @-uri
+            // Add context from new user-specified context items, e.g. @-mentions, active selection, etc.
             const newUserContext = await promptBuilder.tryAddContext('user', this.explicitContext)
             ignoredContext.user += newUserContext.ignored.length
 

vscode/src/commands/context/selection.ts

@@ -11,14 +11,15 @@ import { getEditor } from '../../editor/active-editor'
 import { getSmartSelection } from '../../editor/utils'
 
 import { type Position, Selection } from 'vscode'
+
 /**
- * Gets context file content from the current editor selection.
+ * Gets context file content from the cursor position in the active editor.
  *
  * When no selection is made, try getting the smart selection based on the cursor position.
  * If no smart selection is found, use the visible range of the editor instead.
  */
 export async function getContextFileFromCursor(newCursorPosition?: Position): Promise<ContextItem[]> {
-    return wrapInActiveSpan('commands.context.selection', async span => {
+    return wrapInActiveSpan('commands.context.cursor', async span => {
         try {
             const editor = getEditor()
             const document = editor?.active?.document
@@ -60,3 +61,38 @@ export async function getContextFileFromCursor(newCursorPosition?: Position): Pr
         }
     })
 }
+
+/**
+ * Gets context file content from the current selection in the active editor if any.
+ */
+export async function getContextFileFromSelection(): Promise<ContextItem[]> {
+    return wrapInActiveSpan('commands.context.selection', async span => {
+        try {
+            const editor = getEditor()?.active
+            const document = editor?.document
+            const selection = editor?.selection
+            if (!document || !selection) {
+                return []
+            }
+
+            if (await contextFiltersProvider.isUriIgnored(document.uri)) {
+                return []
+            }
+
+            const content = editor.document.getText(selection)
+            return [
+                {
+                    type: 'file',
+                    uri: document.uri,
+                    content,
+                    source: ContextItemSource.Selection,
+                    range: toRangeData(selection),
+                    size: TokenCounter.countTokens(content),
+                } satisfies ContextItemFile,
+            ]
+        } catch (error) {
+            logError('getContextFileFromCursor', 'failed', { verbose: error })
+            return []
+        }
+    })
+}