fix(amazonq): auto scan use unsaved content

plugins/toolkit/core/src/software/aws/toolkits/core/utils/ZipUtils.kt

@@ -13,10 +13,8 @@ import java.util.zip.ZipOutputStream
  * Adds a new [ZipEntry] with the contents of [file] to the [ZipOutputStream].
  */
 fun ZipOutputStream.putNextEntry(entryName: String, file: Path) {
-    this.putNextEntry(ZipEntry(entryName))
     val bytes = Files.readAllBytes(file)
-    this.write(bytes, 0, bytes.size)
-    this.closeEntry()
+    putNextEntry(entryName, bytes)
 }
 
 /**
@@ -28,6 +26,15 @@ fun ZipOutputStream.putNextEntry(entryName: String, inputStream: InputStream) {
     this.closeEntry()
 }
 
+/**
+ * Adds a new [ZipEntry] with the contents of [data] to the [ZipOutputStream].
+ */
+fun ZipOutputStream.putNextEntry(entryName: String, data: ByteArray) {
+    this.putNextEntry(ZipEntry(entryName))
+    this.write(data, 0, data.size)
+    this.closeEntry()
+}
+
 /**
  * Create a zip file in a temporary location.
  *

plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/listeners/CodeWhispererCodeScanDocumentListener.kt

@@ -8,8 +8,8 @@ import com.intellij.openapi.editor.event.DocumentListener
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.vfs.isFile
+import com.intellij.refactoring.suggested.oldRange
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanIssue
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
@@ -35,7 +35,7 @@ internal class CodeWhispererCodeScanDocumentListener(val project: Project) : Doc
             else -> 0
         }
 
-        val editedTextRange = TextRange.create(event.offset, event.offset + event.oldLength)
+        val editedTextRange = event.oldRange
         scanManager.updateScanNodesForOffSet(file, lineOffset, editedTextRange)
         val nodes = scanManager.getOverlappingScanNodes(file, editedTextRange)
         nodes.forEach {

plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/sessionconfig/CodeScanSessionConfig.kt

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig
 
 import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.guessModuleDir
 import com.intellij.openapi.project.guessProjectDir
@@ -12,6 +13,7 @@ import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vcs.changes.ChangeListManager
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.openapi.vfs.isFile
 import kotlinx.coroutines.runBlocking
 import software.aws.toolkits.core.utils.createTemporaryZipFile
@@ -106,7 +108,10 @@ class CodeScanSessionConfig(
         }
 
         // Copy all the included source files to the source zip
-        val srcZip = zipFiles(payloadMetadata.sourceFiles.map { Path.of(it) })
+        val srcZip = when (scope) {
+            CodeAnalysisScope.PROJECT -> zipProject(payloadMetadata.sourceFiles.map { Path.of(it) })
+            CodeAnalysisScope.FILE -> zipFile(Path.of(payloadMetadata.sourceFiles.first()))
+        }
         val payloadContext = PayloadContext(
             payloadMetadata.language,
             payloadMetadata.linesScanned,
@@ -147,7 +152,7 @@ class CodeScanSessionConfig(
         }
     }
 
-    private fun zipFiles(files: List<Path>): File = createTemporaryZipFile {
+    private fun zipProject(files: List<Path>): File = createTemporaryZipFile {
         files.forEach { file ->
             try {
                 val relativePath = file.relativeTo(projectRoot.toNioPath())
@@ -159,6 +164,23 @@ class CodeScanSessionConfig(
         }
     }.toFile()
 
+    private fun zipFile(filePath: Path): File = createTemporaryZipFile {
+        try {
+            val relativePath = filePath.relativeTo(projectRoot.toNioPath())
+            val virtualFile = VirtualFileManager.getInstance().findFileByNioPath(filePath)
+            virtualFile?.let { file ->
+                val document = runReadAction {
+                    FileDocumentManager.getInstance().getDocument(file)
+                }
+                document?.let { doc ->
+                    it.putNextEntry(relativePath.toString(), doc.text.encodeToByteArray())
+                }
+            }
+        } catch (e: Exception) {
+            cannotFindFile("Zipping error: ${e.message}", filePath.pathString)
+        }
+    }.toFile()
+
     fun getProjectPayloadMetadata(): PayloadMetadata {
         val files = mutableSetOf<String>()
         val traversedDirectories = mutableSetOf<VirtualFile>()

plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeFileScanTest.kt

@@ -3,11 +3,13 @@
 
 package software.aws.toolkits.jetbrains.services.codewhisperer.codescan
 
+import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
 import com.intellij.testFramework.replaceService
+import com.intellij.testFramework.runInEdtAndWait
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import org.apache.commons.codec.digest.DigestUtils
@@ -26,6 +28,7 @@ import org.mockito.kotlin.doNothing
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.inOrder
 import org.mockito.kotlin.isNull
+import org.mockito.kotlin.reset
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
 import org.mockito.kotlin.verify
@@ -41,11 +44,14 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhisperer
 import software.aws.toolkits.jetbrains.utils.isInstanceOf
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.telemetry.CodewhispererLanguage
+import java.io.BufferedInputStream
 import java.io.File
 import java.io.FileInputStream
 import java.util.Base64
+import java.util.Scanner
 import java.util.UUID
 import java.util.zip.ZipFile
+import java.util.zip.ZipInputStream
 import kotlin.io.path.relativeTo
 import kotlin.test.assertNotNull
 
@@ -88,7 +94,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
         )
         file2 = psifile2.virtualFile.toNioPath().toFile()
 
-        psifile = projectRule.fixture.addFileToProject(
+        psifile = projectRule.fixture.configureByText(
             "/test.py",
             """import numpy as np
                import from module1 import helper
@@ -386,6 +392,37 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
         }
     }
 
+    @Test
+    fun `test zipFile uses latest document changes`() {
+        reset(sessionConfigSpy)
+        runInEdtAndWait {
+            WriteCommandAction.runWriteCommandAction(project) {
+                projectRule.fixture.editor.document.insertString(0, "line inserted at beginning of file\n")
+            }
+        }
+
+        val payload = sessionConfigSpy.createPayload()
+        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
+        val zis = ZipInputStream(bufferedInputStream)
+        zis.nextEntry
+        val scanner = Scanner(zis)
+        var contents = ""
+        while (scanner.hasNextLine()) {
+            contents += scanner.nextLine() + "\n"
+        }
+
+        val expectedContents = """line inserted at beginning of file
+import numpy as np
+               import from module1 import helper
+
+               def add(a, b):
+                  return a + b
+
+"""
+
+        assertThat(contents).isEqualTo(expectedContents)
+    }
+
     companion object {
         const val TIMEOUT = 10L * TOTAL_MILLIS_IN_SECOND
     }