Add line ranges to the context files

src/main/kotlin/com/sourcegraph/cody/agent/protocol/ContextFile.kt

@@ -1,6 +1,11 @@
 package com.sourcegraph.cody.agent.protocol
 
-import com.google.gson.*
+import com.google.gson.GsonBuilder
+import com.google.gson.JsonDeserializationContext
+import com.google.gson.JsonDeserializer
+import com.google.gson.JsonElement
+import com.google.gson.JsonObject
+import com.google.gson.JsonSerializer
 import java.lang.reflect.Type
 import java.net.URI
 
@@ -15,6 +20,7 @@ data class ContextFileFile(
     override val uri: URI,
     override val repoName: String?,
     override val revision: String?,
+    val range: Range? = null,
 ) : ContextFile() {
   override val type: String = "file"
 }
@@ -27,8 +33,8 @@ val contextFileDeserializer =
           val uri = context.deserialize<URI>(jsonObject["uri"], URI::class.java)
           val repoName = jsonObject["repoName"]?.asString
           val revision = jsonObject["revision"]?.asString
-
-          ContextFileFile(uri, repoName, revision)
+          val range = gsonMapper.fromJson<Range>(jsonObject["range"], Range::class.java)
+          ContextFileFile(uri, repoName, revision, range)
         }
 
         // TODO(beyang): should throw an exception here, but we don't because the context field is
@@ -37,6 +43,8 @@ val contextFileDeserializer =
       }
     }
 
+private val gsonMapper = GsonBuilder().create()
+
 val uriDeserializer =
     JsonDeserializer { jsonElement: JsonElement, typ: Type, context: JsonDeserializationContext ->
       val j = jsonElement.asJsonObject

src/main/kotlin/com/sourcegraph/cody/chat/ui/ContextFilesPanel.kt

@@ -17,6 +17,7 @@ import com.intellij.util.withFragment
 import com.intellij.util.withQuery
 import com.sourcegraph.cody.agent.protocol.ChatMessage
 import com.sourcegraph.cody.agent.protocol.ContextFile
+import com.sourcegraph.cody.agent.protocol.ContextFileFile
 import com.sourcegraph.cody.agent.protocol.Speaker
 import com.sourcegraph.cody.chat.ChatUIConstants.ASSISTANT_MESSAGE_GRADIENT_WIDTH
 import com.sourcegraph.cody.chat.ChatUIConstants.TEXT_MARGIN
@@ -40,11 +41,16 @@ class ContextFilesPanel(
   }
 
   @RequiresBackgroundThread
-  private fun updateFileList(contextFileNames: Set<String>) {
+  private fun updateFileList(contextFileNames: Map<String, ContextFile>) {
     val filesAvailableInEditor =
         contextFileNames
-            .map(Paths::get)
-            .mapNotNull(VirtualFileManager.getInstance()::findFileByNioPath)
+            .map {
+              val path = Paths.get(it.key)
+              val findFileByNioPath = VirtualFileManager.getInstance().findFileByNioPath(path)
+              findFileByNioPath ?: return@map null
+              Pair(findFileByNioPath, it.value as? ContextFileFile)
+            }
+            .filterNotNull()
             .toList()
 
     ApplicationManager.getApplication().invokeLater {
@@ -54,16 +60,19 @@ class ContextFilesPanel(
       val accordionSection = AccordionSection("Read ${filesAvailableInEditor.size} files")
       accordionSection.isOpaque = false
       accordionSection.border = EmptyBorder(margin)
-      filesAvailableInEditor.forEachIndexed { index, file: VirtualFile ->
-        val filePanel = createFileWithLinkPanel(file)
+      filesAvailableInEditor.forEachIndexed { index, file: Pair<VirtualFile, ContextFileFile?> ->
+        val filePanel = createFileWithLinkPanel(file.first, file.second)
         accordionSection.contentPanel.add(filePanel, index)
       }
       add(accordionSection, BorderLayout.CENTER)
     }
   }
 
   @RequiresEdt
-  private fun createFileWithLinkPanel(file: VirtualFile): JPanel {
+  private fun createFileWithLinkPanel(
+      file: VirtualFile,
+      contextFileFile: ContextFileFile?
+  ): JPanel {
     val projectRelativeFilePath = file.path.removePrefix(project.basePath ?: "")
     val anAction =
         object : DumbAwareAction() {
@@ -73,14 +82,31 @@ class ContextFilesPanel(
             FileEditorManager.getInstance(project).openFile(file, /*focusEditor=*/ true)
           }
         }
-    val goToFile = AnActionLink("@$projectRelativeFilePath", anAction)
+    val actionText = getActionTextForFileLinkAction(contextFileFile, projectRelativeFilePath)
+    val goToFile = AnActionLink(actionText, anAction)
     val panel = JPanel(BorderLayout())
     panel.isOpaque = false
     panel.border = JBUI.Borders.emptyLeft(3)
     panel.add(goToFile, BorderLayout.PAGE_START)
     return panel
   }
 
+  private fun getActionTextForFileLinkAction(
+      contextFileFile: ContextFileFile?,
+      projectRelativeFilePath: String
+  ): String {
+    val startLine = contextFileFile?.range?.start?.line
+    val endLine = contextFileFile?.range?.end?.line
+    return buildString {
+      append("@$projectRelativeFilePath")
+      if (startLine != null && endLine != null) {
+        if (startLine != endLine) {
+          append(":$startLine:$endLine")
+        }
+      }
+    }
+  }
+
   fun updateContentWith(contextFiles: List<ContextFile>?) {
     if (contextFiles.isNullOrEmpty()) {
       return
@@ -89,13 +115,16 @@ class ContextFilesPanel(
     val contextFileNames =
         contextFiles
             .map {
-              if (it.repoName != null) {
-                "${project.basePath}/${it.uri.path}"
-              } else {
-                Paths.get(it.uri.withFragment(null).withQuery(null)).absolutePathString()
-              }
+              val contextFileName =
+                  if (it.repoName != null) {
+                    "${project.basePath}/${it.uri.path}"
+                  } else {
+                    val newUri = it.uri.withFragment(null).withQuery(null)
+                    Paths.get(newUri).absolutePathString()
+                  }
+              contextFileName to it
             }
-            .toSet()
+            .toMap()
 
     ApplicationManager.getApplication().executeOnPooledThread { updateFileList(contextFileNames) }
   }