context transparency multiple fixes (#17)

* context transparency multiple fixes

1. add manually selected context from prompt into context list
2. fix -1 start/end line edge cases
3. Center selction after clicking on the file

* remove console.log

Author: andrewyuq

packages/core/src/amazonq/lsp/lspController.ts

@@ -13,7 +13,6 @@ import fetch from 'node-fetch'
 import request from '../../shared/request'
 import { LspClient } from './lspClient'
 import AdmZip from 'adm-zip'
-import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
 import { makeTemporaryToolkitFolder, tryRemoveFolder } from '../../shared/filesystemUtilities'
 import { activate as activateLsp } from './lspClient'
 import { telemetry } from '../../shared/telemetry/telemetry'
@@ -25,6 +24,7 @@ import { isWeb } from '../../shared/extensionGlobals'
 import { getUserAgent } from '../../shared/telemetry/util'
 import { isAmazonInternalOs } from '../../shared/vscode/env'
 import { sleep } from '../../shared/utilities/timeoutUtils'
+import { RelevantTextDocumentAddition } from '../../codewhispererChat/controllers/chat/model'
 
 export interface Chunk {
     readonly filePath: string
@@ -282,9 +282,9 @@ export class LspController {
         }
     }
 
-    async query(s: string): Promise<RelevantTextDocument[]> {
+    async query(s: string): Promise<RelevantTextDocumentAddition[]> {
         const chunks: Chunk[] | undefined = await LspClient.instance.queryVectorIndex(s)
-        const resp: RelevantTextDocument[] = []
+        const resp: RelevantTextDocumentAddition[] = []
         if (chunks) {
             for (const chunk of chunks) {
                 const text = chunk.context ? chunk.context : chunk.content

packages/core/src/amazonq/webview/ui/main.ts

@@ -351,6 +351,7 @@ export const createMynahUI = (
                     ...(item.relatedContent !== undefined ? { relatedContent: item.relatedContent } : {}),
                     ...(item.followUp !== undefined ? { followUp: item.followUp } : {}),
                     ...(item.fileList !== undefined ? { fileList: item.fileList } : {}),
+                    ...(item.header !== undefined ? { header: item.header } : {}),
                 })
                 if (
                     item.messageId !== undefined &&
@@ -381,7 +382,11 @@ export const createMynahUI = (
                                 file.relativeFilePath,
                                 {
                                     label: file.lineRanges
-                                        .map((range) => `line ${range.first} - ${range.second}`)
+                                        .map((range) =>
+                                            range.first === -1 || range.second === -1
+                                                ? ''
+                                                : `line ${range.first} - ${range.second}`
+                                        )
                                         .join(', '),
                                     description: file.relativeFilePath,
                                     clickable: true,

packages/core/src/codewhisperer/client/user-service-2.json

@@ -1907,12 +1907,6 @@
                 "documentSymbols": {
                     "shape": "DocumentSymbols",
                     "documentation": "<p>DocumentSymbols parsed from a text document</p>"
-                },
-                "startLine": {
-                    "shape": "Integer"
-                },
-                "endLine": {
-                    "shape": "Integer"
                 }
             },
             "documentation": "<p>Represents an IDE retrieved relevant Text Document / File</p>"

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -30,6 +30,7 @@ import {
     QuickCommandGroupActionClick,
     MergedRelevantDocument,
     FileClick,
+    RelevantTextDocumentAddition,
 } from './model'
 import {
     AppToWebViewMessageDispatcher,
@@ -43,7 +44,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
 import { PromptsGenerator } from './prompts/promptsGenerator'
 import { TriggerEventsStorage } from '../../storages/triggerEvents'
 import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
-import { CodeWhispererStreamingServiceException, RelevantTextDocument } from '@amzn/codewhisperer-streaming'
+import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
 import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
 import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
 import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
@@ -600,6 +601,8 @@ export class ChatController {
         if (!lineRanges) {
             return
         }
+
+        // TODO: Fix for multiple workspace setup
         const projectRoot = workspace.workspaceFolders?.[0]?.uri.fsPath
         if (!projectRoot) {
             return
@@ -612,15 +615,25 @@ export class ChatController {
         const editor = await window.showTextDocument(document, ViewColumn.Active)
 
         // Create multiple selections based on line ranges
-        const selections: Selection[] = lineRanges.map(({ first, second }) => {
-            const startPosition = new Position(first - 1, 0) // Convert 1-based to 0-based
-            const endPosition = new Position(second - 1, document.lineAt(second - 1).range.end.character)
-            return new Selection(startPosition.line, startPosition.character, endPosition.line, endPosition.character)
-        })
+        const selections: Selection[] = lineRanges
+            .filter(({ first, second }) => first !== -1 && second !== -1)
+            .map(({ first, second }) => {
+                const startPosition = new Position(first - 1, 0) // Convert 1-based to 0-based
+                const endPosition = new Position(second - 1, document.lineAt(second - 1).range.end.character)
+                return new Selection(
+                    startPosition.line,
+                    startPosition.character,
+                    endPosition.line,
+                    endPosition.character
+                )
+            })
 
         // Apply multiple selections to the editor using the new API
-        editor.selection = selections[0] // Set the first selection as active
-        editor.selections = selections // Apply multiple selections
+        if (selections.length > 0) {
+            editor.selection = selections[0] // Set the first selection as active
+            editor.selections = selections // Apply multiple selections
+            editor.revealRange(selections[0], vscode.TextEditorRevealType.InCenter)
+        }
     }
 
     private processException(e: any, tabID: string) {
@@ -868,22 +881,24 @@ export class ChatController {
         this.messenger.sendStaticTextResponse(responseType, triggerID, tabID)
     }
 
-    private async resolveContextCommandPayload(triggerPayload: TriggerPayload) {
+    private async resolveContextCommandPayload(triggerPayload: TriggerPayload): Promise<string[]> {
         if (triggerPayload.context === undefined || triggerPayload.context.length === 0) {
-            return
+            return []
         }
         const contextCommands: ContextCommandItem[] = []
+        const relativePaths: string[] = []
         for (const context of triggerPayload.context) {
             if (typeof context !== 'string' && context.route && context.route.length === 2) {
                 contextCommands.push({
                     workspaceFolder: context.route?.[0] || '',
                     type: context.icon === 'folder' ? 'folder' : 'file',
                     relativePath: context.route?.[1] || '',
                 })
+                relativePaths.push(context.route[1])
             }
         }
         if (contextCommands.length === 0) {
-            return
+            return []
         }
         const prompts = await LspClient.instance.getContextCommandPrompt(contextCommands)
         if (prompts.length > 0) {
@@ -902,6 +917,7 @@ export class ChatController {
                 `Retrieved chunks of additional context count: ${triggerPayload.additionalContents.length} `
             )
         }
+        return relativePaths
     }
 
     private async generateResponse(
@@ -937,7 +953,7 @@ export class ChatController {
             return
         }
 
-        await this.resolveContextCommandPayload(triggerPayload)
+        const relativePaths = await this.resolveContextCommandPayload(triggerPayload)
         // TODO: resolve the context into real context up to 90k
         triggerPayload.useRelevantDocuments = false
         if (triggerPayload.message) {
@@ -988,11 +1004,24 @@ export class ChatController {
 
         session.currentContextId++
         session.contexts.set(session.currentContextId, new Map())
-        if (triggerPayload.mergedRelevantDocuments) {
-            for (const doc of triggerPayload.mergedRelevantDocuments) {
-                const currentContext = session.contexts.get(session.currentContextId)
-                if (currentContext) {
-                    currentContext.set(doc.relativeFilePath, doc.lineRanges)
+        if (triggerPayload.mergedRelevantDocuments !== undefined) {
+            const relativePathsOfMergedRelevantDocuments = triggerPayload.mergedRelevantDocuments.map(
+                (doc) => doc.relativeFilePath
+            )
+            for (const relativePath of relativePaths) {
+                if (!relativePathsOfMergedRelevantDocuments.includes(relativePath)) {
+                    triggerPayload.mergedRelevantDocuments.push({
+                        relativeFilePath: relativePath,
+                        lineRanges: [{ first: -1, second: -1 }],
+                    })
+                }
+            }
+            if (triggerPayload.mergedRelevantDocuments) {
+                for (const doc of triggerPayload.mergedRelevantDocuments) {
+                    const currentContext = session.contexts.get(session.currentContextId)
+                    if (currentContext) {
+                        currentContext.set(doc.relativeFilePath, doc.lineRanges)
+                    }
                 }
             }
         }
@@ -1037,7 +1066,7 @@ export class ChatController {
     }
 
     private mergeRelevantTextDocuments(
-        documents: RelevantTextDocument[] | undefined
+        documents: RelevantTextDocumentAddition[] | undefined
     ): MergedRelevantDocument[] | undefined {
         if (documents === undefined) {
             return undefined

packages/core/src/codewhispererChat/controllers/chat/model.ts

@@ -187,13 +187,16 @@ export interface TriggerPayload {
     readonly userIntent: UserIntent | undefined
     readonly customization: Customization
     readonly context?: string[] | QuickActionCommand[]
-    relevantTextDocuments?: RelevantTextDocument[]
+    relevantTextDocuments?: RelevantTextDocumentAddition[]
     additionalContents?: AdditionalContentEntry[]
     mergedRelevantDocuments?: MergedRelevantDocument[]
     useRelevantDocuments?: boolean
     traceId?: string
 }
 
+// TODO move this to API definition (or just use this across the codebase)
+export type RelevantTextDocumentAddition = RelevantTextDocument & { startLine: number; endLine: number }
+
 export interface MergedRelevantDocument {
     readonly relativeFilePath: string
     readonly lineRanges: Array<{ first: number; second: number }>